AIK-6314 Attack Wave Detection & Stats

aikido_zen/background_process/aikido_background_process.py

@@ -21,6 +21,7 @@
 )
 from aikido_zen.helpers.urls.get_api_url import get_api_url
 from .commands import process_incoming_command
+from .queue_helpers import ReportingQueueAttackWaveEvent
 
 EMPTY_QUEUE_INTERVAL = 5  # 5 seconds
 
@@ -102,12 +103,21 @@ def send_to_connection_manager(self, event_scheduler):
         )
         while not self.queue.empty():
             queue_attack_item = self.queue.get()
-            self.connection_manager.on_detected_attack(
-                attack=queue_attack_item[0],
-                context=queue_attack_item[1],
-                blocked=queue_attack_item[2],
-                stack=queue_attack_item[3],
-            )
+            # Queue can contain multiple types of events (attack, attack wave)
+            if isinstance(queue_attack_item, ReportingQueueAttackWaveEvent):
+                attack_wave_event: ReportingQueueAttackWaveEvent = queue_attack_item
+                self.connection_manager.on_detected_attack_wave(
+                    attack_wave_event.context, attack_wave_event.metadata
+                )
+            else:
+                # Since the queue item is not of type ReportingQueueAttackWaveEvent,
+                # we default to regular attack reporting.
+                self.connection_manager.on_detected_attack(
+                    attack=queue_attack_item[0],
+                    context=queue_attack_item[1],
+                    blocked=queue_attack_item[2],
+                    stack=queue_attack_item[3],
+                )
 
 
 def add_exit_handlers():

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -4,6 +4,7 @@
 from aikido_zen.background_process.routes import Routes
 from aikido_zen.ratelimiting.rate_limiter import RateLimiter
 from aikido_zen.helpers.logging import logger
+from .on_detected_attack_wave import on_detected_attack_wave
 from .update_firewall_lists import update_firewall_lists
 from ..api.http_api import ReportingApiHTTP
 from ..service_config import ServiceConfig
@@ -20,6 +21,10 @@
 from .update_service_config import update_service_config
 from .on_start import on_start
 from .send_heartbeat import send_heartbeat
+from aikido_zen.vulnerabilities.attack_wave_detection.attack_wave_detector import (
+    AttackWaveDetector,
+)
+from aikido_zen.context import Context
 
 
 class CloudConnectionManager:
@@ -51,6 +56,7 @@ def __init__(self, block, api, token, serverless):
         self.statistics = Statistics()
         self.ai_stats = AIStatistics()
         self.middleware_installed = False
+        self.attack_wave_detector = AttackWaveDetector()
 
         if isinstance(serverless, str) and len(serverless) == 0:
             raise ValueError("Serverless cannot be an empty string")
@@ -86,6 +92,9 @@ def on_detected_attack(self, attack, context, blocked, stack):
         """This will send something to the API when an attack is detected"""
         return on_detected_attack(self, attack, context, blocked, stack)
 
+    def on_detected_attack_wave(self, context: Context, metadata):
+        return on_detected_attack_wave(self, context, metadata)
+
     def on_start(self):
         """This will send out an Event signalling the start to the server"""
         return on_start(self)

aikido_zen/background_process/cloud_connection_manager/on_detected_attack_wave.py

@@ -0,0 +1,33 @@
+from aikido_zen.context import Context
+from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms
+from aikido_zen.helpers.logging import logger
+
+
+def on_detected_attack_wave(connection_manager, context: Context, metadata):
+    if not connection_manager.token:
+        return
+    try:
+        agent_info = connection_manager.get_manager_info()
+        payload = create_attack_wave_payload(agent_info, context, metadata)
+
+        connection_manager.api.report(
+            connection_manager.token,
+            payload,
+            connection_manager.timeout_in_sec,
+        )
+    except Exception as e:
+        logger.info("Failed to report an attack wave (%s)", str(e))
+
+
+def create_attack_wave_payload(agent_info, context: Context, metadata):
+    return {
+        "type": "detected_attack_wave",
+        "time": get_unixtime_ms(),
+        "agent": agent_info,
+        "attack": {"metadata": metadata, "user": context.user},
+        "request": {
+            "ipAddress": context.remote_address,
+            "userAgent": context.get_user_agent(),
+            "source": context.source,
+        },
+    }

aikido_zen/background_process/commands/attack_test.py

@@ -2,6 +2,7 @@
 from queue import Queue
 from unittest.mock import MagicMock
 from .attack import process_attack
+from ..queue_helpers import ReportingQueueAttackWaveEvent
 
 
 class MockCloudConnectionManager:
@@ -57,13 +58,16 @@ def test_process_attack_with_different_data_formats():
     data1 = ("injection_results", "context", True, "stacktrace")
     data2 = ("injection_results", "context", False, "stacktrace")
     data3 = ("injection_results", "context", None, "stacktrace")
+    data4 = ReportingQueueAttackWaveEvent(None, {"a": "b"})
 
     process_attack(connection_manager, data1, queue)
     process_attack(connection_manager, data2, queue)
     process_attack(connection_manager, data3, queue)
+    process_attack(connection_manager, data4, queue)
 
     # Check if all data items are added to the queue
-    assert queue.qsize() == 3
+    assert queue.qsize() == 4
     assert queue.get() == data1
     assert queue.get() == data2
     assert queue.get() == data3
+    assert queue.get() == data4

aikido_zen/background_process/commands/check_firewall_lists.py

@@ -12,34 +12,38 @@ 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}
+    Checks whether an IP is blocked and whether there is an ongoing attack wave
+    data: {"ip": string, "user-agent": string, "is_attack_wave_request" ?: bool}
+    returns -> CheckFirewallListsRes
     """
     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"}
+            return CheckFirewallListsRes(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,
-            }
+            return CheckFirewallListsRes(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,
-    }
+            return CheckFirewallListsRes(blocked=True, type="bot-blocking")
+
+    is_attack_wave_request = data.get("is_attack_wave_request", None)
+    if is_attack_wave_request and ip is not None:
+        if connection_manager.attack_wave_detector.is_attack_wave(ip):
+            return CheckFirewallListsRes(blocked=False, is_attack_wave=True)
+
+    return CheckFirewallListsRes()
+
+
+class CheckFirewallListsRes:
+    def __init__(self, blocked=False, is_attack_wave=False, type=None, reason=None):
+        self.blocked = blocked
+        self.is_attack_wave = is_attack_wave
+        self.type = type
+        self.reason = reason

aikido_zen/background_process/commands/sync_data_test.py

@@ -93,6 +93,7 @@ def test_process_sync_data_initialization(setup_connection_manager):
     assert connection_manager.statistics.get_record()["requests"] == {
         "aborted": 0,
         "attacksDetected": {"blocked": 0, "total": 5},
+        "attackWaves": {"total": 0, "blocked": 0},
         "total": 10,
         "rateLimited": 0,
     }
@@ -168,6 +169,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana
     assert connection_manager.statistics.get_record()["requests"] == {
         "aborted": 0,
         "attacksDetected": {"blocked": 0, "total": 5},
+        "attackWaves": {"total": 0, "blocked": 0},
         "total": 10,
         "rateLimited": 0,
     }
@@ -255,6 +257,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager
     assert connection_manager.statistics.get_record()["requests"] == {
         "aborted": 0,
         "attacksDetected": {"blocked": 0, "total": 10},
+        "attackWaves": {"total": 0, "blocked": 0},
         "total": 20,
         "rateLimited": 0,
     }

aikido_zen/background_process/queue_helpers.py

@@ -0,0 +1,7 @@
+from aikido_zen.context import Context
+
+
+class ReportingQueueAttackWaveEvent:
+    def __init__(self, context: Context, metadata):
+        self.context = context
+        self.metadata = metadata

aikido_zen/context/__init__.py

@@ -99,8 +99,9 @@ def __reduce__(self):
 
     def set_as_current_context(self):
         """
-        Set the current context
+        Set the current context, called every time we change the context.
         """
+        self.reset_cache()
         current_context.set(self)
 
     def set_cookies(self, cookies):
@@ -112,6 +113,9 @@ def set_body(self, body):
         except Exception as e:
             logger.debug("Exception occurred whilst setting body: %s", e)
 
+    def reset_cache(self):
+        self.parsed_userinput = {}
+
     def set_body_internal(self, body):
         """Sets the body and checks if it's possibly JSON"""
         self.body = body

aikido_zen/ratelimiting/lru_cache.py

@@ -3,7 +3,7 @@
 """
 
 from collections import OrderedDict
-from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms
+import aikido_zen.helpers.get_current_unixtime_ms as internal_time
 
 
 class LRUCache:
@@ -24,7 +24,8 @@ def get(self, key):
         if key in self.cache:
             # Check if the item is still valid based on TTL
             if (
-                get_unixtime_ms(monotonic=True) - self.cache[key]["startTime"]
+                internal_time.get_unixtime_ms(monotonic=True)
+                - self.cache[key]["startTime"]
                 < self.time_to_live_in_ms
             ):
                 return self.cache[key]["value"]  # Return the actual value
@@ -41,7 +42,7 @@ def set(self, key, value):
             self.cache.popitem(last=False)  # Remove the oldest item
         self.cache[key] = {
             "value": value,
-            "startTime": get_unixtime_ms(monotonic=True),
+            "startTime": internal_time.get_unixtime_ms(monotonic=True),
         }  # Store value and timestamp
 
     def clear(self):

aikido_zen/sources/functions/request_handler.py

@@ -7,6 +7,9 @@
 from aikido_zen.thread.thread_cache import get_cache
 from .ip_allowed_to_access_route import ip_allowed_to_access_route
 import aikido_zen.background_process.comms as c
+from ...background_process.commands.check_firewall_lists import CheckFirewallListsRes
+from ...background_process.queue_helpers import ReportingQueueAttackWaveEvent
+from ...vulnerabilities.attack_wave_detection.is_web_scanner import is_web_scanner
 
 
 def request_handler(stage, status_code=0):
@@ -49,7 +52,11 @@ def pre_response():
             message += f" (Your IP: {context.remote_address})"
         return message, 403
 
-    # Do a check on firewall lists, this happens in background because of the heavy data.
+    is_attack_wave_request = is_web_scanner(context)
+    if is_attack_wave_request:
+        logger.debug("Web scan detected for %s:%s", context.method, context.route)
+
+    # Do a check on firewall lists & attack waves, this happens in background because of the heavy data.
     # For the timeout we notice the request during heavy loads usually takes 2ms - 2.5ms, we set timeout at 10ms.
     # That way we have a very small timeout with very little risk of not blocking ips.
     comms = c.get_comms()
@@ -58,28 +65,37 @@ def pre_response():
         obj={
             "ip": context.remote_address,
             "user-agent": context.get_user_agent(),
+            "is_attack_wave_request": bool(is_attack_wave_request),
         },
         receive=True,
         timeout_in_sec=(10 / 1000),
     )
-    if not check_fw_lists_res["success"] or not check_fw_lists_res["data"]["blocked"]:
+    if not check_fw_lists_res.get("success", False):
         return
+    res: CheckFirewallListsRes = check_fw_lists_res.get("data")
 
-    block_type = check_fw_lists_res["data"]["type"]
-
-    if block_type == "allowlist":
+    if res.blocked and res.type == "allowlist":
         message = "Your IP address is not allowed."
         message += " (Your IP: " + context.remote_address + ")"
         return message, 403
-    if block_type == "blocklist":
+    if res.blocked and res.type == "blocklist":
         message = "Your IP address is blocked due to "
-        message += check_fw_lists_res["data"]["reason"]
+        message += res.reason
         message += " (Your IP: " + context.remote_address + ")"
         return message, 403
-    if block_type == "bot-blocking":
+    if res.blocked and res.type == "bot-blocking":
         msg = "You are not allowed to access this resource because you have been identified as a bot."
         return msg, 403
 
+    # We only check for attack waves after IP/Bot blocking, the reason being that if you already block the scanner
+    # There is no attack wave happening.
+    if res.is_attack_wave:
+        # Report to core & increase stats
+        comms.send_data_to_bg_process(
+            "ATTACK", ReportingQueueAttackWaveEvent(context, metadata={})
+        )
+        cache.stats.on_detected_attack_wave(blocked=res.blocked)
+
 
 def post_response(status_code):
     """Checks if the current route is useful"""