Add cidr helpers to templating ctx (#1)

noir/templating.py

@@ -1,11 +1,12 @@
 import base64
 import datetime
 import hashlib
+import ipaddress
 import json
 import os
 import random
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import tomli_w
 import yaml
@@ -39,6 +40,48 @@ def templater(delimiters: Tuple[str, str]):
     )
 
 
+def _cidr_host(prefix: str, hostnum: int) -> str:
+    net = ipaddress.ip_network(prefix, strict=False)
+    return str(net._address_class(int(net.network_address) + hostnum))
+
+
+def _cidr_netmask(prefix: str) -> str:
+    return str(ipaddress.ip_network(prefix, strict=False).netmask)
+
+
+def _cidr_subnet(prefix: str, newbits: int, netnum: int) -> str:
+    net = ipaddress.ip_network(prefix, strict=False)
+    return str(
+        net.__class__(
+            (
+                int(net.network_address) +
+                (((int(net.hostmask) + 1) >> newbits) * netnum),
+                net._prefixlen + newbits
+            )
+        )
+    )
+
+
+def _cidr_subnets(prefix: str, *newbits: int) -> List[str]:
+    rv = []
+    net = ipaddress.ip_network(prefix, strict=False)
+    net_address = int(net.network_address)
+    shift = 0
+    for newbit in newbits:
+        offset = (int(net.hostmask) + 1) >> newbit
+        delta = 0
+        if shift % offset:
+            step = int(offset)
+            while step < shift:
+                step += offset
+            delta = step - shift
+            offset += delta
+        new_prefixlen = net._prefixlen + newbit
+        rv.append(str(net.__class__((net_address + shift + delta, new_prefixlen))))
+        shift += offset
+    return rv
+
+
 def _indent(text: str, spaces: int = 2) -> str:
     offset = " " * spaces
     rv = f"\n{offset}".join(text.split("\n"))
@@ -64,6 +107,12 @@ def base_ctx(ctx: Dict[str, Any]):
         hashlib=hashlib,
         random=random,
         env=obj_to_adict(os.environ),
+        cidr=adict(
+            host=_cidr_host,
+            netmask=_cidr_netmask,
+            subnet=_cidr_subnet,
+            subnets=_cidr_subnets
+        ),
         indent=_indent,
         to_json=_to_json,
         to_toml=_to_toml,

tests/test_template_ctx.py

@@ -0,0 +1,26 @@
+from noir.templating import _cidr_host, _cidr_netmask, _cidr_subnet, _cidr_subnets
+
+
+def test_cidr_host():
+    assert _cidr_host("10.12.112.0/20", 16) == "10.12.112.16"
+    assert _cidr_host("10.12.112.0/20", 268) == "10.12.113.12"
+    assert _cidr_host("fd00:fd12:3456:7890:00a2::/72", 34) == "fd00:fd12:3456:7890::22"
+
+
+def test_cidr_netmask():
+    assert _cidr_netmask("172.16.0.0/12") == "255.240.0.0"
+
+
+def test_cidr_subnet():
+    assert _cidr_subnet("172.16.0.0/12", 4, 2) == "172.18.0.0/16"
+    assert _cidr_subnet("10.1.2.0/24", 4, 15) == "10.1.2.240/28"
+    assert _cidr_subnet("fd00:fd12:3456:7890::/56", 16, 162) == "fd00:fd12:3456:7800:a200::/72"
+
+
+def test_cidr_subnets():
+    assert set(_cidr_subnets("10.1.0.0/16", 4, 4, 8, 4)) == {
+        "10.1.0.0/20",
+        "10.1.16.0/20",
+        "10.1.32.0/24",
+        "10.1.48.0/20",
+    }