Use zen-internals for IP matching

aikido_zen/helpers/get_lib_path.py

@@ -0,0 +1,27 @@
+import os
+import platform
+
+
+def get_binary_path():
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    return os.path.normpath(os.path.join(current_dir, "../libs", get_file_name()))
+
+
+def get_file_name():
+    os_name = platform.system().lower()
+    machine = platform.machine().lower()
+    file_name = "libzen_internals_"
+
+    if "arm64" in machine or "aarch64" in machine:
+        file_name += "aarch64-"
+    elif "x86_64" in machine or "amd64" in machine:
+        file_name += "x86_64-"
+
+    if os_name == "windows":
+        file_name += "pc-windows-gnu.dll"
+    elif os_name == "darwin":
+        file_name += "apple-darwin.dylib"
+    elif os_name in ["linux", "linux2"]:
+        file_name += "unknown-linux-gnu.so"
+
+    return file_name

aikido_zen/helpers/ip_matcher/__init__.py

@@ -1,64 +1,75 @@
 import ipaddress
 
-try:
-    import pytricia
+from aikido_zen.helpers.ip_matcher.native import create_ip_matcher
+from aikido_zen.helpers.ip_matcher_fallback import IPMatcher as FallbackIPMatcher
 
-    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):
+    candidate = network.strip()
+    if candidate.startswith("["):
+        closing_bracket = candidate.rfind("]")
+        if closing_bracket < 0:
+            return network
+        candidate = candidate[1:closing_bracket]
+    if ":" not in candidate or "%" in candidate or "ffff" not in candidate.lower():
+        return network
 
-
-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)
+        mapped = ipaddress.IPv6Address(candidate).ipv4_mapped
     except ValueError:
+        return network
+    return str(mapped) if mapped else network
+
+
+def _collect_networks(networks):
+    if networks is None:
+        return ()
+
+    collected = []
+    try:
+        for network in networks:
+            if isinstance(network, str):
+                collected.append(network)
+    except Exception:
         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
-
-        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
+    return tuple(collected)
+
+
+def _create_fallback(networks):
+    try:
+        return FallbackIPMatcher(networks)
+    except Exception:
+        return FallbackIPMatcher()
+
+
+class IPMatcher:
+    def __init__(self, networks=None):
+        self._networks = _collect_networks(networks)
+        self._native = create_ip_matcher(self._networks)
+        self._fallback = (
+            _create_fallback(self._networks) if self._native is None else None
+        )
+
+    def has(self, network):
+        if not isinstance(network, str):
+            return False
+
+        try:
+            if self._has(network):
+                return True
+            mapped_ipv4 = preparse(network)
+            return mapped_ipv4 != network and self._has(mapped_ipv4)
+        except Exception:
+            return False
+
+    def _has(self, network):
+        matcher = self._native or self._fallback
+        return matcher is not None and matcher.has(network)
+
+    def is_empty(self):
+        if self._fallback is None:
+            self._fallback = _create_fallback(self._networks)
+        return self._fallback.is_empty()
+
+    def __reduce__(self):
+        return self.__class__, (self._networks,)

aikido_zen/helpers/ip_matcher/init_test.py

@@ -1,4 +1,8 @@
+import pickle
+
 import pytest
+
+import aikido_zen.helpers.ip_matcher.native as native_ip_matcher
 from . import IPMatcher
 
 
@@ -56,6 +60,8 @@ def test_with_invalid_ranges():
         "123.123.123.123/1999",
         "",
         ",,,",
+        None,
+        "\ud800",
         "192.168.0.124/32",
         "192.168.0.125/32",
         "192.168.0.170/32",
@@ -72,15 +78,20 @@ def test_with_invalid_ranges():
     assert matcher.has("10.0.0.1") == False
     assert matcher.has("192.168.0.255") == True
     assert matcher.has("") == False
+    assert matcher.has(None) == False
+    assert matcher.has("\ud800") == False
     assert matcher.has("1") == False
     assert matcher.has("192.168.0.1/32") == True
+    assert matcher.is_empty() is False
 
 
 def test_with_empty_ranges():
     input_list = []
     matcher = IPMatcher(input_list)
     assert matcher.has("192.168.2.1") == False
     assert matcher.has("foobar") == False
+    assert matcher.is_empty() is True
+    assert IPMatcher(["not-an-address", None, "\ud800"]).is_empty() is True
 
 
 def test_with_ipv6_ranges():
@@ -144,10 +155,11 @@ 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") == True
+    assert matcher.has("::ffff:123") == False
     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
+    assert IPMatcher(["127.0.0.0/8"]).has("::ffff:127.0.0.1") is True
 
 
 def test_different_cidr_ranges():
@@ -195,9 +207,42 @@ 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
+    assert IPMatcher(["::/0"]).has("::ffff:192.0.2.1") is 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
+
+
+def test_pickle_round_trip():
+    matcher = pickle.loads(pickle.dumps(IPMatcher(["192.0.2.0/24"])))
+
+    assert matcher.has("::ffff:192.0.2.42") is True
+
+
+def test_uses_fallback_when_native_symbols_are_unavailable(monkeypatch):
+    class LibraryWithoutIPMatcher:
+        pass
+
+    native_ip_matcher._load_library.cache_clear()
+    monkeypatch.setattr(
+        native_ip_matcher.ctypes,
+        "CDLL",
+        lambda _path: LibraryWithoutIPMatcher(),
+    )
+    matcher = IPMatcher(["192.0.2.0/24", "not-an-address"])
+    assert matcher.has("192.0.2.42") is True
+    assert matcher.has("198.51.100.1") is False
+
+
+def test_uses_native_matcher_when_symbols_are_available():
+    native_ip_matcher._load_library.cache_clear()
+    try:
+        native_ip_matcher._load_library()
+    except Exception:
+        pytest.skip("native IP matcher symbols are unavailable")
+
+    matcher = IPMatcher(["192.0.2.1"])
+    assert matcher.has("192.0.2.1:80") is True

aikido_zen/helpers/ip_matcher/native.py

@@ -0,0 +1,78 @@
+import ctypes
+import weakref
+from functools import lru_cache
+
+from aikido_zen.helpers.get_lib_path import get_binary_path
+
+
+class IpMatcherByteSlice(ctypes.Structure):
+    _fields_ = [("ptr", ctypes.c_char_p), ("len", ctypes.c_size_t)]
+
+
+@lru_cache(maxsize=1)
+def _load_library():
+    library = ctypes.CDLL(get_binary_path())
+    library.ip_matcher_create.argtypes = [
+        ctypes.POINTER(IpMatcherByteSlice),
+        ctypes.c_size_t,
+    ]
+    library.ip_matcher_create.restype = ctypes.c_void_p
+    library.ip_matcher_has.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_char_p,
+        ctypes.c_size_t,
+    ]
+    library.ip_matcher_has.restype = ctypes.c_int
+    library.ip_matcher_free.argtypes = [ctypes.c_void_p]
+    library.ip_matcher_free.restype = None
+    return library
+
+
+def _release_handle(library, handle):
+    try:
+        library.ip_matcher_free(handle)
+    except Exception:
+        pass
+
+
+class NativeIPMatcher:
+    def __init__(self, networks):
+        library = _load_library()
+        encoded_networks = []
+        for network in networks:
+            try:
+                encoded_networks.append(network.encode("utf-8"))
+            except UnicodeError:
+                continue
+
+        descriptors = (IpMatcherByteSlice * len(encoded_networks))(
+            *(IpMatcherByteSlice(network, len(network)) for network in encoded_networks)
+        )
+
+        handle = library.ip_matcher_create(
+            descriptors if descriptors else None, len(encoded_networks)
+        )
+        if not handle:
+            raise RuntimeError("Native IP matcher creation failed")
+
+        self._library = library
+        self._handle = handle
+        self._finalizer = weakref.finalize(
+            self, _release_handle, self._library, self._handle
+        )
+
+    def has(self, network):
+        encoded_network = network.encode("utf-8")
+        return (
+            self._library.ip_matcher_has(
+                self._handle, encoded_network, len(encoded_network)
+            )
+            == 1
+        )
+
+
+def create_ip_matcher(networks):
+    try:
+        return NativeIPMatcher(networks)
+    except Exception:
+        return None

aikido_zen/vulnerabilities/sql_injection/get_lib_path.py

@@ -1,36 +1,3 @@
 """Exports get_binary_path"""
 
-import platform
-import os
-
-
-def get_binary_path():
-    """Returns an absolute path for Rust binary file"""
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    lib_path = os.path.join(current_dir, "../../libs", get_file_name())
-    return lib_path
-
-
-def get_file_name():
-    """Gives you the file name for the binary based on platform info"""
-    os_name = platform.system().lower()
-    machine = platform.machine().lower()
-    file_name = "libzen_internals_"
-
-    # On macOS, platform.machine() returns "arm64" for Apple Silicon
-    # On Linux, platform.machine() returns "aarch64" for ARM64
-    if "arm64" in machine or "aarch64" in machine:
-        file_name += "aarch64-"
-    # On macOS, platform.machine() returns "x86_64" for Intel
-    # On Linux, platform.machine() returns "x86_64" for AMD64
-    elif "x86_64" in machine or "amd64" in machine:
-        file_name += "x86_64-"  # x86_64 or AMD64
-
-    if os_name == "windows":
-        file_name += "pc-windows-gnu.dll"  # Windows
-    elif os_name == "darwin":
-        file_name += "apple-darwin.dylib"  # macOS
-    elif os_name in ["linux", "linux2"]:
-        file_name += "unknown-linux-gnu.so"  # Linux
-
-    return file_name
+from aikido_zen.helpers.get_lib_path import get_binary_path, get_file_name