Merge branch 'main' into improved-wrapping-main

Author: bitterpanda63

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -1,8 +1,5 @@
 """ This file simply exports the CloudConnectionManager class"""
 
-import time
-from aikido_zen.helpers.token import Token
-from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms
 from aikido_zen.background_process.heartbeats import send_heartbeats_every_x_secs
 from aikido_zen.background_process.routes import Routes
 from aikido_zen.ratelimiting.rate_limiter import RateLimiter
@@ -11,7 +8,7 @@
 from ..api.http_api import ReportingApiHTTP
 from ..service_config import ServiceConfig
 from ..users import Users
-from ..hostnames import Hostnames
+from aikido_zen.storage.hostnames import Hostnames
 from ..realtime.start_polling_for_changes import start_polling_for_changes
 from ..statistics import Statistics
 

aikido_zen/background_process/cloud_connection_manager/init_test.py

@@ -3,7 +3,7 @@
 from aikido_zen.background_process.api.http_api import ReportingApiHTTP
 from aikido_zen.background_process.service_config import ServiceConfig
 from aikido_zen.background_process.users import Users
-from aikido_zen.background_process.hostnames import Hostnames
+from aikido_zen.storage.hostnames import Hostnames
 from aikido_zen.ratelimiting.rate_limiter import RateLimiter
 from aikido_zen.background_process.statistics import Statistics
 from . import CloudConnectionManager

aikido_zen/background_process/cloud_connection_manager/on_detected_attack.py

@@ -1,10 +1,10 @@
 """Mainly exports on_detected_attack"""
 
 import json
+
 from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms
 from aikido_zen.helpers.logging import logger
 from aikido_zen.helpers.limit_length_metadata import limit_length_metadata
-from aikido_zen.helpers.get_ua_from_context import get_ua_from_context
 from aikido_zen.helpers.serialize_to_json import serialize_to_json
 
 
@@ -31,7 +31,7 @@ def on_detected_attack(connection_manager, attack, context, blocked, stack):
                 "method": context.method,
                 "url": context.url,
                 "ipAddress": context.remote_address,
-                "userAgent": get_ua_from_context(context),
+                "userAgent": context.get_user_agent(),
                 "body": context.body,
                 "headers": context.headers,
                 "source": context.source,

aikido_zen/background_process/cloud_connection_manager/on_detected_attack_test.py

@@ -1,6 +1,7 @@
 import pytest
 from unittest.mock import MagicMock, patch
 from .on_detected_attack import on_detected_attack
+from ...context import Context
 
 
 @pytest.fixture
@@ -16,16 +17,22 @@ def mock_connection_manager():
 
 @pytest.fixture
 def mock_context():
-    class Context:
-        method = "POST"
-        url = "http://example.com/api"
-        remote_address = "192.168.1.1"
-        body = "test body"
-        headers = {"Content-Type": "application/json"}
-        source = "test_source"
-        route = "test_route"
+    basic_wsgi_req = {
+        "REQUEST_METHOD": "GET",
+        "HTTP_HEADER_1": "header 1 value",
+        "HTTP_HEADER_2": "Header 2 value",
+        "RANDOM_VALUE": "Random value",
+        "HTTP_COOKIE": "sessionId=abc123xyz456;",
+        "wsgi.url_scheme": "http",
+        "HTTP_HOST": "localhost:8080",
+        "PATH_INFO": "/hello",
+        "QUERY_STRING": "user=JohnDoe&age=30&age=35",
+        "CONTENT_TYPE": "application/json",
+        "REMOTE_ADDR": "198.51.100.23",
+        "HTTP_USER_AGENT": "Mozilla/5.0",
+    }
 
-    return Context()
+    return Context(req=basic_wsgi_req, body=123, source="django")
 
 
 def test_on_detected_attack_no_token(mock_context):
@@ -114,3 +121,64 @@ def test_on_detected_attack_with_blocked_and_stack(
     assert attack["blocked"] is True
     assert attack["stack"] == stack
     assert mock_connection_manager.api.report.call_count == 1
+
+
+def test_on_detected_attack_request_data_and_attack_data(
+    mock_connection_manager, mock_context
+):
+    attack = {
+        "payload": {"key": "value"},
+        "metadata": {"test": "true"},
+    }
+
+    on_detected_attack(
+        mock_connection_manager, attack, mock_context, blocked=False, stack=None
+    )
+
+    # Extract the call arguments for the report method
+    _, event, _ = mock_connection_manager.api.report.call_args[0]
+
+    # Verify the request attribute in the payload
+    request_data = event["request"]
+
+    assert request_data["method"] == "GET"
+    assert request_data["url"] == "http://localhost:8080/hello"
+    assert request_data["ipAddress"] == "198.51.100.23"
+    assert request_data["body"] == 123
+    assert request_data["headers"] == {
+        "CONTENT_TYPE": "application/json",
+        "USER_AGENT": "Mozilla/5.0",
+        "COOKIE": "sessionId=abc123xyz456;",
+        "HEADER_1": "header 1 value",
+        "HEADER_2": "Header 2 value",
+        "HOST": "localhost:8080",
+    }
+    assert request_data["source"] == "django"
+    assert request_data["route"] == "/hello"
+    assert request_data["userAgent"] == "Mozilla/5.0"
+
+    attack_data = event["attack"]
+    assert attack_data["blocked"] == False
+    assert attack_data["metadata"] == {"test": "true"}
+    assert attack_data["payload"] == '{"key": "value"}'
+    assert attack_data["stack"] is None
+    assert attack_data["user"] is None
+
+
+def test_on_detected_attack_with_user(mock_connection_manager, mock_context):
+    attack = {
+        "payload": {"key": "value"},
+        "metadata": {},
+    }
+    # Simulate a user in the context
+    mock_context.user = "test_user"
+
+    on_detected_attack(
+        mock_connection_manager, attack, mock_context, blocked=False, stack=None
+    )
+
+    # Extract the call arguments for the report method
+    _, event, _ = mock_connection_manager.api.report.call_args[0]
+
+    # Verify the user is included in the attack data
+    assert event["attack"]["user"] == "test_user"

aikido_zen/background_process/commands/__init__.py

@@ -6,7 +6,6 @@
 from .user import process_user
 from .should_ratelimit import process_should_ratelimit
 from .kill import process_kill
-from .hostnames_add import process_hostnames_add
 from .statistics import process_statistics
 from .ping import process_ping
 from .sync_data import process_sync_data
@@ -18,7 +17,6 @@
     "USER": (process_user, False),
     "KILL": (process_kill, False),
     "STATISTICS": (process_statistics, False),
-    "HOSTNAMES_ADD": (process_hostnames_add, False),
     # Commands that return data :
     "SYNC_DATA": (process_sync_data, True),
     "READ_PROPERTY": (process_read_property, True),

aikido_zen/background_process/commands/hostnames_add.py

@@ -8,7 +8,7 @@ def process_sync_data(connection_manager, data, conn, queue=None):
     """
     Synchronizes data between the thread-local cache (with a TTL of usually 1 minute) and the
     background thread. Which data gets synced?
-    Thread -> BG Process : Hits, request statistics, api specs
+    Thread -> BG Process : Hits, request statistics, api specs, hostnames
     BG Process -> Thread : Routes, endpoints, bypasssed ip's, blocked users
     """
     routes = connection_manager.routes
@@ -31,6 +31,15 @@ def process_sync_data(connection_manager, data, conn, queue=None):
     # Save middleware installed :
     if data.get("middleware_installed", False):
         connection_manager.middleware_installed = True
+
+    # Sync hostnames
+    for hostnames_entry in data.get("hostnames", list()):
+        connection_manager.hostnames.add(
+            hostnames_entry["hostname"],
+            hostnames_entry["port"],
+            hostnames_entry["hits"],
+        )
+
     if connection_manager.conf.last_updated_at > 0:
         # Only report data if the config has been fetched.
         return {

aikido_zen/background_process/commands/sync_data.py

@@ -5,13 +5,15 @@
 from .sync_data import process_sync_data
 from aikido_zen.background_process.routes import Routes
 from aikido_zen.helpers.iplist import IPList
+from ...storage.hostnames import Hostnames
 
 
 @pytest.fixture
 def setup_connection_manager():
     """Fixture to set up a mock connection manager."""
     connection_manager = MagicMock()
     connection_manager.routes = Routes()
+    connection_manager.hostnames = Hostnames()
     connection_manager.conf.endpoints = ["endpoint1", "endpoint2"]
 
     connection_manager.conf.bypassed_ips = IPList()
@@ -26,6 +28,10 @@ def setup_connection_manager():
 def test_process_sync_data_initialization(setup_connection_manager):
     """Test the initialization of routes and hits update."""
     connection_manager = setup_connection_manager
+    test_hostnames = Hostnames()
+    test_hostnames.add("example2.com", 443, 15)
+    test_hostnames.add("bumblebee.com", 8080)
+
     data = {
         "current_routes": {
             "route1": {
@@ -43,6 +49,7 @@ def test_process_sync_data_initialization(setup_connection_manager):
         },
         "reqs": 10,  # Total requests to be added
         "middleware_installed": False,
+        "hostnames": test_hostnames.as_array(),
     }
 
     result = process_sync_data(connection_manager, data, None)
@@ -69,6 +76,10 @@ def test_process_sync_data_initialization(setup_connection_manager):
     assert result["routes"] == dict(connection_manager.routes.routes)
     assert result["config"] == connection_manager.conf
     assert connection_manager.middleware_installed == False
+    assert connection_manager.hostnames.as_array() == [
+        {"hits": 15, "hostname": "example2.com", "port": 443},
+        {"hits": 1, "hostname": "bumblebee.com", "port": 8080},
+    ]
 
 
 def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_manager):
@@ -114,13 +125,22 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana
     # Check that the total requests were updated
     assert connection_manager.statistics.requests["total"] == 10
     assert connection_manager.middleware_installed == True
+    assert len(connection_manager.hostnames.as_array()) == 0
     # Check that the return value is correct
     assert result == {}
 
 
-def test_process_sync_data_existing_route(setup_connection_manager):
+def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager):
     """Test updating an existing route's hit count."""
     connection_manager = setup_connection_manager
+    connection_manager.hostnames.add("example.com", 443, 200)
+    connection_manager.hostnames.add("example.org", 443)
+    connection_manager.hostnames.add("a.com", 443)
+
+    hostnames_sync = Hostnames()
+    hostnames_sync.add("example.com", 443, 15)
+    hostnames_sync.add("c.com", 443)
+
     data = {
         "current_routes": {
             "route1": {
@@ -131,6 +151,7 @@ def test_process_sync_data_existing_route(setup_connection_manager):
             }
         },
         "reqs": 5,  # Total requests to be added
+        "hostnames": hostnames_sync.as_array(),
     }
 
     # First call to initialize the route
@@ -162,6 +183,12 @@ def test_process_sync_data_existing_route(setup_connection_manager):
     # Check that the total requests were updated
     assert connection_manager.statistics.requests["total"] == 20  # 5 + 15
     assert connection_manager.middleware_installed == False
+    assert connection_manager.hostnames.as_array() == [
+        {"hits": 215, "hostname": "example.com", "port": 443},
+        {"hits": 1, "hostname": "example.org", "port": 443},
+        {"hits": 1, "hostname": "a.com", "port": 443},
+        {"hits": 1, "hostname": "c.com", "port": 443},
+    ]
 
     # Check that the return value is correct
     assert result["routes"] == dict(connection_manager.routes.routes)

aikido_zen/background_process/commands/sync_data_test.py

@@ -59,6 +59,15 @@ def get(self, route_metadata):
         key = route_to_key(route_metadata)
         return self.routes.get(key)
 
+    def get_routes_with_hits(self):
+        """Gets you all routes with a positive hits delta"""
+        result = dict()
+        for key, route in self.routes.items():
+            if route["hits_delta_since_sync"] <= 0:
+                continue  # do not add routes without a hit delta
+            result[key] = route
+        return result
+
     def clear(self):
         """Deletes all routes"""
         self.routes = {}