From 66d901c52afd405e0e629d7cd3f17600b4d50591 Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Wed, 7 Oct 2026 12:43:22 +0300 Subject: [PATCH 1/3] Stamp queued events with the time they were built --- .../cloud_connection_manager/init_test.py | 30 +++++++++++++++++++ .../background_process/commands/put_event.py | 2 +- aikido_zen/context/track_test.py | 8 ++++- .../helpers/create_attack_wave_event.py | 2 ++ .../helpers/create_attack_wave_event_test.py | 17 +++++++++++ aikido_zen/helpers/create_custom_event.py | 2 ++ .../helpers/create_custom_event_test.py | 14 ++++++++- .../create_detected_attack_api_event.py | 2 ++ .../create_detected_attack_api_event_test.py | 9 ++++-- 9 files changed, 81 insertions(+), 5 deletions(-) diff --git a/aikido_zen/background_process/cloud_connection_manager/init_test.py b/aikido_zen/background_process/cloud_connection_manager/init_test.py index f9b15cfc9..b9e1558dc 100644 --- a/aikido_zen/background_process/cloud_connection_manager/init_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/init_test.py @@ -19,6 +19,36 @@ def setup_cloud_connection_manager(): return CloudConnectionManager(block, api, token, serverless) +def test_report_api_event_keeps_the_event_time(setup_cloud_connection_manager): + manager = setup_cloud_connection_manager + manager.api.report = MagicMock(return_value={"success": True}) + + with patch( + "aikido_zen.background_process.cloud_connection_manager.get_unixtime_ms", + return_value=9999, + ): + manager.report_api_event({"type": "detected_attack", "time": 111}) + + payload = manager.api.report.call_args[0][1] + assert payload["time"] == 111 + + +def test_report_api_event_stamps_a_send_time_when_the_event_has_none( + setup_cloud_connection_manager, +): + manager = setup_cloud_connection_manager + manager.api.report = MagicMock(return_value={"success": True}) + + with patch( + "aikido_zen.background_process.cloud_connection_manager.get_unixtime_ms", + return_value=9999, + ): + manager.report_api_event({"type": "heartbeat"}) + + payload = manager.api.report.call_args[0][1] + assert payload["time"] == 9999 + + def test_cloud_connection_manager_initialization(setup_cloud_connection_manager): manager = setup_cloud_connection_manager diff --git a/aikido_zen/background_process/commands/put_event.py b/aikido_zen/background_process/commands/put_event.py index a2e51a94b..638f3a644 100644 --- a/aikido_zen/background_process/commands/put_event.py +++ b/aikido_zen/background_process/commands/put_event.py @@ -4,7 +4,7 @@ class PutEventReq: def __init__(self, event): # Event is a dictionary containing data that is going to be reported to core - # "time" and "agent" fields are added by default from the CloudConnectionManager + # The CloudConnectionManager adds "agent", and "time" if the event has none self.event = event diff --git a/aikido_zen/context/track_test.py b/aikido_zen/context/track_test.py index 790a2e760..7134f6dba 100644 --- a/aikido_zen/context/track_test.py +++ b/aikido_zen/context/track_test.py @@ -95,7 +95,12 @@ def test_track_sends_event_over_ipc(): with patch( "aikido_zen.thread.thread_cache.get_cache", return_value=cache_with_bypassed_ips(), - ), patch("aikido_zen.background_process.comms.get_comms", return_value=comms): + ), patch( + "aikido_zen.background_process.comms.get_comms", return_value=comms + ), patch( + "aikido_zen.helpers.create_custom_event.get_unixtime_ms", + return_value=111, + ): track("my-custom-event") comms.send_data_to_bg_process.assert_called_once() @@ -104,6 +109,7 @@ def test_track_sends_event_over_ipc(): assert returns_data is False assert request.event == { "type": "custom", + "time": 111, "name": "my-custom-event", "request": { "method": "POST", diff --git a/aikido_zen/helpers/create_attack_wave_event.py b/aikido_zen/helpers/create_attack_wave_event.py index 8650d0b30..43b7bf5b3 100644 --- a/aikido_zen/helpers/create_attack_wave_event.py +++ b/aikido_zen/helpers/create_attack_wave_event.py @@ -1,4 +1,5 @@ import json +from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms from aikido_zen.helpers.limit_length_metadata import limit_length_metadata from aikido_zen.helpers.logging import logger from aikido_zen.storage.attack_wave_detector_store import attack_wave_detector_store @@ -17,6 +18,7 @@ def create_attack_wave_event(context): return { "type": "detected_attack_wave", + "time": get_unixtime_ms(), "attack": { "user": getattr(context, "user", None), "metadata": limit_length_metadata(metadata, 4096), diff --git a/aikido_zen/helpers/create_attack_wave_event_test.py b/aikido_zen/helpers/create_attack_wave_event_test.py index 05e9dc25a..b2cc2b438 100644 --- a/aikido_zen/helpers/create_attack_wave_event_test.py +++ b/aikido_zen/helpers/create_attack_wave_event_test.py @@ -28,6 +28,23 @@ def test_create_attack_wave_event_success(): assert event["request"] is not None +@patch( + "aikido_zen.helpers.create_attack_wave_event.get_unixtime_ms", + return_value=111, +) +def test_create_attack_wave_event_sets_the_time(mock_time): + context = test_utils.generate_context() + + with patch.object( + attack_wave_detector_store, "get_samples_for_ip", return_value=None + ), patch.object( + attack_wave_detector_store, "clear_samples_for_ip", return_value=None + ): + event = create_attack_wave_event(context) + + assert event["time"] == 111 + + def test_create_attack_wave_event_with_samples(): """Test attack wave event creation with samples from store""" context = test_utils.generate_context() diff --git a/aikido_zen/helpers/create_custom_event.py b/aikido_zen/helpers/create_custom_event.py index ccb4349cc..80212328e 100644 --- a/aikido_zen/helpers/create_custom_event.py +++ b/aikido_zen/helpers/create_custom_event.py @@ -1,3 +1,4 @@ +from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms from aikido_zen.helpers.logging import logger @@ -5,6 +6,7 @@ def create_custom_event(event_name, context): try: event = { "type": "custom", + "time": get_unixtime_ms(), "name": event_name, "request": extract_request_if_possible(context), } diff --git a/aikido_zen/helpers/create_custom_event_test.py b/aikido_zen/helpers/create_custom_event_test.py index 41239ee69..a9c644469 100644 --- a/aikido_zen/helpers/create_custom_event_test.py +++ b/aikido_zen/helpers/create_custom_event_test.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from .create_custom_event import ( create_custom_event, extract_request_if_possible, @@ -18,6 +18,18 @@ def test_create_custom_event_success(): assert event["request"] is not None +@patch( + "aikido_zen.helpers.create_custom_event.get_unixtime_ms", + return_value=111, +) +def test_create_custom_event_sets_the_time(mock_time): + context = test_utils.generate_context() + + event = create_custom_event("user.login_failed", context) + + assert event["time"] == 111 + + def test_create_custom_event_with_user(): """Test custom event creation with user information""" context = test_utils.generate_context(user={"id": "user-1", "name": "Jane Doe"}) diff --git a/aikido_zen/helpers/create_detected_attack_api_event.py b/aikido_zen/helpers/create_detected_attack_api_event.py index c236092fc..3d8b66e3a 100644 --- a/aikido_zen/helpers/create_detected_attack_api_event.py +++ b/aikido_zen/helpers/create_detected_attack_api_event.py @@ -1,5 +1,6 @@ import json +from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms from aikido_zen.helpers.limit_length_metadata import limit_length_metadata from aikido_zen.helpers.logging import logger @@ -8,6 +9,7 @@ def create_detected_attack_api_event(attack, context, blocked, stack): try: return { "type": "detected_attack", + "time": get_unixtime_ms(), "attack": { **attack, "user": getattr(context, "user", None), diff --git a/aikido_zen/helpers/create_detected_attack_api_event_test.py b/aikido_zen/helpers/create_detected_attack_api_event_test.py index 9918dabdc..bdac3d6a3 100644 --- a/aikido_zen/helpers/create_detected_attack_api_event_test.py +++ b/aikido_zen/helpers/create_detected_attack_api_event_test.py @@ -1,5 +1,5 @@ import pytest -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from .create_detected_attack_api_event import create_detected_attack_api_event from aikido_zen.context import Context import aikido_zen.test_utils as test_utils @@ -39,7 +39,11 @@ def test_create_attack_event_with_long_metadata(): assert event["attack"]["metadata"]["test"] == long_metadata[:4096] -def test_create_attack_event_success(): +@patch( + "aikido_zen.helpers.create_detected_attack_api_event.get_unixtime_ms", + return_value=111, +) +def test_create_attack_event_success(mock_time): attack = { "payload": {"key": "value"}, "metadata": {}, @@ -68,6 +72,7 @@ def test_create_attack_event_success(): "userAgent": None, }, "type": "detected_attack", + "time": 111, } From 644062d5a17d8bfe846ebf4de5506586a087a2d8 Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Wed, 7 Oct 2026 15:13:11 +0300 Subject: [PATCH 2/3] Stamp the heartbeat and startup events in the same way --- .../cloud_connection_manager/__init__.py | 6 +---- .../cloud_connection_manager/init_test.py | 26 +++++++------------ .../cloud_connection_manager/on_start.py | 2 +- .../cloud_connection_manager/on_start_test.py | 9 +++++++ .../send_heartbeat.py | 2 ++ .../helpers/create_attack_wave_event.py | 2 +- aikido_zen/helpers/create_custom_event.py | 2 +- .../create_detected_attack_api_event.py | 2 +- 8 files changed, 25 insertions(+), 26 deletions(-) diff --git a/aikido_zen/background_process/cloud_connection_manager/__init__.py b/aikido_zen/background_process/cloud_connection_manager/__init__.py index 593d7e39b..dd41119a8 100644 --- a/aikido_zen/background_process/cloud_connection_manager/__init__.py +++ b/aikido_zen/background_process/cloud_connection_manager/__init__.py @@ -11,7 +11,6 @@ from aikido_zen.storage.hostnames import Hostnames from ..realtime.start_polling_for_changes import start_polling_for_changes from ..realtime.listen_for_config_updates import listen_for_config_updates -from ...helpers.get_current_unixtime_ms import get_unixtime_ms from ...storage.ai_statistics import AIStatistics from ...storage.firewall_lists import FirewallLists from ...storage.statistics import Statistics @@ -88,10 +87,7 @@ def report_api_event(self, event): if not self.token: return {"success": False, "error": "invalid_token"} try: - payload = { - "time": get_unixtime_ms(), - "agent": get_manager_info(self), - } + payload = {"agent": get_manager_info(self)} payload.update(event) # Merge default fields with event fields result = self.api.report(self.token, payload, self.timeout_in_sec) diff --git a/aikido_zen/background_process/cloud_connection_manager/init_test.py b/aikido_zen/background_process/cloud_connection_manager/init_test.py index b9e1558dc..16b2f18ba 100644 --- a/aikido_zen/background_process/cloud_connection_manager/init_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/init_test.py @@ -8,6 +8,7 @@ from aikido_zen.ratelimiting.rate_limiter import RateLimiter from aikido_zen.storage.statistics import Statistics from . import CloudConnectionManager +import aikido_zen.test_utils as test_utils @pytest.fixture @@ -19,34 +20,25 @@ def setup_cloud_connection_manager(): return CloudConnectionManager(block, api, token, serverless) -def test_report_api_event_keeps_the_event_time(setup_cloud_connection_manager): +def test_report_api_event_only_adds_the_agent(setup_cloud_connection_manager): manager = setup_cloud_connection_manager manager.api.report = MagicMock(return_value={"success": True}) - with patch( - "aikido_zen.background_process.cloud_connection_manager.get_unixtime_ms", - return_value=9999, - ): - manager.report_api_event({"type": "detected_attack", "time": 111}) + manager.report_api_event({"type": "detected_attack", "time": 111}) payload = manager.api.report.call_args[0][1] assert payload["time"] == 111 + assert "agent" in payload -def test_report_api_event_stamps_a_send_time_when_the_event_has_none( - setup_cloud_connection_manager, -): +def test_send_heartbeat_sets_the_time(setup_cloud_connection_manager): manager = setup_cloud_connection_manager - manager.api.report = MagicMock(return_value={"success": True}) + manager.report_api_event = MagicMock(return_value={"success": True}) - with patch( - "aikido_zen.background_process.cloud_connection_manager.get_unixtime_ms", - return_value=9999, - ): - manager.report_api_event({"type": "heartbeat"}) + with test_utils.patch_time(time_ms=111): + manager.send_heartbeat() - payload = manager.api.report.call_args[0][1] - assert payload["time"] == 9999 + assert manager.report_api_event.call_args[0][0]["time"] == 111 def test_cloud_connection_manager_initialization(setup_cloud_connection_manager): diff --git a/aikido_zen/background_process/cloud_connection_manager/on_start.py b/aikido_zen/background_process/cloud_connection_manager/on_start.py index 6407ee7cf..08cd40b20 100644 --- a/aikido_zen/background_process/cloud_connection_manager/on_start.py +++ b/aikido_zen/background_process/cloud_connection_manager/on_start.py @@ -5,7 +5,7 @@ def on_start(connection_manager): - event = {"type": "started"} + event = {"type": "started", "time": get_unixtime_ms()} res = connection_manager.report_api_event(event) if not res.get("success", True): diff --git a/aikido_zen/background_process/cloud_connection_manager/on_start_test.py b/aikido_zen/background_process/cloud_connection_manager/on_start_test.py index fd9580d4e..2f14bf4f2 100644 --- a/aikido_zen/background_process/cloud_connection_manager/on_start_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/on_start_test.py @@ -1,6 +1,7 @@ import pytest from unittest.mock import MagicMock, patch from .on_start import on_start +import aikido_zen.test_utils as test_utils @pytest.fixture @@ -22,6 +23,14 @@ def test_on_start_no_token(): connection_manager.report_api_event.assert_called() +def test_on_start_sets_the_time(mock_connection_manager): + """Test that the event carries the time it was built""" + with test_utils.patch_time(time_ms=111): + on_start(mock_connection_manager) + + assert mock_connection_manager.report_api_event.call_args[0][0]["time"] == 111 + + def test_on_start_success(mock_connection_manager, caplog): """Test that the API call is made successfully and the service config is updated.""" on_start(mock_connection_manager) 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 18fa732fa..1a29f1123 100644 --- a/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py +++ b/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py @@ -1,6 +1,7 @@ """Exports the send_heartbeat function""" from aikido_zen.background_process.packages import PackagesStore +from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms from aikido_zen.helpers.logging import logger @@ -35,6 +36,7 @@ def send_heartbeat(connection_manager): "routes": routes, "users": users, "middlewareInstalled": connection_manager.middleware_installed, + "time": get_unixtime_ms(), } ) connection_manager.update_service_config(res) diff --git a/aikido_zen/helpers/create_attack_wave_event.py b/aikido_zen/helpers/create_attack_wave_event.py index 43b7bf5b3..6efb56b3f 100644 --- a/aikido_zen/helpers/create_attack_wave_event.py +++ b/aikido_zen/helpers/create_attack_wave_event.py @@ -18,12 +18,12 @@ def create_attack_wave_event(context): return { "type": "detected_attack_wave", - "time": get_unixtime_ms(), "attack": { "user": getattr(context, "user", None), "metadata": limit_length_metadata(metadata, 4096), }, "request": extract_request_if_possible(context), + "time": get_unixtime_ms(), } except Exception as e: logger.error("Failed to create detected_attack_wave API event: %s", str(e)) diff --git a/aikido_zen/helpers/create_custom_event.py b/aikido_zen/helpers/create_custom_event.py index 80212328e..d7ce8dcf1 100644 --- a/aikido_zen/helpers/create_custom_event.py +++ b/aikido_zen/helpers/create_custom_event.py @@ -6,9 +6,9 @@ def create_custom_event(event_name, context): try: event = { "type": "custom", - "time": get_unixtime_ms(), "name": event_name, "request": extract_request_if_possible(context), + "time": get_unixtime_ms(), } user = getattr(context, "user", None) if user: diff --git a/aikido_zen/helpers/create_detected_attack_api_event.py b/aikido_zen/helpers/create_detected_attack_api_event.py index 3d8b66e3a..49949ba69 100644 --- a/aikido_zen/helpers/create_detected_attack_api_event.py +++ b/aikido_zen/helpers/create_detected_attack_api_event.py @@ -9,7 +9,6 @@ def create_detected_attack_api_event(attack, context, blocked, stack): try: return { "type": "detected_attack", - "time": get_unixtime_ms(), "attack": { **attack, "user": getattr(context, "user", None), @@ -19,6 +18,7 @@ def create_detected_attack_api_event(attack, context, blocked, stack): "stack": stack, }, "request": extract_request_if_possible(context), + "time": get_unixtime_ms(), } except Exception as e: logger.error("Failed to create detected_attack API event: %s", str(e)) From 78422a34e803c09c04e5f453cb13accb7fc76bc5 Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Wed, 7 Oct 2026 16:56:02 +0300 Subject: [PATCH 3/3] Correct the put_event comment --- aikido_zen/background_process/commands/put_event.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/background_process/commands/put_event.py b/aikido_zen/background_process/commands/put_event.py index 638f3a644..790df47a3 100644 --- a/aikido_zen/background_process/commands/put_event.py +++ b/aikido_zen/background_process/commands/put_event.py @@ -4,7 +4,7 @@ class PutEventReq: def __init__(self, event): # Event is a dictionary containing data that is going to be reported to core - # The CloudConnectionManager adds "agent", and "time" if the event has none + # The CloudConnectionManager adds the "agent" field self.event = event