diff --git a/kartograf/util.py b/kartograf/util.py index f89d9e6..0770dcb 100644 --- a/kartograf/util.py +++ b/kartograf/util.py @@ -138,12 +138,17 @@ def format_pfx(pfx): """ 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)) except ValueError: return pfx + def get_root_network(pfx): """ Extract the top-level network from an IPv4 or IPv6 address. diff --git a/tests/merge_base_class_test.py b/tests/merge_base_class_test.py index 9b36912..17a63af 100644 --- a/tests/merge_base_class_test.py +++ b/tests/merge_base_class_test.py @@ -14,7 +14,6 @@ def _df_from_networks(networks, asn=123): for network in networks: ipn = ipaddress.ip_network(network) root_net = get_root_network(network) - print(root_net) network_int = int(ipn.network_address) df.loc[len(df)] = [network_int, asn, str(ipn), root_net] return df diff --git a/tests/test_util.py b/tests/test_util.py new file mode 100644 index 0000000..87a650d --- /dev/null +++ b/tests/test_util.py @@ -0,0 +1,54 @@ +import pytest + +from kartograf.util import format_pfx, get_root_network + +def test_valid_ipv4_network(): + pfx = "192.144.11.0/21" + assert format_pfx(pfx) == pfx + + +def test_valid_ipv4_addr(): + pfx = "192.144.11.0" + assert format_pfx(pfx) == pfx + + +def test_valid_ipv6_network(): + pfx = "2001:db8::/64" + assert format_pfx(pfx) == pfx + + +def test_valid_ipv6_addr(): + pfx = "2001:db8::1" + 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 + + +def test_invalid_input(): + pfx = "no.slash" + assert format_pfx(pfx) == pfx + + +def test_get_root_network(): + ipv4 = "192.144.11.0/24" + assert get_root_network(ipv4) == 192 + ipv6 = "2001:db8::/64" + assert get_root_network(ipv6) == int("2001", 16) + invalid = "not.a.network" + with pytest.raises(ValueError): + get_root_network(invalid)