IPMatcher: Switch node logic out in favour of pytricia

aikido_zen/background_process/service_config.py

@@ -67,9 +67,7 @@ def get_endpoints(self, route_metadata):
 
     def set_bypassed_ips(self, bypassed_ips):
         """Creates an IPMatcher from the given bypassed ip set"""
-        self.bypassed_ips = IPMatcher()
-        for ip in bypassed_ips:
-            self.bypassed_ips.add(ip)
+        self.bypassed_ips = IPMatcher(bypassed_ips)
 
     def is_bypassed_ip(self, ip):
         """Checks if the IP is on the bypass list"""

aikido_zen/helpers/ip_matcher/__init__.py

@@ -1,49 +1,43 @@
-"""
-Based on https://github.com/demskie/netparser
-MIT License - Copyright (c) 2019 alex
-"""
+import ipaddress
 
-from .shared import parse_base_network, sort_networks, summarize_sorted_networks
-from .sort import binary_search_for_insertion_index
+import pytricia
+
+
+def preparse(network: str) -> str:
+    # Remove the brackets around IPv6 addresses if they are there.
+    network = network.strip("[]")
+    try:
+        ip = ipaddress.IPv6Address(network)
+        if ip.ipv4_mapped:
+            return str(ip.ipv4_mapped)
+    except ValueError:
+        pass
+    return network
 
 
 class IPMatcher:
     def __init__(self, networks=None):
-        self.sorted = []
+        self.trie = pytricia.PyTricia(128)
         if networks is not None:
-            subnets = []
             for s in networks:
-                net = parse_base_network(s, False)
-                if net and net.is_valid():
-                    subnets.append(net)
-            sort_networks(subnets)
-            self.sorted = summarize_sorted_networks(subnets)
+                self._add(s)
+        # We freeze in constructor ensuring that after initialization the IPMatcher is always frozen.
+        self.trie.freeze()
 
     def has(self, network):
-        """
-        Checks if the given IP address is in the list of networks.
-        """
-        net = parse_base_network(network, False)
-        if not net or not net.is_valid():
+        try:
+            return self.trie.get(preparse(network)) is not None
+        except ValueError:
             return False
-        idx = binary_search_for_insertion_index(net, self.sorted)
-        if idx < 0:
-            return False
-        if idx < len(self.sorted) and self.sorted[idx].contains(net):
-            return True
-        if idx - 1 >= 0 and self.sorted[idx - 1].contains(net):
-            return True
-        return False
-
-    def add(self, network):
-        net = parse_base_network(network, False)
-        if not net or not net.is_valid():
-            return self
-        idx = binary_search_for_insertion_index(net, self.sorted)
-        if idx < len(self.sorted) and self.sorted[idx].compare(net) == 0:
-            return self
-        self.sorted.insert(idx, net)
+
+    def _add(self, network):
+        try:
+            self.trie[preparse(network)] = True
+        except ValueError:
+            pass
+        except SystemError:
+            pass
         return self
 
     def is_empty(self):
-        return len(self.sorted) == 0
+        return len(self.trie) == 0

aikido_zen/helpers/ip_matcher/init_test.py

@@ -130,8 +130,7 @@ def test_add_ips_later():
     matcher = IPMatcher()
     assert matcher.has("2001:db8::0") == False
     assert matcher.has("2002:db8::1") == False
-    for ip in input_list:
-        matcher.add(ip)
+    matcher = IPMatcher(input_list)
     assert matcher.has("2001:db8::1") == False
     assert matcher.has("2001:db8::0") == False
     assert matcher.has("2002:db8::1") == True
@@ -145,7 +144,7 @@ def test_strange_ips():
     matcher = IPMatcher(input_list)
     assert matcher.has("::ffff:0.0.0.0") == True
     assert matcher.has("::ffff:127.0.0.1") == True
-    assert matcher.has("::ffff:123") == False
+    assert matcher.has("::ffff:123") == True
     assert matcher.has("2001:db8::1") == False
     assert matcher.has("[::ffff:0.0.0.0]") == True
     assert matcher.has("::ffff:0:0:0:0") == True