From 820af4c0e41c7436ed1833f4c265d207fe56716f Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Wed, 4 Dec 2024 13:38:36 +0100 Subject: [PATCH 1/5] Linting --- aikido_zen/context/__init__.py | 6 +++--- aikido_zen/vulnerabilities/init_test.py | 3 --- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/aikido_zen/context/__init__.py b/aikido_zen/context/__init__.py index 251221e6..3ced8b78 100644 --- a/aikido_zen/context/__init__.py +++ b/aikido_zen/context/__init__.py @@ -50,9 +50,9 @@ def __init__(self, context_obj=None, body=None, req=None, source=None): self.set_body(body) # Parse WSGI/ASGI/... request : - self.cookies = self.method = self.remote_address = self.query = self.headers = ( - self.url - ) = None + self.cookies = ( + self.method + ) = self.remote_address = self.query = self.headers = self.url = None if source in WSGI_SOURCES: set_wsgi_attributes_on_context(self, req) elif source in ASGI_SOURCES: diff --git a/aikido_zen/vulnerabilities/init_test.py b/aikido_zen/vulnerabilities/init_test.py index 0c1ac9fc..1b79d7cb 100644 --- a/aikido_zen/vulnerabilities/init_test.py +++ b/aikido_zen/vulnerabilities/init_test.py @@ -89,7 +89,6 @@ def test_lifecycle_cache_bypassed_ip(caplog, get_context): def test_sql_injection(caplog, get_context, monkeypatch): - get_context.set_as_current_context() cache = ThreadCache() monkeypatch.setenv("AIKIDO_BLOCKING", "1") @@ -114,7 +113,6 @@ def test_sql_injection_with_route_params(caplog, get_context, monkeypatch): def test_sql_injection_with_comms(caplog, get_context, monkeypatch): - get_context.set_as_current_context() cache = ThreadCache() cache.last_renewal = 9999999999999999999999 @@ -140,7 +138,6 @@ def test_sql_injection_with_comms(caplog, get_context, monkeypatch): def test_ssrf_with_comms_hostnames_add(caplog, get_context, monkeypatch): - get_context.set_as_current_context() monkeypatch.setenv("AIKIDO_BLOCKING", "1") with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: From ea8943ebba574a0a2f96d9866ed45c7e5cea08fb Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Wed, 4 Dec 2024 13:38:52 +0100 Subject: [PATCH 2/5] Send middleware installed in heartbeat --- .../cloud_connection_manager/__init__.py | 1 + .../cloud_connection_manager/send_heartbeat.py | 1 + aikido_zen/background_process/commands/sync_data.py | 5 +++++ aikido_zen/middleware/should_block_request.py | 4 ++++ aikido_zen/thread/thread_cache.py | 7 ++++++- 5 files changed, 17 insertions(+), 1 deletion(-) diff --git a/aikido_zen/background_process/cloud_connection_manager/__init__.py b/aikido_zen/background_process/cloud_connection_manager/__init__.py index 5e253c0a..1cd0e28b 100644 --- a/aikido_zen/background_process/cloud_connection_manager/__init__.py +++ b/aikido_zen/background_process/cloud_connection_manager/__init__.py @@ -52,6 +52,7 @@ def __init__(self, block, api, token, serverless): self.statistics = Statistics( max_perf_samples_in_mem=5000, max_compressed_stats_in_mem=100 ) + self.middleware_installed = False if isinstance(serverless, str) and len(serverless) == 0: raise ValueError("Serverless cannot be an empty string") diff --git a/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py b/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py index 6028a964..4b97a388 100644 --- a/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py +++ b/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py @@ -30,6 +30,7 @@ def send_heartbeat(connection_manager): "hostnames": outgoing_domains, "routes": routes, "users": users, + "middlewareInstalled": connection_manager.middleware_installed, }, connection_manager.timeout_in_sec, ) diff --git a/aikido_zen/background_process/commands/sync_data.py b/aikido_zen/background_process/commands/sync_data.py index 36747e1c..63757a66 100644 --- a/aikido_zen/background_process/commands/sync_data.py +++ b/aikido_zen/background_process/commands/sync_data.py @@ -25,7 +25,12 @@ def process_sync_data(connection_manager, data, conn, queue=None): # Update API Spec : update_route_info(route["apispec"], existing_route) + # Save request data : connection_manager.statistics.requests["total"] += data.get("reqs", 0) + + # Save middleware installed : + if data.get("middleware_installed", False): + connection_manager.middleware_installed = True if connection_manager.conf.last_updated_at > 0: # Only report data if the config has been fetched. return { diff --git a/aikido_zen/middleware/should_block_request.py b/aikido_zen/middleware/should_block_request.py index 134a1276..85b57914 100644 --- a/aikido_zen/middleware/should_block_request.py +++ b/aikido_zen/middleware/should_block_request.py @@ -17,11 +17,15 @@ def should_block_request(): cache = get_cache() if not context or not cache: return {"block": False} + context.executed_middleware = ( True # Update context with middleware execution set to true ) context.set_as_current_context() + # Make sure we set middleware installed to true (reports back to core) : + cache.middleware_installed = True + # Blocked users: if context.user and cache.is_user_blocked(context.user["id"]): return {"block": True, "type": "blocked", "trigger": "user"} diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index 0ba884f7..e84c3b5e 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -68,6 +68,7 @@ def reset(self): ) self.reqs = 0 self.last_renewal = 0 + self.middleware_installed = False def renew(self): """ @@ -77,7 +78,11 @@ def renew(self): return res = comms.get_comms().send_data_to_bg_process( action="SYNC_DATA", - obj={"current_routes": dict(self.routes.routes), "reqs": self.reqs}, + obj={ + "current_routes": dict(self.routes.routes), + "reqs": self.reqs, + "middleware_installed": self.middleware_installed, + }, receive=True, ) From 8964a996ec785504b10438feedcb556871ec0d69 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Wed, 4 Dec 2024 13:43:35 +0100 Subject: [PATCH 3/5] Linting --- aikido_zen/context/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/aikido_zen/context/__init__.py b/aikido_zen/context/__init__.py index 3ced8b78..251221e6 100644 --- a/aikido_zen/context/__init__.py +++ b/aikido_zen/context/__init__.py @@ -50,9 +50,9 @@ def __init__(self, context_obj=None, body=None, req=None, source=None): self.set_body(body) # Parse WSGI/ASGI/... request : - self.cookies = ( - self.method - ) = self.remote_address = self.query = self.headers = self.url = None + self.cookies = self.method = self.remote_address = self.query = self.headers = ( + self.url + ) = None if source in WSGI_SOURCES: set_wsgi_attributes_on_context(self, req) elif source in ASGI_SOURCES: From a0a1a52f1121cbb87d863f2b57fbc2a476d1d5d4 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Wed, 4 Dec 2024 13:46:49 +0100 Subject: [PATCH 4/5] Add some unit tests --- aikido_zen/background_process/commands/sync_data_test.py | 7 ++++++- aikido_zen/middleware/init_test.py | 5 +++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/aikido_zen/background_process/commands/sync_data_test.py b/aikido_zen/background_process/commands/sync_data_test.py index 1f1899ea..ad89f1ff 100644 --- a/aikido_zen/background_process/commands/sync_data_test.py +++ b/aikido_zen/background_process/commands/sync_data_test.py @@ -16,6 +16,7 @@ def setup_connection_manager(): connection_manager.conf.blocked_uids = ["user1", "user2"] connection_manager.conf.last_updated_at = 200 connection_manager.statistics.requests = {"total": 0} # Initialize total requests + connection_manager.middleware_installed = False return connection_manager @@ -38,6 +39,7 @@ def test_process_sync_data_initialization(setup_connection_manager): }, }, "reqs": 10, # Total requests to be added + "middleware_installed": False, } result = process_sync_data(connection_manager, data, None) @@ -63,6 +65,7 @@ def test_process_sync_data_initialization(setup_connection_manager): # Check that the return value is correct assert result["routes"] == dict(connection_manager.routes.routes) assert result["config"] == connection_manager.conf + assert connection_manager.middleware_installed == False def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_manager): @@ -85,6 +88,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana }, }, "reqs": 10, # Total requests to be added + "middleware_installed": True, } result = process_sync_data(connection_manager, data, None) @@ -106,7 +110,7 @@ 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 # Check that the return value is correct assert result == {} @@ -154,6 +158,7 @@ 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 # Check that the return value is correct assert result["routes"] == dict(connection_manager.routes.routes) diff --git a/aikido_zen/middleware/init_test.py b/aikido_zen/middleware/init_test.py index 5a30e9ec..dea1e69a 100644 --- a/aikido_zen/middleware/init_test.py +++ b/aikido_zen/middleware/init_test.py @@ -56,12 +56,14 @@ def test_with_context_with_cache(): thread_cache.config.blocked_uids = ["123"] assert get_current_context().executed_middleware == False + assert thread_cache.middleware_installed == False assert should_block_request() == { "block": True, "trigger": "user", "type": "blocked", } assert get_current_context().executed_middleware == True + assert thread_cache.middleware_installed == True thread_cache.config.blocked_uids = [] assert should_block_request() == {"block": False} @@ -69,6 +71,7 @@ def test_with_context_with_cache(): thread_cache.config.blocked_uids = ["23", "234", "456"] assert should_block_request() == {"block": False} assert get_current_context().executed_middleware == True + assert thread_cache.middleware_installed == True def test_cache_comms_with_endpoints(): @@ -88,11 +91,13 @@ def test_cache_comms_with_endpoints(): } ] assert get_current_context().executed_middleware == False + assert thread_cache.middleware_installed == False with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: mock_get_comms.return_value = None # Set the return value of get_comms assert should_block_request() == {"block": False} assert get_current_context().executed_middleware == True + assert thread_cache.middleware_installed == True with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: mock_comms = MagicMock() From 0a85024a2d3e36672c393a61020322a9df21ff4d Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Wed, 4 Dec 2024 13:56:05 +0100 Subject: [PATCH 5/5] Add extra unit tests for cloud connection manager --- .../cloud_connection_manager/init_test.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 aikido_zen/background_process/cloud_connection_manager/init_test.py diff --git a/aikido_zen/background_process/cloud_connection_manager/init_test.py b/aikido_zen/background_process/cloud_connection_manager/init_test.py new file mode 100644 index 00000000..c3699b8f --- /dev/null +++ b/aikido_zen/background_process/cloud_connection_manager/init_test.py @@ -0,0 +1,48 @@ +import pytest +from aikido_zen.helpers.token import Token +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.ratelimiting.rate_limiter import RateLimiter +from aikido_zen.background_process.statistics import Statistics +from . import CloudConnectionManager + + +@pytest.fixture +def setup_cloud_connection_manager(): + block = None # Replace with appropriate mock or object if needed + api = ReportingApiHTTP(None) # Mock or create an instance of ReportingApiHTTP + token = Token("AIK_TOKEN_TEST") # Mock or create an instance of Token + serverless = "some_value" # Valid serverless value + return CloudConnectionManager(block, api, token, serverless) + + +def test_cloud_connection_manager_initialization(setup_cloud_connection_manager): + manager = setup_cloud_connection_manager + + # Check that instance variables are initialized correctly + assert manager.block is None + assert isinstance(manager.api, ReportingApiHTTP) + assert isinstance(manager.token, Token) + assert len(manager.routes) == 0 + assert isinstance(manager.hostnames, Hostnames) + assert isinstance(manager.conf, ServiceConfig) + assert isinstance(manager.rate_limiter, RateLimiter) + assert isinstance(manager.users, Users) + assert isinstance(manager.statistics, Statistics) + assert manager.middleware_installed is False + assert manager.serverless == "some_value" + + +def test_cloud_connection_manager_empty_serverless(): + block = None # Replace with appropriate mock or object if needed + api = ReportingApiHTTP(None) # Mock or create an instance of ReportingApiHTTP + token = Token("Hellow") # Mock or create an instance of Token + serverless = "" # Invalid serverless value + + with pytest.raises(ValueError, match="Serverless cannot be an empty string"): + CloudConnectionManager(block, api, token, serverless) + + +# Additional tests can be added here for other edge cases or scenarios