Add ability to filter get_network_list() by ipv4 or ipv6, and add tests

This commit is contained in:
Pim van Pelt
2023-01-16 12:15:24 +00:00
parent 4e2354c3d8
commit 8a7c690ee5
6 changed files with 61 additions and 11 deletions

View File

@ -165,17 +165,25 @@ def is_ip(ip_string):
return False return False
def get_network_list(yaml, network_string): def get_network_list(yaml, network_string, want_ipv4=True, want_ipv6=True):
"""Return the full list of source or destination address(es). This function resolves the """Return the full list of source or destination address(es). This function resolves the
'source' or 'destination' field, which can either be an IP address, a Prefix, or the name 'source' or 'destination' field, which can either be an IP address, a Prefix, or the name
of a Prefix List. It returns a list of ip_network() objects, including prefix. IP addresses of a Prefix List. It returns a list of ip_network() objects, including prefix. IP addresses
will receive prefixlen /32 or /128.""" will receive prefixlen /32 or /128. Optionally, want_ipv4 or want_ipv6 can be set to False
to filter the list."""
ret = []
if is_ip(network_string): if is_ip(network_string):
ipn = ipaddress.ip_network(network_string, strict=False) ipn = ipaddress.ip_network(network_string, strict=False)
return [ipn] if ipn.version == 4 and want_ipv4:
ret = [ipn]
if ipn.version == 6 and want_ipv6:
ret = [ipn]
return ret
return prefixlist.get_network_list(yaml, network_string) return prefixlist.get_network_list(
yaml, network_string, want_ipv4=want_ipv4, want_ipv6=want_ipv6
)
def get_protocol(protostring): def get_protocol(protostring):

View File

@ -36,17 +36,20 @@ def get_by_name(yaml, plname):
return None, None return None, None
def get_network_list(yaml, plname): def get_network_list(yaml, plname, want_ipv4=True, want_ipv6=True):
"""Returns a list of 0 or more ip_network elements, that represent the members """Returns a list of 0 or more ip_network elements, that represent the members
in a prefixlist of given name. Return the empty list if the prefixlist doesn't in a prefixlist of given name. Return the empty list if the prefixlist doesn't
exist""" exist. Optionally, want_ipv4 or want_ipv6 can be set to False to filter the list."""
ret = [] ret = []
plname, pl = get_by_name(yaml, plname) plname, pl = get_by_name(yaml, plname)
if not pl: if not pl:
return ret return ret
for m in pl["members"]: for m in pl["members"]:
ipn = ipaddress.ip_network(m, strict=False) ipn = ipaddress.ip_network(m, strict=False)
ret.append(ipn) if ipn.version == 4 and want_ipv4:
ret.append(ipn)
if ipn.version == 6 and want_ipv6:
ret.append(ipn)
return ret return ret

View File

@ -135,7 +135,19 @@ class TestACLMethods(unittest.TestCase):
l = acl.get_network_list(self.cfg, "trusted") l = acl.get_network_list(self.cfg, "trusted")
self.assertIsInstance(l, list) self.assertIsInstance(l, list)
self.assertEquals(4, len(l)) self.assertEquals(5, len(l))
l = acl.get_network_list(self.cfg, "trusted", want_ipv6=False)
self.assertIsInstance(l, list)
self.assertEquals(2, len(l))
l = acl.get_network_list(self.cfg, "trusted", want_ipv4=False)
self.assertIsInstance(l, list)
self.assertEquals(3, len(l))
l = acl.get_network_list(self.cfg, "trusted", want_ipv4=False, want_ipv6=False)
self.assertIsInstance(l, list)
self.assertEquals(0, len(l))
l = acl.get_network_list(self.cfg, "pl-notexist") l = acl.get_network_list(self.cfg, "pl-notexist")
self.assertIsInstance(l, list) self.assertIsInstance(l, list)

View File

@ -41,7 +41,7 @@ class TestACLMethods(unittest.TestCase):
def test_count(self): def test_count(self):
v4, v6 = prefixlist.count(self.cfg, "trusted") v4, v6 = prefixlist.count(self.cfg, "trusted")
self.assertEqual(2, v4) self.assertEqual(2, v4)
self.assertEqual(2, v6) self.assertEqual(3, v6)
v4, v6 = prefixlist.count(self.cfg, "empty") v4, v6 = prefixlist.count(self.cfg, "empty")
self.assertEqual(0, v4) self.assertEqual(0, v4)
@ -57,7 +57,7 @@ class TestACLMethods(unittest.TestCase):
self.assertEqual(0, prefixlist.count_ipv4(self.cfg, "pl-noexist")) self.assertEqual(0, prefixlist.count_ipv4(self.cfg, "pl-noexist"))
def test_count_ipv6(self): def test_count_ipv6(self):
self.assertEqual(2, prefixlist.count_ipv6(self.cfg, "trusted")) self.assertEqual(3, prefixlist.count_ipv6(self.cfg, "trusted"))
self.assertEqual(0, prefixlist.count_ipv6(self.cfg, "empty")) self.assertEqual(0, prefixlist.count_ipv6(self.cfg, "empty"))
self.assertEqual(0, prefixlist.count_ipv6(self.cfg, "pl-noexist")) self.assertEqual(0, prefixlist.count_ipv6(self.cfg, "pl-noexist"))
@ -79,7 +79,21 @@ class TestACLMethods(unittest.TestCase):
def test_get_network_list(self): def test_get_network_list(self):
l = prefixlist.get_network_list(self.cfg, "trusted") l = prefixlist.get_network_list(self.cfg, "trusted")
self.assertIsInstance(l, list) self.assertIsInstance(l, list)
self.assertEquals(4, len(l)) self.assertEquals(5, len(l))
l = prefixlist.get_network_list(self.cfg, "trusted", want_ipv6=False)
self.assertIsInstance(l, list)
self.assertEquals(2, len(l))
l = prefixlist.get_network_list(self.cfg, "trusted", want_ipv4=False)
self.assertIsInstance(l, list)
self.assertEquals(3, len(l))
l = prefixlist.get_network_list(
self.cfg, "trusted", want_ipv4=False, want_ipv6=False
)
self.assertIsInstance(l, list)
self.assertEquals(0, len(l))
l = prefixlist.get_network_list(self.cfg, "pl-notexist") l = prefixlist.get_network_list(self.cfg, "pl-notexist")
self.assertIsInstance(l, list) self.assertIsInstance(l, list)

View File

@ -1,7 +1,19 @@
prefixlists:
trusted:
members:
- 192.0.2.1
- 192.0.2.0/24
- 2001:db8::1
- 2001:db8::/64
- 2001:db8::/48
acls: acls:
acl01: acl01:
description: "Test ACL #1" description: "Test ACL #1"
terms: terms:
- description: "Allow a Prefixlist"
action: permit
source: trusted
- description: "Allow a specific IPv6 TCP flow" - description: "Allow a specific IPv6 TCP flow"
action: permit action: permit
source: 2001:db8::/64 source: 2001:db8::/64

View File

@ -6,6 +6,7 @@ prefixlists:
- 192.0.2.0/24 - 192.0.2.0/24
- 2001:db8::1 - 2001:db8::1
- 2001:db8::/64 - 2001:db8::/64
- 2001:db8::/48
deny-all: deny-all:
description: "Default for IPv4 and IPv6" description: "Default for IPv4 and IPv6"
members: members: