Use the zen-internals IP matcher

Makefile

@@ -44,7 +44,7 @@ benchmark:
 
 
 # Binaries cache :
-BASE_URL = https://github.com/AikidoSec/zen-internals/releases/download/v0.1.74
+BASE_URL = https://github.com/AikidoSec/zen-internals/releases/download/v0.1.75-beta.3
 FILES = \
     libzen_internals_aarch64-apple-darwin.dylib \
     libzen_internals_aarch64-apple-darwin.dylib.sha256sum \

aikido_zen/background_process/cloud_connection_manager/update_firewall_lists.py

@@ -24,9 +24,9 @@ def update_firewall_lists(connection_manager):
         if not isinstance(blocked_user_agents, str):
             return
 
-        connection_manager.firewall_lists.set_blocked_ips(blocked_ips)
-        connection_manager.firewall_lists.set_allowed_ips(allowed_ips)
-        connection_manager.firewall_lists.set_blocked_user_agents(blocked_user_agents)
+        connection_manager.firewall_lists.update(
+            blocked_ips, allowed_ips, blocked_user_agents
+        )
 
     except Exception as e:
         logger.debug("Exception in update_firewall_lists: %s", e)

aikido_zen/background_process/cloud_connection_manager/update_firewall_lists_test.py

@@ -4,6 +4,9 @@
 from aikido_zen.background_process.cloud_connection_manager import (
     CloudConnectionManager,
 )
+from aikido_zen.background_process.commands.check_firewall_lists import (
+    process_check_firewall_lists,
+)
 from .update_firewall_lists import update_firewall_lists
 
 
@@ -72,6 +75,27 @@ def fetch_firewall_lists(self, token):
         }
 
 
+class MockApiInvalidNetwork:
+    """Mock API class to simulate a list with one network that is not valid."""
+
+    def __init__(self, invalid_network):
+        self.invalid_network = invalid_network
+
+    def fetch_firewall_lists(self, token):
+        return {
+            "success": True,
+            "blockedIPAddresses": [
+                {
+                    "source": "new",
+                    "description": "Example description",
+                    "ips": ["10.0.0.0/8", self.invalid_network],
+                }
+            ],
+            "allowedIPAddresses": [],
+            "blockedUserAgents": "NewBot",
+        }
+
+
 @pytest.fixture
 def connection_manager():
     """Fixture to create a CloudConnectionManager instance with a mock API."""
@@ -107,6 +131,93 @@ def test_update_firewall_lists_success(connection_manager):
     assert not connection_manager.firewall_lists.is_user_agent_blocked(None)
 
 
+def test_update_firewall_lists_keeps_allowed_ips_when_rebuild_fails(
+    connection_manager, monkeypatch
+):
+    update_firewall_lists(connection_manager)
+    assert not connection_manager.firewall_lists.is_allowed_ip("1.1.1.1")
+
+    monkeypatch.setattr(
+        "aikido_zen.helpers.zen_internals._ip_matcher_create", lambda *args: None
+    )
+    update_firewall_lists(connection_manager)
+
+    assert not connection_manager.firewall_lists.is_allowed_ip("1.1.1.1")
+    assert connection_manager.firewall_lists.is_allowed_ip("192.168.2.50")
+    assert connection_manager.firewall_lists.is_blocked_ip("192.168.1.1")
+
+
+@pytest.mark.parametrize("invalid_network", ["not-a-network", "\udc80.1.1.1"])
+def test_update_firewall_lists_skips_invalid_networks_and_logs_them(
+    connection_manager, caplog, invalid_network
+):
+    update_firewall_lists(connection_manager)
+
+    connection_manager.api = MockApiInvalidNetwork(invalid_network)
+    update_firewall_lists(connection_manager)
+
+    assert "IP list new has 1 invalid network(s)" in caplog.text
+    assert connection_manager.firewall_lists.is_blocked_ip("10.0.0.1")
+    assert not connection_manager.firewall_lists.is_blocked_ip("192.168.1.1")
+    assert connection_manager.firewall_lists.is_allowed_ip("1.1.1.1")
+    assert connection_manager.firewall_lists.is_user_agent_blocked("NewBot")
+    assert not connection_manager.firewall_lists.is_user_agent_blocked("BadBot")
+
+
+@pytest.mark.parametrize(
+    "allowed_entry",
+    [
+        {"source": "list1", "description": "Example description"},
+        {"source": "list1", "description": "Example description", "ips": None},
+        {"source": "list1", "description": "Example description", "ips": []},
+    ],
+)
+def test_update_firewall_lists_skips_an_allowlist_entry_without_ips(
+    connection_manager, allowed_entry
+):
+    connection_manager.firewall_lists.update(
+        [{"source": "new", "description": "Example", "ips": ["10.0.0.0/8"]}],
+        [allowed_entry],
+        "NewBot",
+    )
+
+    assert connection_manager.firewall_lists.is_allowed_ip("8.8.8.8")
+    assert connection_manager.firewall_lists.is_blocked_ip("10.0.0.1")
+    assert connection_manager.firewall_lists.is_user_agent_blocked("NewBot")
+
+
+@pytest.mark.parametrize(
+    "blocked_entry",
+    [
+        {"source": "new", "ips": ["10.0.0.0/8"]},
+        {"source": "new", "description": "", "ips": ["10.0.0.0/8"]},
+    ],
+)
+def test_update_firewall_lists_blocks_ips_from_a_list_without_description(
+    connection_manager, blocked_entry
+):
+    connection_manager.firewall_lists.update([blocked_entry], [], "")
+
+    result = process_check_firewall_lists(
+        connection_manager, {"ip": "10.0.0.1", "user-agent": None}
+    )
+    assert result == {
+        "blocked": True,
+        "type": "blocklist",
+        "reason": "a firewall list",
+    }
+
+
+def test_update_firewall_lists_without_ip_matcher(connection_manager, monkeypatch):
+    monkeypatch.setattr("aikido_zen.helpers.zen_internals._ip_matcher_create", None)
+
+    update_firewall_lists(connection_manager)
+
+    assert not connection_manager.firewall_lists.is_blocked_ip("192.168.1.1")
+    assert connection_manager.firewall_lists.is_allowed_ip("1.1.1.1")
+    assert connection_manager.firewall_lists.is_user_agent_blocked("BadBot")
+
+
 def test_update_firewall_lists_no_ua(connection_manager):
     # Call the function to test
     connection_manager.api = MockApiNoUA()

aikido_zen/background_process/service_config.py

@@ -61,7 +61,11 @@ def set_endpoints(self, endpoints):
                 #  Skip empty allowlist
                 continue
 
-            endpoint["allowedIPAddresses"] = IPMatcher(endpoint["allowedIPAddresses"])
+            allowed_ips = IPMatcher(endpoint["allowedIPAddresses"])
+            if not allowed_ips.is_built():
+                # An allowlist that failed to build would block the visitors it should allow.
+                continue
+            endpoint["allowedIPAddresses"] = allowed_ips
 
     def get_endpoints(self, route_metadata):
         """

aikido_zen/background_process/service_config_test.py

@@ -1,6 +1,9 @@
 import pytest
 from .service_config import ServiceConfig
 from aikido_zen.helpers.ip_matcher import IPMatcher
+from aikido_zen.sources.functions.ip_allowed_to_access_route import (
+    ip_allowed_to_access_route,
+)
 
 
 def test_service_config_outbound_blocking_initialization():
@@ -321,6 +324,30 @@ def test_service_config_with_empty_allowlist():
     assert len(admin_endpoint["allowedIPAddresses"]) == 0
 
 
+def test_service_config_skips_allowlist_that_failed_to_build(monkeypatch):
+    monkeypatch.setattr("aikido_zen.helpers.zen_internals._ip_matcher_create", None)
+    endpoints = [
+        {
+            "graphql": False,
+            "method": "GET",
+            "route": "/admin",
+            "rate_limiting": {
+                "enabled": False,
+                "max_requests": 10,
+                "window_size_in_ms": 1000,
+            },
+            "allowedIPAddresses": ["1.2.3.4"],
+            "force_protection_off": False,
+        },
+    ]
+    service_config = ServiceConfig(endpoints, "2023-10-01", [], [], True)
+
+    route_metadata = {"method": "GET", "route": "/admin", "url": "/admin"}
+    admin_endpoints = service_config.get_endpoints(route_metadata)
+    assert admin_endpoints[0]["allowedIPAddresses"] == ["1.2.3.4"]
+    assert ip_allowed_to_access_route("8.8.8.8", route_metadata, admin_endpoints)
+
+
 def test_service_config_enabled_features():
     """Test the update_enabled_features and is_feature_enabled methods"""
     config = ServiceConfig(

aikido_zen/helpers/ip_matcher/__init__.py

@@ -1,64 +1,45 @@
-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
-
-        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
+import weakref
+
+from aikido_zen.helpers import zen_internals
+
+
+class IPMatcher:
+    def __init__(self, networks=None):
+        # json.loads can return lone surrogates such as "\udc80", and .encode() raises on
+        # them, which would fail the whole list update. Valid networks are ASCII.
+        # The library splits on newlines, so a newline would turn one entry into two.
+        networks = list(networks or [])
+        self._networks = [
+            network
+            for network in networks
+            if isinstance(network, str) and network.isascii() and "\n" not in network
+        ]
+        self._skipped_count = len(networks) - len(self._networks)
+        self._handle = zen_internals.ip_matcher_create(
+            "\n".join(self._networks).encode()
+        )
+        if self._handle is not None:
+            weakref.finalize(self, zen_internals.ip_matcher_free, self._handle)
+
+    def is_built(self):
+        return self._handle is not None
+
+    def rejected_count(self):
+        if self._handle is None:
+            return self._skipped_count
+        return self._skipped_count + zen_internals.ip_matcher_rejected_count(
+            self._handle
+        )
+
+    def has(self, network):
+        if self._handle is None:
+            return False
+        if not isinstance(network, str) or not network.isascii():
+            return False
+        return zen_internals.ip_matcher_has(self._handle, network.encode())
+
+    def is_empty(self):
+        return len(self._networks) == 0
+
+    def __reduce__(self):
+        return self.__class__, (self._networks,)