Stamp queued events with the time they were built

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)

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,6 +20,27 @@ def setup_cloud_connection_manager():
     return CloudConnectionManager(block, api, token, serverless)
 
 
+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})
+
+    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_send_heartbeat_sets_the_time(setup_cloud_connection_manager):
+    manager = setup_cloud_connection_manager
+    manager.report_api_event = MagicMock(return_value={"success": True})
+
+    with test_utils.patch_time(time_ms=111):
+        manager.send_heartbeat()
+
+    assert manager.report_api_event.call_args[0][0]["time"] == 111
+
+
 def test_cloud_connection_manager_initialization(setup_cloud_connection_manager):
     manager = setup_cloud_connection_manager
 

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):

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)

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)

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 the "agent" field
         self.event = event
 
 

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",

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
@@ -22,6 +23,7 @@ def create_attack_wave_event(context):
                 "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))

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()

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
 
 
@@ -7,6 +8,7 @@ def create_custom_event(event_name, context):
             "type": "custom",
             "name": event_name,
             "request": extract_request_if_possible(context),
+            "time": get_unixtime_ms(),
         }
         user = getattr(context, "user", None)
         if user:

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"})

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
 
@@ -17,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))

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,
     }