Only check firewall lists in background process

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -11,6 +11,7 @@
 from aikido_zen.storage.hostnames import Hostnames
 from ..realtime.start_polling_for_changes import start_polling_for_changes
 from ...storage.ai_statistics import AIStatistics
+from ...storage.firewall_lists import FirewallLists
 from ...storage.statistics import Statistics
 
 # Import functions :
@@ -41,6 +42,7 @@ def __init__(self, block, api, token, serverless):
             bypassed_ips=[],
             received_any_stats=True,
         )
+        self.firewall_lists = FirewallLists()
         self.rate_limiter = RateLimiter(
             max_items=5000, time_to_live_in_ms=120 * 60 * 1000  # 120 minutes
         )

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.conf.set_blocked_ips(blocked_ips)
-        connection_manager.conf.set_allowed_ips(allowed_ips)
-        connection_manager.conf.set_blocked_user_agents(blocked_user_agents)
+        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)
 
     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

@@ -5,7 +5,6 @@
     CloudConnectionManager,
 )
 from .update_firewall_lists import update_firewall_lists
-from ...helpers.is_ip_allowed_by_allowlist import is_ip_allowed_by_allowlist
 
 
 class MockApi:
@@ -93,17 +92,19 @@ def test_update_firewall_lists_success(connection_manager):
     update_firewall_lists(connection_manager)
 
     # Check that the blocked IPs were set correctly
-    assert connection_manager.conf.is_blocked_ip("192.168.1.1")
-    assert connection_manager.conf.is_blocked_ip("192.168.1.2")
+    assert connection_manager.firewall_lists.is_blocked_ip("192.168.1.1")
+    assert connection_manager.firewall_lists.is_blocked_ip("192.168.1.2")
 
     # Check that the allowed IPs were set correctly
-    assert is_ip_allowed_by_allowlist(connection_manager.conf, "192.168.1.3")
-    assert is_ip_allowed_by_allowlist(connection_manager.conf, "192.168.2.50")
+    assert connection_manager.firewall_lists.is_allowed_ip("192.168.1.3")
+    assert connection_manager.firewall_lists.is_allowed_ip("192.168.2.50")
 
     # Check that the blocked user agents were set correctly
-    assert connection_manager.conf.is_user_agent_blocked("bAdBoT test woop wop")
-    assert not connection_manager.conf.is_user_agent_blocked("")
-    assert not connection_manager.conf.is_user_agent_blocked(None)
+    assert connection_manager.firewall_lists.is_user_agent_blocked(
+        "bAdBoT test woop wop"
+    )
+    assert not connection_manager.firewall_lists.is_user_agent_blocked("")
+    assert not connection_manager.firewall_lists.is_user_agent_blocked(None)
 
 
 def test_update_firewall_lists_no_ua(connection_manager):
@@ -112,13 +113,15 @@ def test_update_firewall_lists_no_ua(connection_manager):
     update_firewall_lists(connection_manager)
 
     # Check that the blocked IPs were set correctly
-    assert connection_manager.conf.is_blocked_ip("192.168.1.1")
-    assert connection_manager.conf.is_blocked_ip("192.168.1.2")
+    assert connection_manager.firewall_lists.is_blocked_ip("192.168.1.1")
+    assert connection_manager.firewall_lists.is_blocked_ip("192.168.1.2")
 
     # Check that the blocked user agents were set correctly
-    assert not connection_manager.conf.is_user_agent_blocked("bAdBoT test woop wop")
-    assert not connection_manager.conf.is_user_agent_blocked("")
-    assert connection_manager.conf.blocked_user_agent_regex is None
+    assert not connection_manager.firewall_lists.is_user_agent_blocked(
+        "bAdBoT test woop wop"
+    )
+    assert not connection_manager.firewall_lists.is_user_agent_blocked("")
+    assert connection_manager.firewall_lists.blocked_user_agent_regex is None
 
 
 def test_update_firewall_lists_invalid_regex(connection_manager):
@@ -127,13 +130,15 @@ def test_update_firewall_lists_invalid_regex(connection_manager):
     update_firewall_lists(connection_manager)
 
     # Check that the blocked IPs were set correctly
-    assert connection_manager.conf.is_blocked_ip("192.168.1.1")
-    assert connection_manager.conf.is_blocked_ip("192.168.1.2")
+    assert connection_manager.firewall_lists.is_blocked_ip("192.168.1.1")
+    assert connection_manager.firewall_lists.is_blocked_ip("192.168.1.2")
 
     # Check that the blocked user agents were set correctly
-    assert not connection_manager.conf.is_user_agent_blocked("bAdBoT test woop wop")
-    assert not connection_manager.conf.is_user_agent_blocked("")
-    assert connection_manager.conf.blocked_user_agent_regex is None
+    assert not connection_manager.firewall_lists.is_user_agent_blocked(
+        "bAdBoT test woop wop"
+    )
+    assert not connection_manager.firewall_lists.is_user_agent_blocked("")
+    assert connection_manager.firewall_lists.blocked_user_agent_regex is None
 
 
 def test_update_firewall_lists_no_token(connection_manager):

aikido_zen/background_process/commands/__init__.py

@@ -2,6 +2,7 @@
 
 from aikido_zen.helpers.logging import logger
 from .attack import process_attack
+from .check_firewall_lists import process_check_firewall_lists
 from .read_property import process_read_property
 from .should_ratelimit import process_should_ratelimit
 from .ping import process_ping
@@ -16,6 +17,7 @@
     "READ_PROPERTY": (process_read_property, True),
     "SHOULD_RATELIMIT": (process_should_ratelimit, True),
     "PING": (process_ping, True),
+    "CHECK_FIREWALL_LISTS": (process_check_firewall_lists, True),
 }
 
 

aikido_zen/background_process/commands/check_firewall_lists.py

@@ -0,0 +1,45 @@
+"""Exports process_check_firewall_lists"""
+
+from aikido_zen.api_discovery.update_route_info import update_route_info
+from aikido_zen.background_process.cloud_connection_manager import (
+    CloudConnectionManager,
+)
+from aikido_zen.background_process.packages import PackagesStore
+from aikido_zen.helpers.logging import logger
+
+
+def process_check_firewall_lists(
+    connection_manager: CloudConnectionManager, data, conn, queue=None
+):
+    """
+    Checks whether an IP is blocked
+    data: {"ip": string, "user-agent": string}
+    returns -> {"blocked": boolean, "type": string, "reason": string}
+    """
+    ip = data["ip"]
+    if ip is not None and isinstance(ip, str):
+        # Global IP Allowlist (e.g. for geofencing)
+        if not connection_manager.firewall_lists.is_allowed_ip(ip):
+            return {"blocked": True, "type": "allowlist"}
+
+        # Global IP Blocklist (e.g. blocking known threat actors)
+        reason = connection_manager.firewall_lists.is_blocked_ip(ip)
+        if reason:
+            return {
+                "blocked": True,
+                "type": "blocklist",
+                "reason": reason,
+            }
+
+    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):
+            return {
+                "blocked": True,
+                "type": "bot-blocking",
+            }
+
+    return {
+        "blocked": False,
+    }

aikido_zen/background_process/comms.py

@@ -44,15 +44,15 @@ def __init__(self, address, key):
         global comms
         comms = self
 
-    def send_data_to_bg_process(self, action, obj, receive=False):
+    def send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0.1):
         """Try-catched send_data_to_bg_process"""
         try:
-            return self._send_data_to_bg_process(action, obj, receive)
+            return self._send_data_to_bg_process(action, obj, receive, timeout_in_sec)
         except Exception as e:
             logger.debug("Exception happened in send_data_to_bg_process : %s", e)
             return {"success": False, "error": "unknown"}
 
-    def _send_data_to_bg_process(self, action, obj, receive=False):
+    def _send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0.1):
         """
         This creates a new client for comms to the background process
         """
@@ -88,7 +88,7 @@ def target(address, key, receive, data, result_obj):
 
         # Start and join the thread for 100ms, afterwards the thread is forced to close (daemon=True)
         t.start()
-        t.join(timeout=0.1)
+        t.join(timeout=timeout_in_sec)
         if not result_obj[0]:
             logger.debug(
                 " Failure in communication to background process, %s(%s)", action, obj

aikido_zen/background_process/service_config.py

@@ -4,8 +4,6 @@
 
 from typing import Pattern
 
-import regex as re
-
 from aikido_zen.helpers.ip_matcher import IPMatcher
 from aikido_zen.helpers.match_endpoints import match_endpoints
 
@@ -26,9 +24,6 @@ def __init__(
         self.update(
             endpoints, last_updated_at, blocked_uids, bypassed_ips, received_any_stats
         )
-        self.blocked_ips = []
-        self.allowed_ips = []
-        self.blocked_user_agent_regex: Pattern = None
 
     def update(
         self,
@@ -81,39 +76,3 @@ def set_bypassed_ips(self, bypassed_ips):
     def is_bypassed_ip(self, ip):
         """Checks if the IP is on the bypass list"""
         return self.bypassed_ips.has(ip)
-
-    def set_blocked_ips(self, blocked_ip_entries):
-        self.blocked_ips = list(map(parse_ip_entry, blocked_ip_entries))
-
-    def set_allowed_ips(self, allowed_ip_entries):
-        self.allowed_ips = list(map(parse_ip_entry, allowed_ip_entries))
-
-    def is_blocked_ip(self, ip):
-        for entry in self.blocked_ips:
-            if entry["iplist"].has(ip):
-                return entry["description"]
-        return False
-
-    def set_blocked_user_agents(self, blocked_user_agents: str):
-        if not blocked_user_agents:
-            self.blocked_user_agent_regex = None
-            return
-        self.blocked_user_agent_regex = re.compile(blocked_user_agents, re.IGNORECASE)
-
-    def is_user_agent_blocked(self, ua: str):
-        if not ua:
-            return False
-        if not self.blocked_user_agent_regex:
-            return False
-        return self.blocked_user_agent_regex.search(ua)
-
-
-def parse_ip_entry(entry):
-    """
-    Converts ip entry: {"source": "example", "description": "Example description", "ips": []}
-    """
-    return {
-        "source": entry["source"],
-        "description": entry["description"],
-        "iplist": IPMatcher(entry["ips"]),
-    }

aikido_zen/background_process/service_config_test.py

@@ -130,31 +130,6 @@ def test_ip_blocking():
         bypassed_ips=["192.168.1.1", "10.0.0.0/16", "::1/128"],
         received_any_stats=True,
     )
-    config.set_blocked_ips(
-        [
-            {
-                "source": "geoip",
-                "description": "description",
-                "ips": [
-                    "1.2.3.4",
-                    "192.168.2.1/24",
-                    "fd00:1234:5678:9abc::1",
-                    "fd00:3234:5678:9abc::1/64",
-                    "5.6.7.8/32",
-                ],
-            }
-        ]
-    )
-
-    assert config.is_blocked_ip("1.2.3.4") is "description"
-    assert config.is_blocked_ip("2.3.4.5") is False
-    assert config.is_blocked_ip("192.168.2.2") is "description"
-    assert config.is_blocked_ip("fd00:1234:5678:9abc::1") is "description"
-    assert config.is_blocked_ip("fd00:1234:5678:9abc::2") is False
-    assert config.is_blocked_ip("fd00:3234:5678:9abc::1") is "description"
-    assert config.is_blocked_ip("fd00:3234:5678:9abc::2") is "description"
-    assert config.is_blocked_ip("5.6.7.8") is "description"
-    assert config.is_blocked_ip("1.2") is False
 
     assert config.is_bypassed_ip("192.168.1.1")
     assert config.is_bypassed_ip("10.0.0.1")

aikido_zen/sources/functions/request_handler.py

@@ -6,7 +6,7 @@
 from aikido_zen.helpers.logging import logger
 from aikido_zen.thread.thread_cache import get_cache
 from .ip_allowed_to_access_route import ip_allowed_to_access_route
-from ...helpers.is_ip_allowed_by_allowlist import is_ip_allowed_by_allowlist
+import aikido_zen.background_process.comms as c
 
 
 def request_handler(stage, status_code=0):
@@ -49,21 +49,32 @@ def pre_response():
             message += f" (Your IP: {context.remote_address})"
         return message, 403
 
-    # Global IP Allowlist (e.g. for geofencing)
-    if not is_ip_allowed_by_allowlist(cache.config, context.remote_address):
+    # Do a check on firewall lists, this happens in background because of the heavy data.
+    comms = c.get_comms()
+    check_fw_lists_res = comms.send_data_to_bg_process(
+        action="CHECK_FIREWALL_LISTS",
+        obj={
+            "ip": context.remote_address,
+            "user-agent": context.get_user_agent(),
+        },
+        receive=True,
+        timeout_in_sec=(20 / 1000),
+    )
+    if not check_fw_lists_res["success"] or not check_fw_lists_res["data"]["blocked"]:
+        return
+
+    block_type = check_fw_lists_res["data"]["type"]
+
+    if block_type == "allowlist":
         message = "Your IP address is not allowed."
         message += " (Your IP: " + context.remote_address + ")"
         return message, 403
-
-    # Global IP Blocklist (e.g. blocking known threat actors)
-    reason = cache.config.is_blocked_ip(context.remote_address)
-    if reason:
-        message = "Your IP address is blocked due to " + reason
+    if block_type == "blocklist":
+        message = "Your IP address is blocked due to "
+        message += check_fw_lists_res["data"]["reason"]
         message += " (Your IP: " + context.remote_address + ")"
         return message, 403
-
-    # User agent blocking (e.g. blocking AI scrapers)
-    if cache.config.is_user_agent_blocked(context.get_user_agent()):
+    if block_type == "bot-blocking":
         msg = "You are not allowed to access this resource because you have been identified as a bot."
         return msg, 403