diff --git a/kartograf/util.py b/kartograf/util.py index 0770dcb..5691610 100644 --- a/kartograf/util.py +++ b/kartograf/util.py @@ -133,15 +133,11 @@ def wait_for_launch(wait): def format_pfx(pfx): """ - We have seen some formatting issues like leading zeros in the prefix, - which can cause problems. + Attempt to format an IP network or address. + If invalid, return the input unchanged. """ try: if "/" in pfx: - pattern = r"^0+" - match = re.search(pattern, pfx) - if match: - pfx = re.sub(pattern, "", pfx) formatted_pfx = str(ipaddress.ip_network(pfx)) return f"{formatted_pfx}" return str(ipaddress.ip_address(pfx)) @@ -149,6 +145,21 @@ def format_pfx(pfx): return pfx +def is_valid_pfx(pfx): + """ + Check whether the IP network or address provided is valid. + """ + try: + if "/" in pfx: + ipaddress.ip_network(pfx) + return True + else: + ipaddress.ip_address(pfx) + return True + except ValueError: + return False + + def get_root_network(pfx): """ Extract the top-level network from an IPv4 or IPv6 address. diff --git a/tests/util_test.py b/tests/util_test.py index 45431c5..1c23ec5 100644 --- a/tests/util_test.py +++ b/tests/util_test.py @@ -1,4 +1,4 @@ -from kartograf.util import format_pfx, get_root_network +from kartograf.util import format_pfx, is_valid_pfx, get_root_network def test_valid_ipv4_network(): pfx = "192.144.11.0/21" @@ -20,18 +20,6 @@ def test_valid_ipv6_addr(): assert format_pfx(pfx) == pfx -def test_ipv4_prefix_with_leading_zeros(): - pfx = "010.10.00.00/16" - expected_output = "10.10.00.00/16" - assert format_pfx(pfx) == expected_output - - -def test_ipv6_prefix_with_leading_zeros(): - pfx = "001:db8::0/24" - expected_output = "1:db8::0/24" - assert format_pfx(pfx) == expected_output - - def test_invalid_ip_network(): pfx = "192.1/asdf" assert format_pfx(pfx) == pfx @@ -42,6 +30,18 @@ def test_invalid_input(): assert format_pfx(pfx) == pfx +def test_ipv4_prefix_with_leading_zeros(): + pfx = "010.10.00.00/16" + assert format_pfx(pfx) == pfx + assert not is_valid_pfx(pfx) + + +def test_ipv6_prefix_with_leading_zeros(): + pfx = "001:db8::0/24" + assert format_pfx(pfx) == pfx + assert not is_valid_pfx(pfx) + + def test_get_root_network(): ipv4 = "192.144.11.0/24" assert get_root_network(ipv4) == 192