Report firewall list hits in the heartbeat

aikido_zen/background_process/api/http_api.py

@@ -35,7 +35,10 @@ def fetch_firewall_lists(self, token):
         If successful the current API returns :
         - `allowedIPAddresses` : An array with iplist entries which are the only ones allowed
         - `blockedIPAddresses` : An array with iplist entries that are blocked
+        - `monitoredIPAddresses` : An array with iplist entries that are counted but not blocked
         - `blockedUserAgents` : A string with a simple regex to match user agents with
+        - `monitoredUserAgents` : A string with a simple regex to match monitored user agents with
+        - `userAgentDetails` : An array of {key, pattern}, to report which user agent matched
         """
         try:
             res = requests.get(

aikido_zen/background_process/api/http_api_test.py

@@ -56,7 +56,24 @@ def test_local_gzip():
         "success": True,
         "allowedIPAddresses": [],
         "blockedIPAddresses": [
-            {"description": "geo restrictions", "ips": ["1.2.3.4"], "source": "geoip"}
+            {
+                "key": "geoip/blocked",
+                "description": "geo restrictions",
+                "ips": ["1.2.3.4"],
+                "source": "geoip",
+            }
+        ],
+        "monitoredIPAddresses": [
+            {
+                "key": "tor/exit_nodes",
+                "description": "Tor exit nodes",
+                "ips": ["1.2.3.4", "9.9.9.9"],
+                "source": "tor",
+            }
         ],
         "blockedUserAgents": "",
+        "monitoredUserAgents": "MonitoredTestBot",
+        "userAgentDetails": [
+            {"key": "monitoredtestbot", "pattern": "MonitoredTestBot"}
+        ],
     }

aikido_zen/background_process/cloud_connection_manager/update_firewall_lists.py

@@ -15,6 +15,10 @@ def update_firewall_lists(connection_manager):
         blocked_ips = res.get("blockedIPAddresses")
         allowed_ips = res.get("allowedIPAddresses")
         blocked_user_agents = res.get("blockedUserAgents")
+        # Only used for heartbeat stats, so these may be missing.
+        monitored_ips = res.get("monitoredIPAddresses") or []
+        monitored_user_agents = res.get("monitoredUserAgents")
+        user_agent_details = res.get("userAgentDetails") or []
 
         # Validate response
         if not res.get("success"):
@@ -27,6 +31,11 @@ def update_firewall_lists(connection_manager):
         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.set_monitored_ips(monitored_ips)
+        connection_manager.firewall_lists.set_monitored_user_agents(
+            monitored_user_agents
+        )
+        connection_manager.firewall_lists.set_user_agent_details(user_agent_details)
 
     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

@@ -179,3 +179,77 @@ def fetch_firewall_lists(self, token):
     update_firewall_lists(connection_manager)
     # No changes should be made, so we check the initial state
     assert connection_manager.conf.bypassed_ips.is_empty()
+
+
+class MockApiMonitoredLists:
+    def fetch_firewall_lists(self, token):
+        return {
+            "success": True,
+            "blockedIPAddresses": [
+                {
+                    "key": "geoip/blocked",
+                    "source": "geoip",
+                    "description": "geo restrictions",
+                    "ips": ["1.2.3.4"],
+                }
+            ],
+            "allowedIPAddresses": [],
+            "monitoredIPAddresses": [
+                {
+                    "key": "tor/exit_nodes",
+                    "source": "tor",
+                    "description": "Tor exit nodes",
+                    "ips": ["9.9.9.9"],
+                }
+            ],
+            "blockedUserAgents": "GPTBot",
+            "monitoredUserAgents": "ClaudeBot",
+            "userAgentDetails": [
+                {"key": "gptbot", "pattern": "GPTBot"},
+                {"key": "claudebot", "pattern": "ClaudeBot"},
+                {"key": "invalid", "pattern": "[abc"},
+            ],
+        }
+
+
+def test_update_firewall_lists_monitored_lists(connection_manager):
+    connection_manager.api = MockApiMonitoredLists()
+    update_firewall_lists(connection_manager)
+    firewall_lists = connection_manager.firewall_lists
+
+    assert firewall_lists.get_matching_monitored_ip_list_keys("9.9.9.9") == [
+        "tor/exit_nodes"
+    ]
+    assert firewall_lists.get_matching_blocked_ip_list_keys("1.2.3.4") == [
+        "geoip/blocked"
+    ]
+    assert firewall_lists.is_monitored_user_agent("Mozilla/5.0 (ClaudeBot/1.0)")
+    assert not firewall_lists.is_user_agent_blocked("Mozilla/5.0 (ClaudeBot/1.0)")
+    assert firewall_lists.get_matching_user_agent_keys("claudebot/1.0") == ["claudebot"]
+    assert [d["key"] for d in firewall_lists.user_agent_details] == [
+        "gptbot",
+        "claudebot",
+    ]
+
+
+def test_update_firewall_lists_without_monitored_lists(connection_manager):
+    update_firewall_lists(connection_manager)
+    firewall_lists = connection_manager.firewall_lists
+
+    assert firewall_lists.is_blocked_ip("192.168.1.1")
+    assert firewall_lists.monitored_ips == []
+    assert firewall_lists.monitored_user_agent_regex is None
+    assert firewall_lists.user_agent_details == []
+
+
+def test_update_firewall_lists_invalid_monitored_user_agents(connection_manager):
+    connection_manager.api = MockApiMonitoredLists()
+    lists = connection_manager.api.fetch_firewall_lists("token")
+    lists["monitoredUserAgents"] = "[abc"
+    connection_manager.api.fetch_firewall_lists = lambda token: lists
+    update_firewall_lists(connection_manager)
+    firewall_lists = connection_manager.firewall_lists
+
+    assert firewall_lists.monitored_user_agent_regex is None
+    assert firewall_lists.get_matching_user_agent_keys("GPTBot/1.0") == ["gptbot"]
+    assert firewall_lists.is_user_agent_blocked("GPTBot/1.0")

aikido_zen/background_process/commands/check_firewall_lists.py

@@ -3,7 +3,7 @@
 
 def process_check_firewall_lists(connection_manager, data):
     """
-    Checks whether an IP is blocked
+    Checks whether an IP is blocked, and counts the lists it matches for the heartbeat
     data: {"ip": string, "user-agent": string}
     returns -> {"blocked": boolean, "type": string, "reason": string}
     """
@@ -15,7 +15,14 @@ def process_check_firewall_lists(connection_manager, data):
 
         # Global IP Blocklist (e.g. blocking known threat actors)
         reason = connection_manager.firewall_lists.is_blocked_ip(ip)
+
+        connection_manager.statistics.on_ip_address_matches(
+            connection_manager.firewall_lists.get_matching_monitored_ip_list_keys(ip)
+        )
         if reason:
+            connection_manager.statistics.on_ip_address_matches(
+                connection_manager.firewall_lists.get_matching_blocked_ip_list_keys(ip)
+            )
             return {
                 "blocked": True,
                 "type": "blocklist",
@@ -25,7 +32,17 @@ def process_check_firewall_lists(connection_manager, data):
     user_agent = data["user-agent"]
     if user_agent is not None and isinstance(user_agent, str):
         # User agent blocking (e.g. blocking AI scrapers)
-        if connection_manager.firewall_lists.is_user_agent_blocked(user_agent):
+        is_blocked = connection_manager.firewall_lists.is_user_agent_blocked(user_agent)
+        is_monitored = connection_manager.firewall_lists.is_monitored_user_agent(
+            user_agent
+        )
+
+        if is_blocked or is_monitored:
+            keys = connection_manager.firewall_lists.get_matching_user_agent_keys(
+                user_agent
+            )
+            connection_manager.statistics.on_user_agent_matches(keys)
+        if is_blocked:
             return {
                 "blocked": True,
                 "type": "bot-blocking",

aikido_zen/sources/functions/request_handler_test.py

@@ -9,6 +9,7 @@
 from ...context import Context, current_context
 from ...helpers.headers import Headers
 from ...storage.firewall_lists import FirewallLists
+from ...storage.statistics import Statistics
 from ...vulnerabilities.attack_wave_detection.attack_wave_detector import (
     AttackWaveDetector,
 )
@@ -43,6 +44,7 @@ def __init__(self):
         self.firewall_lists = FirewallLists()
         self.conn_manager = MagicMock()
         self.conn_manager.firewall_lists = self.firewall_lists
+        self.conn_manager.statistics = Statistics()
         self.conn_manager.attack_wave_detector = AttackWaveDetector()
         self.attacks = []
 
@@ -220,6 +222,8 @@ def wrapper(*args, **kwargs):
                 kwargs["attacks"] = comms.attacks
             if "firewall_lists" in sig.parameters:
                 kwargs["firewall_lists"] = comms.firewall_lists
+            if "statistics" in sig.parameters:
+                kwargs["statistics"] = comms.conn_manager.statistics
 
             return func(*args, **kwargs)
 
@@ -978,3 +982,163 @@ def test_attack_wave_samples_json_format(firewall_lists):
     detector = attack_wave_detector_store._get_detector()
     detector.suspicious_requests_map.clear()
     detector.sent_events_map.clear()
+
+
+def set_monitored_and_blocked_lists(firewall_lists):
+    firewall_lists.set_blocked_ips(
+        [
+            {
+                "key": "geoip/blocked",
+                "source": "geoip",
+                "description": "geo restrictions",
+                "ips": ["1.2.3.4"],
+            }
+        ]
+    )
+    firewall_lists.set_monitored_ips(
+        [
+            {
+                "key": "tor/exit_nodes",
+                "source": "tor",
+                "description": "Tor exit nodes",
+                "ips": ["1.2.3.4", "9.9.9.9"],
+            },
+            {
+                "key": "crowdsec/scanners",
+                "source": "crowdsec",
+                "description": "scanners",
+                "ips": ["9.9.9.9"],
+            },
+        ]
+    )
+    firewall_lists.set_blocked_user_agents("GPTBot")
+    firewall_lists.set_monitored_user_agents("ClaudeBot")
+    firewall_lists.set_user_agent_details(
+        [
+            {"key": "gptbot", "pattern": "GPTBot"},
+            {"key": "claudebot", "pattern": "ClaudeBot"},
+        ]
+    )
+
+
+@patch_firewall_lists
+def test_monitored_ip_is_counted_but_not_blocked(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    set_context("9.9.9.9")
+
+    assert request_handler("pre_response") is None
+    assert statistics.get_record()["ipAddresses"] == {
+        "breakdown": {"tor/exit_nodes": 1, "crowdsec/scanners": 1}
+    }
+
+
+@patch_firewall_lists
+def test_blocked_ip_counts_monitored_and_blocked_lists(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    set_context("1.2.3.4")
+
+    assert request_handler("pre_response") == (
+        "Your IP address is blocked due to geo restrictions (Your IP: 1.2.3.4)",
+        403,
+    )
+    assert statistics.get_record()["ipAddresses"] == {
+        "breakdown": {"tor/exit_nodes": 1, "geoip/blocked": 1}
+    }
+
+
+@patch_firewall_lists
+def test_ip_not_in_allowlist_is_not_counted(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    firewall_lists.set_allowed_ips(
+        [{"source": "geoip", "description": "allowed", "ips": ["4.4.4.4"]}]
+    )
+    set_context("9.9.9.9")
+
+    assert request_handler("pre_response") == (
+        "Your IP address is not allowed. (Your IP: 9.9.9.9)",
+        403,
+    )
+    assert statistics.get_record()["ipAddresses"] == {"breakdown": {}}
+
+
+@patch_firewall_lists
+def test_monitored_user_agent_is_counted_but_not_blocked(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    set_context("5.5.5.5", "Mozilla/5.0 (compatible; ClaudeBot/1.0)")
+
+    assert request_handler("pre_response") is None
+    assert statistics.get_record()["userAgents"] == {"breakdown": {"claudebot": 1}}
+
+
+@patch_firewall_lists
+def test_blocked_user_agent_is_counted(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    set_context("5.5.5.5", "Mozilla/5.0 (compatible; GPTBot/1.0)")
+
+    assert request_handler("pre_response") == (
+        "You are not allowed to access this resource because you have been identified as a bot.",
+        403,
+    )
+    assert statistics.get_record()["userAgents"] == {"breakdown": {"gptbot": 1}}
+
+
+@patch_firewall_lists
+def test_other_user_agent_is_not_counted(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    set_context("5.5.5.5", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/120.0")
+
+    assert request_handler("pre_response") is None
+    assert statistics.get_record()["userAgents"] == {"breakdown": {}}
+    assert statistics.get_record()["ipAddresses"] == {"breakdown": {}}
+
+
+@patch_firewall_lists
+def test_bypassed_ip_is_not_counted(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    get_cache().config = ServiceConfig(
+        endpoints=[],
+        last_updated_at=1,
+        blocked_uids=set(),
+        bypassed_ips=["9.9.9.9"],
+        received_any_stats=False,
+    )
+    set_context("9.9.9.9", "ClaudeBot")
+
+    assert request_handler("pre_response") is None
+    assert statistics.get_record()["ipAddresses"] == {"breakdown": {}}
+    assert statistics.get_record()["userAgents"] == {"breakdown": {}}
+
+
+@patch_firewall_lists
+def test_blocked_ip_stops_before_user_agent_is_counted(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    set_context("1.2.3.4", "ClaudeBot/1.0")
+
+    assert request_handler("pre_response")[1] == 403
+    assert statistics.get_record()["userAgents"] == {"breakdown": {}}
+
+
+@patch_firewall_lists
+def test_ip_not_allowed_by_endpoint_is_not_counted(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    create_service_config()
+    set_context("9.9.9.9", "ClaudeBot/1.0")
+
+    assert request_handler("pre_response") == (
+        "Your IP address is not allowed to access this resource. (Your IP: 9.9.9.9)",
+        403,
+    )
+    assert statistics.get_record()["ipAddresses"] == {"breakdown": {}}
+    assert statistics.get_record()["userAgents"] == {"breakdown": {}}
+
+
+@patch_firewall_lists
+def test_ip_and_user_agent_counted_in_one_request(firewall_lists, statistics):
+    set_monitored_and_blocked_lists(firewall_lists)
+    set_context("9.9.9.9", "ClaudeBot/1.0")
+
+    assert request_handler("pre_response") is None
+    assert statistics.get_record()["ipAddresses"] == {
+        "breakdown": {"tor/exit_nodes": 1, "crowdsec/scanners": 1}
+    }
+    assert statistics.get_record()["userAgents"] == {"breakdown": {"claudebot": 1}}