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 f9b15cfc9..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,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 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/background_process/commands/put_event.py b/aikido_zen/background_process/commands/put_event.py index a2e51a94b..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 - # "time" and "agent" fields are added by default from the CloudConnectionManager + # The CloudConnectionManager adds the "agent" field 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..6efb56b3f 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 @@ -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)) 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..d7ce8dcf1 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 @@ -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: 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..49949ba69 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 @@ -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)) 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, }