From 4e98595240765a96303fe0798443ec0349900100 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 13 Jun 2025 11:18:07 +0200 Subject: [PATCH 1/7] Add on_rate_limit() to statistics object --- aikido_zen/storage/statistics/__init__.py | 7 ++++ aikido_zen/storage/statistics/init_test.py | 47 ++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/aikido_zen/storage/statistics/__init__.py b/aikido_zen/storage/statistics/__init__.py index bedf42ca..41daf0c3 100644 --- a/aikido_zen/storage/statistics/__init__.py +++ b/aikido_zen/storage/statistics/__init__.py @@ -12,6 +12,7 @@ def __init__(self): self.total_hits = 0 self.attacks_detected = 0 self.attacks_blocked = 0 + self.rate_limited_hits = 0 self.started_at = t.get_unixtime_ms() self.operations = Operations() @@ -19,6 +20,7 @@ def clear(self): self.total_hits = 0 self.attacks_detected = 0 self.attacks_blocked = 0 + self.rate_limited_hits = 0 self.started_at = t.get_unixtime_ms() self.operations.clear() @@ -31,6 +33,9 @@ def on_detected_attack(self, blocked, operation): self.attacks_blocked += 1 self.operations.on_detected_attack(blocked, operation) + def on_rate_limit(self): + self.rate_limited_hits += 1 + def get_record(self): current_time = t.get_unixtime_ms() return { @@ -38,6 +43,7 @@ def get_record(self): "endedAt": current_time, "requests": { "total": self.total_hits, + "rate_limited": self.rate_limited_hits, "aborted": 0, # statistic currently not in use "attacksDetected": { "total": self.attacks_detected, @@ -50,6 +56,7 @@ def get_record(self): def import_from_record(self, record): attacks_detected = record.get("requests", {}).get("attacksDetected", {}) self.total_hits += record.get("requests", {}).get("total", 0) + self.rate_limited_hits += record.get("requests", {}).get("rate_limited", 0) self.attacks_detected += attacks_detected.get("total", 0) self.attacks_blocked += attacks_detected.get("blocked", 0) self.operations.update(record.get("operations", {})) diff --git a/aikido_zen/storage/statistics/init_test.py b/aikido_zen/storage/statistics/init_test.py index 5ca481d0..25737678 100644 --- a/aikido_zen/storage/statistics/init_test.py +++ b/aikido_zen/storage/statistics/init_test.py @@ -68,6 +68,8 @@ def test_get_record(monkeypatch): stats = Statistics() stats.total_hits = 10 + stats.on_rate_limit() + stats.on_rate_limit() stats.operations.register_call("test.test", "nosql_op") stats.on_detected_attack(blocked=True, operation="test.test") stats.attacks_detected = 5 @@ -77,6 +79,7 @@ def test_get_record(monkeypatch): assert record["startedAt"] == stats.started_at assert record["endedAt"] == mock_time assert record["requests"]["total"] == 10 + assert record["requests"]["rate_limited"] == 2 assert record["requests"]["aborted"] == 0 assert record["requests"]["attacksDetected"]["total"] == 5 assert record["requests"]["attacksDetected"]["blocked"] == 3 @@ -97,6 +100,7 @@ def test_import_from_record(): record = { "requests": { "total": 10, + "rate_limited": 5, "attacksDetected": { "total": 5, "blocked": 3, @@ -117,6 +121,7 @@ def test_import_from_record(): } stats.import_from_record(record) assert stats.total_hits == 10 + assert stats.rate_limited_hits == 5 assert stats.attacks_detected == 5 assert stats.attacks_blocked == 3 assert stats.operations == { @@ -152,6 +157,7 @@ def test_multiple_imports(stats): record1 = { "requests": { "total": 10, + "rate_limited": 20, "attacksDetected": { "total": 5, "blocked": 3, @@ -168,6 +174,7 @@ def test_multiple_imports(stats): record2 = { "requests": { "total": 20, + "rate_limited": 5, "attacksDetected": { "total": 10, "blocked": 7, @@ -184,6 +191,7 @@ def test_multiple_imports(stats): stats.import_from_record(record1) stats.import_from_record(record2) assert stats.total_hits == 30 + assert stats.rate_limited_hits == 25 assert stats.attacks_detected == 15 assert stats.attacks_blocked == 10 assert stats.operations == { @@ -204,6 +212,7 @@ def test_import_empty_record(stats): record = {"requests": {}} stats.import_from_record(record) assert stats.total_hits == 0 + assert stats.rate_limited_hits == 0 assert stats.attacks_detected == 0 assert stats.attacks_blocked == 0 assert stats.operations == {} @@ -213,6 +222,7 @@ def test_import_partial_record(stats): record = {"requests": {"total": 10}} stats.import_from_record(record) assert stats.total_hits == 10 + assert stats.rate_limited_hits == 0 assert stats.attacks_detected == 0 assert stats.attacks_blocked == 0 assert stats.operations == {} @@ -242,3 +252,40 @@ def test_multiple_increments_and_detects(stats): "kind": "sql_op", "total": 1, } + + stats.on_rate_limit() + assert stats.rate_limited_hits == 1 + + stats.on_rate_limit() + assert stats.rate_limited_hits == 2 + + +def test_multiple_rate_limits(stats): + """Test multiple rate limit calls""" + for _ in range(5): + stats.on_rate_limit() + assert stats.rate_limited_hits == 5 + + +def test_rate_limit_in_get_record(): + """Test that rate_limited_hits is included in get_record output""" + stats = Statistics() + stats.total_hits = 10 + stats.on_rate_limit() + stats.on_rate_limit() + stats.on_rate_limit() + + record = stats.get_record() + assert record["requests"]["rate_limited"] == 3 + assert record["requests"]["total"] == 10 + + +def test_rate_limit_clear(): + """Test that clear() resets rate_limited_hits""" + stats = Statistics() + stats.on_rate_limit() + stats.on_rate_limit() + assert stats.rate_limited_hits == 2 + + stats.clear() + assert stats.rate_limited_hits == 0 From ed82846f4d77908cd113b6d600c8d602e5ad50aa Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 13 Jun 2025 11:21:42 +0200 Subject: [PATCH 2/7] should_block_request count rate_limited_hits --- aikido_zen/middleware/init_test.py | 4 ++++ aikido_zen/middleware/should_block_request.py | 1 + 2 files changed, 5 insertions(+) diff --git a/aikido_zen/middleware/init_test.py b/aikido_zen/middleware/init_test.py index 8ee6f2a2..22716deb 100644 --- a/aikido_zen/middleware/init_test.py +++ b/aikido_zen/middleware/init_test.py @@ -61,6 +61,7 @@ def test_with_context_with_cache(): } assert get_current_context().executed_middleware == True assert thread_cache.middleware_installed == True + assert thread_cache.stats.rate_limited_hits == 0 thread_cache.config.blocked_uids = [] assert should_block_request() == {"block": False} @@ -69,6 +70,7 @@ def test_with_context_with_cache(): assert should_block_request() == {"block": False} assert get_current_context().executed_middleware == True assert thread_cache.middleware_installed == True + assert thread_cache.stats.rate_limited_hits == 0 def test_cache_comms_with_endpoints(): @@ -158,9 +160,11 @@ def test_cache_comms_with_endpoints(): "success": True, "data": {"block": True, "trigger": "my_trigger"}, } + assert thread_cache.stats.rate_limited_hits == 0 assert should_block_request() == { "block": True, "ip": "::1", "type": "ratelimited", "trigger": "my_trigger", } + assert thread_cache.stats.rate_limited_hits == 1 diff --git a/aikido_zen/middleware/should_block_request.py b/aikido_zen/middleware/should_block_request.py index 20e4d143..360d5fde 100644 --- a/aikido_zen/middleware/should_block_request.py +++ b/aikido_zen/middleware/should_block_request.py @@ -51,6 +51,7 @@ def should_block_request(): receive=True, ) if ratelimit_res["success"] and ratelimit_res["data"]["block"]: + cache.stats.on_rate_limit() return { "block": True, "type": "ratelimited", From dd77188b83402794a8533d50ae25bc5120841b1f Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 13 Jun 2025 11:25:22 +0200 Subject: [PATCH 3/7] rate_limited -> rateLimited --- aikido_zen/storage/statistics/__init__.py | 4 ++-- aikido_zen/storage/statistics/init_test.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/aikido_zen/storage/statistics/__init__.py b/aikido_zen/storage/statistics/__init__.py index 41daf0c3..567018e5 100644 --- a/aikido_zen/storage/statistics/__init__.py +++ b/aikido_zen/storage/statistics/__init__.py @@ -43,7 +43,7 @@ def get_record(self): "endedAt": current_time, "requests": { "total": self.total_hits, - "rate_limited": self.rate_limited_hits, + "rateLimited": self.rate_limited_hits, "aborted": 0, # statistic currently not in use "attacksDetected": { "total": self.attacks_detected, @@ -56,7 +56,7 @@ def get_record(self): def import_from_record(self, record): attacks_detected = record.get("requests", {}).get("attacksDetected", {}) self.total_hits += record.get("requests", {}).get("total", 0) - self.rate_limited_hits += record.get("requests", {}).get("rate_limited", 0) + self.rate_limited_hits += record.get("requests", {}).get("rateLimited", 0) self.attacks_detected += attacks_detected.get("total", 0) self.attacks_blocked += attacks_detected.get("blocked", 0) self.operations.update(record.get("operations", {})) diff --git a/aikido_zen/storage/statistics/init_test.py b/aikido_zen/storage/statistics/init_test.py index 25737678..7d54ec1c 100644 --- a/aikido_zen/storage/statistics/init_test.py +++ b/aikido_zen/storage/statistics/init_test.py @@ -79,7 +79,7 @@ def test_get_record(monkeypatch): assert record["startedAt"] == stats.started_at assert record["endedAt"] == mock_time assert record["requests"]["total"] == 10 - assert record["requests"]["rate_limited"] == 2 + assert record["requests"]["rateLimited"] == 2 assert record["requests"]["aborted"] == 0 assert record["requests"]["attacksDetected"]["total"] == 5 assert record["requests"]["attacksDetected"]["blocked"] == 3 @@ -100,7 +100,7 @@ def test_import_from_record(): record = { "requests": { "total": 10, - "rate_limited": 5, + "rateLimited": 5, "attacksDetected": { "total": 5, "blocked": 3, @@ -157,7 +157,7 @@ def test_multiple_imports(stats): record1 = { "requests": { "total": 10, - "rate_limited": 20, + "rateLimited": 20, "attacksDetected": { "total": 5, "blocked": 3, @@ -174,7 +174,7 @@ def test_multiple_imports(stats): record2 = { "requests": { "total": 20, - "rate_limited": 5, + "rateLimited": 5, "attacksDetected": { "total": 10, "blocked": 7, @@ -276,7 +276,7 @@ def test_rate_limit_in_get_record(): stats.on_rate_limit() record = stats.get_record() - assert record["requests"]["rate_limited"] == 3 + assert record["requests"]["rateLimited"] == 3 assert record["requests"]["total"] == 10 From 4e26ceb38f5b6e9fa9758d1b046c09e2cc6a7f5f Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 13 Jun 2025 11:27:20 +0200 Subject: [PATCH 4/7] Update SYNC_DATA test cases to include rateLimited field --- aikido_zen/background_process/commands/sync_data_test.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aikido_zen/background_process/commands/sync_data_test.py b/aikido_zen/background_process/commands/sync_data_test.py index b12451a0..4827947f 100644 --- a/aikido_zen/background_process/commands/sync_data_test.py +++ b/aikido_zen/background_process/commands/sync_data_test.py @@ -54,6 +54,7 @@ def test_process_sync_data_initialization(setup_connection_manager): "endedAt": 1, "requests": { "total": 10, + "rateLimited": 0, "aborted": 0, "attacksDetected": { "total": 5, @@ -94,6 +95,7 @@ def test_process_sync_data_initialization(setup_connection_manager): "aborted": 0, "attacksDetected": {"blocked": 0, "total": 5}, "total": 10, + "rateLimited": 0, } # Check that the return value is correct @@ -135,6 +137,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana "endedAt": 1, "requests": { "total": 10, + "rateLimited": 0, "aborted": 0, "attacksDetected": { "total": 5, @@ -167,6 +170,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana "aborted": 0, "attacksDetected": {"blocked": 0, "total": 5}, "total": 10, + "rateLimited": 0, } assert connection_manager.middleware_installed == True assert len(connection_manager.hostnames.as_array()) == 0 @@ -199,6 +203,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager "endedAt": 1, "requests": { "total": 5, + "rateLimited": 0, "aborted": 0, "attacksDetected": { "total": 5, @@ -227,6 +232,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager "endedAt": 1, "requests": { "total": 15, + "rateLimited": 0, "aborted": 0, "attacksDetected": { "total": 5, @@ -251,6 +257,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager "aborted": 0, "attacksDetected": {"blocked": 0, "total": 10}, "total": 20, + "rateLimited": 0, } assert connection_manager.middleware_installed == False assert connection_manager.hostnames.as_array() == [ From 3ef9c1bd8df0818851e3d74d29607e14d6f5b8d7 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 13 Jun 2025 11:33:11 +0200 Subject: [PATCH 5/7] Update thread_cache_test cases to include RateLimited field --- aikido_zen/thread/thread_cache_test.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index c64e6fad..861124b6 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -42,6 +42,7 @@ def test_initialization(thread_cache: ThreadCache): assert thread_cache.config.blocked_uids == set() assert thread_cache.stats.get_record()["requests"] == { "total": 0, + "rateLimited": 0, "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, } @@ -77,6 +78,7 @@ def test_reset(thread_cache: ThreadCache): assert thread_cache.config.blocked_uids == set() assert thread_cache.stats.get_record()["requests"] == { "total": 0, + "rateLimited": 0, "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, } @@ -100,6 +102,7 @@ def test_renew_with_no_comms(thread_cache: ThreadCache): assert thread_cache.config.blocked_uids == set() assert thread_cache.stats.get_record()["requests"] == { "total": 0, + "rateLimited": 0, "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, } @@ -285,6 +288,7 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach "endedAt": -1, "requests": { "total": 2, + "rateLimited": 0, "aborted": 0, "attacksDetected": {"blocked": 1, "total": 3}, }, @@ -369,6 +373,7 @@ def test_sync_data_for_users(mock_get_comms, thread_cache: ThreadCache): "endedAt": -1, "requests": { "total": 1, + "rateLimited": 0, "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, }, @@ -421,6 +426,7 @@ def test_renew_called_with_empty_routes(mock_get_comms, thread_cache: ThreadCach "endedAt": -1, "requests": { "total": 0, + "rateLimited": 0, "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, }, @@ -461,6 +467,7 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache "endedAt": -1, "requests": { "total": 0, + "rateLimited": 0, "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, }, From 6d467e3b228f38dcdb0fe8e82fa12f0853e21906 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 13 Jun 2025 11:43:56 +0200 Subject: [PATCH 6/7] Update e2e test to add rateLimited field --- end2end/django_mysql_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/end2end/django_mysql_test.py b/end2end/django_mysql_test.py index db606019..b293d110 100644 --- a/end2end/django_mysql_test.py +++ b/end2end/django_mysql_test.py @@ -88,6 +88,6 @@ def test_initial_heartbeat(): "method": "POST", "path": "/app/create" }], - {"aborted":0,"attacksDetected":{"blocked":2,"total":2},"total":3}, + {"aborted":0,"attacksDetected":{"blocked":2,"total":2},"total":3, 'rateLimited': 0}, {'asgiref', 'regex', 'mysqlclient', 'sqlparse', 'aikido_zen', 'django'} ) From ad14623414f9ad615b62bcbb1036b916991104cf Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 13 Jun 2025 11:46:29 +0200 Subject: [PATCH 7/7] Update statistics comment --- aikido_zen/storage/statistics/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aikido_zen/storage/statistics/__init__.py b/aikido_zen/storage/statistics/__init__.py index 567018e5..47e7dec3 100644 --- a/aikido_zen/storage/statistics/__init__.py +++ b/aikido_zen/storage/statistics/__init__.py @@ -4,8 +4,8 @@ class Statistics: """ - Keeps track of total and aborted requests - and total and blocked attacks + Stores: hits, counts of attacks (split up in detected/blocked), count of rate-limited requests, + statistics for operations (i.e. how many times did we see a query being executed) """ def __init__(self):