From ebe217e01f160d00af2c57405b8694455bfe0b58 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Thu, 3 Sep 2026 19:35:27 +0200 Subject: [PATCH] Use zen-internals for IP matching --- aikido_zen/helpers/get_lib_path.py | 27 ++++ aikido_zen/helpers/ip_matcher/__init__.py | 123 ++++++++++-------- aikido_zen/helpers/ip_matcher/init_test.py | 47 ++++++- aikido_zen/helpers/ip_matcher/native.py | 78 +++++++++++ .../sql_injection/get_lib_path.py | 35 +---- poetry.lock | 14 +- pyproject.toml | 1 - 7 files changed, 220 insertions(+), 105 deletions(-) create mode 100644 aikido_zen/helpers/get_lib_path.py create mode 100644 aikido_zen/helpers/ip_matcher/native.py diff --git a/aikido_zen/helpers/get_lib_path.py b/aikido_zen/helpers/get_lib_path.py new file mode 100644 index 000000000..bf586fb2b --- /dev/null +++ b/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 diff --git a/aikido_zen/helpers/ip_matcher/__init__.py b/aikido_zen/helpers/ip_matcher/__init__.py index 98ca78a97..721ec6337 100644 --- a/aikido_zen/helpers/ip_matcher/__init__.py +++ b/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,) diff --git a/aikido_zen/helpers/ip_matcher/init_test.py b/aikido_zen/helpers/ip_matcher/init_test.py index e06a3ca4d..e06fad727 100644 --- a/aikido_zen/helpers/ip_matcher/init_test.py +++ b/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,8 +78,11 @@ 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(): @@ -81,6 +90,8 @@ def test_with_empty_ranges(): 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 diff --git a/aikido_zen/helpers/ip_matcher/native.py b/aikido_zen/helpers/ip_matcher/native.py new file mode 100644 index 000000000..8ed2a630e --- /dev/null +++ b/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 diff --git a/aikido_zen/vulnerabilities/sql_injection/get_lib_path.py b/aikido_zen/vulnerabilities/sql_injection/get_lib_path.py index b89c385a3..d6dbc5f66 100644 --- a/aikido_zen/vulnerabilities/sql_injection/get_lib_path.py +++ b/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 diff --git a/poetry.lock b/poetry.lock index 155981067..e8fb4245a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3045,18 +3045,6 @@ files = [ [package.dependencies] six = ">=1.5" -[[package]] -name = "pytricia" -version = "1.3.0" -description = "An efficient IP address storage and lookup module for Python." -optional = false -python-versions = "*" -groups = ["main"] -markers = "sys_platform != \"win32\"" -files = [ - {file = "pytricia-1.3.0.tar.gz", hash = "sha256:1c3a3d6909e10d4c9c2f0fe4542a2481e109d29aab99cc027ca7fe93f8c8853f"}, -] - [[package]] name = "pytz" version = "2025.2" @@ -4106,4 +4094,4 @@ cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and pyt [metadata] lock-version = "2.1" python-versions = ">=3.8,<3.15" -content-hash = "b0a2f0b885ebf712bda6b5cc09b705144a21352a70584f35671d2f8bfdade61a" +content-hash = "5f1249f8b42723dff85a3e5a94a64eb6241adb70eeaa109ed19745afe18317ad" diff --git a/pyproject.toml b/pyproject.toml index 5a4318222..c1be78998 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,7 +68,6 @@ regex = [ ] packaging = "^24.1" wrapt = "^1.17.2" -pytricia = { version = "^1.3.0", markers = "sys_platform != 'win32'" } [tool.poetry.group.dev.dependencies] black = "^24.4.2"