Use pytricia for ip matching on non-windows machines

aikido_zen/background_process/service_config.py

@@ -69,9 +69,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,64 @@
-"""
-Based on https://github.com/demskie/netparser
-MIT License - Copyright (c) 2019 alex
-"""
-
-from .shared import parse_base_network, sort_networks, summarize_sorted_networks
-from .sort import binary_search_for_insertion_index
-
-
-class IPMatcher:
-    def __init__(self, networks=None):
-        self.sorted = []
-        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)
-
-    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():
-            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:
+import ipaddress
+
+try:
+    import pytricia
+
+    PYTRICIA_AVAILABLE = True
+except ImportError:
+    PYTRICIA_AVAILABLE = False
+    from aikido_zen.helpers.logging import logger
+
+    logger.warning(
+        "pytricia is not available. This happens on windows devices where pytricia is not supported yet."
+        "Using fallback, this may result in slower performance."
+        "You can try to install pytricia for better performance: pip install 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
+
+
+if PYTRICIA_AVAILABLE:
+
+    class IPMatcher:
+        def __init__(self, networks=None):
+            self.trie = pytricia.PyTricia(128)
+            if networks is not None:
+                for s in networks:
+                    self._add(s)
+            # We freeze in constructor ensuring that after initialization the IPMatcher is always frozen.
+            self.trie.freeze()
+
+        def has(self, network):
+            try:
+                return self.trie.get(preparse(network)) is not None
+            except ValueError:
+                return False
+
+        def _add(self, network):
+            try:
+                self.trie[preparse(network)] = True
+            except ValueError:
+                pass
+            except SystemError:
+                # SystemError's have been known to occur in the PyTricia library (see issue #34 e.g.),
+                # best to play it safe and catch these errors.
+                pass
             return self
-        self.sorted.insert(idx, net)
-        return self
 
-    def is_empty(self):
-        return len(self.sorted) == 0
+        def is_empty(self):
+            return len(self.trie) == 0
+
+else:
+    # Fallback to pure Python implementation - this happens on windows machines since pytricia is not
+    # fully supported there.
+    from aikido_zen.helpers.ip_matcher_fallback import IPMatcher  # noqa: F401

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
@@ -196,3 +195,9 @@ def test_allow_all_ips():
     assert matcher.has("10.0.0.1") == True
     assert matcher.has("10.0.0.255") == True
     assert matcher.has("192.168.1.1") == True
+
+
+def test_edge_cases():
+    matcher1 = IPMatcher(["224.0.0.0/4"])
+    assert matcher1.has("224.0.0.1") == True
+    assert matcher1.has("240.0.0.0") == False

aikido_zen/helpers/ip_matcher_fallback/__init__.py

@@ -0,0 +1,49 @@
+"""
+Based on https://github.com/demskie/netparser
+MIT License - Copyright (c) 2019 alex
+"""
+
+from .shared import parse_base_network, sort_networks, summarize_sorted_networks
+from .sort import binary_search_for_insertion_index
+
+
+class IPMatcher:
+    def __init__(self, networks=None):
+        self.sorted = []
+        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)
+
+    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():
+            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)
+        return self
+
+    def is_empty(self):
+        return len(self.sorted) == 0

aikido_zen/helpers/ip_matcher/address.py → aikido_zen/helpers/ip_matcher_fallback/address.py

@@ -3,7 +3,7 @@
 MIT License - Copyright (c) 2019 alex
 """
 
-from aikido_zen.helpers.ip_matcher import parse
+from aikido_zen.helpers.ip_matcher_fallback import parse
 
 # Constants
 BEFORE = -1