diff --git a/noir/templating.py b/noir/templating.py index def4d08..151b4c3 100644 --- a/noir/templating.py +++ b/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, diff --git a/tests/test_template_ctx.py b/tests/test_template_ctx.py new file mode 100644 index 0000000..5b56dcd --- /dev/null +++ b/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", + }