Improved command data structure & handler

aikido_zen/background_process/aikido_background_process.py

@@ -102,13 +102,7 @@ def send_to_connection_manager(self, event_scheduler):
             EMPTY_QUEUE_INTERVAL, 1, self.send_to_connection_manager, (event_scheduler,)
         )
         while not self.queue.empty():
-            queue_attack_item = self.queue.get()
-            self.connection_manager.on_detected_attack(
-                attack=queue_attack_item[0],
-                context=queue_attack_item[1],
-                blocked=queue_attack_item[2],
-                stack=queue_attack_item[3],
-            )
+            self.connection_manager.report_api_event(self.queue.get())
 
 
 def add_exit_handlers():

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -10,12 +10,12 @@
 from aikido_zen.storage.users import Users
 from aikido_zen.storage.hostnames import Hostnames
 from ..realtime.start_polling_for_changes import start_polling_for_changes
+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
 
 # Import functions :
-from .on_detected_attack import on_detected_attack
 from .get_manager_info import get_manager_info
 from .update_service_config import update_service_config
 from .on_start import on_start
@@ -57,7 +57,7 @@ def __init__(self, block, api, token, serverless):
 
     def start(self, event_scheduler):
         """Send out start event and add heartbeats"""
-        res = self.on_start()
+        res = on_start(self)
         if res.get("error", None) == "invalid_token":
             logger.info(
                 "Token was invalid, not starting heartbeats and realtime polling."
@@ -81,26 +81,37 @@ def report_initial_stats(self):
         if should_report_initial_stats:
             self.send_heartbeat()
 
-    def on_detected_attack(self, attack, context, blocked, stack):
-        """This will send something to the API when an attack is detected"""
-        return on_detected_attack(self, attack, context, blocked, stack)
-
-    def on_start(self):
-        """This will send out an Event signalling the start to the server"""
-        return on_start(self)
-
     def send_heartbeat(self):
         """This will send a heartbeat to the server"""
         return send_heartbeat(self)
 
-    def get_manager_info(self):
-        """This returns info about the connection_manager"""
-        return get_manager_info(self)
-
     def update_service_config(self, res):
         """Update configuration based on the server's response"""
         return update_service_config(self, res)
 
     def update_firewall_lists(self):
         """Will update service config with blocklist of IP addresses"""
         return update_firewall_lists(self)
+
+    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.update(event)  # Merge default fields with event fields
+
+            result = self.api.report(self.token, payload, self.timeout_in_sec)
+            if not result.get("success", True):
+                logger.error(
+                    "CloudConnectionManager: Reporting to api failed, error=%s",
+                    result.get("error", "unknown"),
+                )
+            return result
+        except Exception as e:
+            logger.debug(e)
+            logger.error(
+                "CloudConnectionManager: Reporting to api failed, unexpected error (see debug logs)"
+            )

aikido_zen/background_process/cloud_connection_manager/on_start.py

@@ -5,26 +5,16 @@
 
 
 def on_start(connection_manager):
-    """
-    This will send out an Event signalling the start to the server
-    """
-    if not connection_manager.token:
-        return
-    res = connection_manager.api.report(
-        connection_manager.token,
-        {
-            "type": "started",
-            "time": get_unixtime_ms(),
-            "agent": connection_manager.get_manager_info(),
-        },
-        connection_manager.timeout_in_sec,
-    )
+    event = {"type": "started"}
+    res = connection_manager.report_api_event(event)
+
     if not res.get("success", True):
-        # Update config time even in failure :
-        connection_manager.conf.last_updated_at = get_unixtime_ms()
-        logger.error("Failed to communicate with Aikido Server : %s", res["error"])
+        connection_manager.conf.last_updated_at = (
+            get_unixtime_ms()
+        )  # Update config time even in failure
     else:
         connection_manager.update_service_config(res)
         connection_manager.update_firewall_lists()
         logger.info("Established connection with Aikido Server")
+
     return res

aikido_zen/background_process/cloud_connection_manager/on_start_test.py

@@ -8,7 +8,7 @@ def mock_connection_manager():
     connection_manager = MagicMock()
     connection_manager.token = "test_token"
     connection_manager.timeout_in_sec = 5
-    connection_manager.api.report = MagicMock(return_value={"success": True})
+    connection_manager.report_api_event = MagicMock(return_value={"success": True})
     connection_manager.get_manager_info = lambda: {}
     connection_manager.update_service_config = MagicMock()
     return connection_manager
@@ -19,15 +19,15 @@ def test_on_start_no_token():
     connection_manager = MagicMock()
     connection_manager.token = None
     on_start(connection_manager)
-    connection_manager.api.report.assert_not_called()
+    connection_manager.report_api_event.assert_called()
 
 
 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)
 
     # Check that the API report method was called
-    mock_connection_manager.api.report.assert_called_once()
+    mock_connection_manager.report_api_event.assert_called_once()
 
     # Check that the service config was updated
     mock_connection_manager.update_service_config.assert_called_once()
@@ -38,18 +38,15 @@ def test_on_start_success(mock_connection_manager, caplog):
 
 def test_on_start_failure(mock_connection_manager, caplog):
     """Test that an error is logged when the API call fails."""
-    mock_connection_manager.api.report.return_value = {
+    mock_connection_manager.report_api_event.return_value = {
         "success": False,
         "error": "Some error",
     }
 
     on_start(mock_connection_manager)
 
     # Check that the API report method was called
-    mock_connection_manager.api.report.assert_called_once()
+    mock_connection_manager.report_api_event.assert_called_once()
 
     # Check that the service config was not updated
     mock_connection_manager.update_service_config.assert_not_called()
-
-    # Check that the error log was called
-    assert "Failed to communicate with Aikido Server : Some error" in caplog.text

aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py

@@ -2,7 +2,6 @@
 
 from aikido_zen.background_process.packages import PackagesStore
 from aikido_zen.helpers.logging import logger
-from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms
 
 
 def send_heartbeat(connection_manager):
@@ -26,20 +25,16 @@ def send_heartbeat(connection_manager):
     connection_manager.ai_stats.clear()
     PackagesStore.clear()
 
-    res = connection_manager.api.report(
-        connection_manager.token,
+    res = connection_manager.report_api_event(
         {
             "type": "heartbeat",
-            "time": get_unixtime_ms(),
-            "agent": connection_manager.get_manager_info(),
             "stats": stats,
             "ai": ai_stats,
             "hostnames": outgoing_domains,
             "packages": packages,
             "routes": routes,
             "users": users,
             "middlewareInstalled": connection_manager.middleware_installed,
-        },
-        connection_manager.timeout_in_sec,
+        }
     )
     connection_manager.update_service_config(res)

aikido_zen/background_process/cloud_connection_manager/update_service_config.py

@@ -9,8 +9,8 @@ def update_service_config(connection_manager, res):
     Update configuration based on the server's response
     """
     if res.get("success", False) is False:
-        logger.debug(res)
         return
+
     if "block" in res.keys() and res["block"] != connection_manager.block:
         logger.debug("Updating blocking, setting blocking to : %s", res["block"])
         connection_manager.block = bool(res["block"])

aikido_zen/background_process/commands/__init__.py

@@ -1,34 +1,34 @@
-"""Commands __init__.py file"""
-
 from aikido_zen.helpers.logging import logger
-from .attack import process_attack
+from aikido_zen.helpers.ipc.command_types import CommandContext
+from .put_event import PutEventCommand
 from .check_firewall_lists import process_check_firewall_lists
 from .read_property import process_read_property
 from .should_ratelimit import process_should_ratelimit
 from .ping import process_ping
 from .sync_data import process_sync_data
 
 commands_map = {
-    # This maps to a tuple : (function, returns_data?)
-    # Commands that don't return data :
-    "ATTACK": (process_attack, False),
-    # Commands that return data :
-    "SYNC_DATA": (process_sync_data, True),
-    "READ_PROPERTY": (process_read_property, True),
-    "SHOULD_RATELIMIT": (process_should_ratelimit, True),
-    "PING": (process_ping, True),
-    "CHECK_FIREWALL_LISTS": (process_check_firewall_lists, True),
+    "SYNC_DATA": process_sync_data,
+    "READ_PROPERTY": process_read_property,
+    "SHOULD_RATELIMIT": process_should_ratelimit,
+    "PING": process_ping,
+    "CHECK_FIREWALL_LISTS": process_check_firewall_lists,
 }
 
+modern_commands = [PutEventCommand]
+
 
 def process_incoming_command(connection_manager, obj, conn, queue):
-    """Processes an incoming command"""
-    action = obj[0]
-    data = obj[1]
-    if action in commands_map:
-        func, returns_data = commands_map[action]
-        if returns_data:
-            return conn.send(func(connection_manager, data, queue))
-        func(connection_manager, data, queue)
-    else:
-        logger.debug("Command : `%s` not found, aborting", action)
+    inbound_identifier = obj[0]
+    inbound_request = obj[1]
+    if inbound_identifier in commands_map:
+        func = commands_map[inbound_identifier]
+        return conn.send(func(connection_manager, inbound_request))
+
+    for cmd in modern_commands:
+        if cmd.identifier() == inbound_identifier:
+            cmd.run(CommandContext(connection_manager, queue, conn), inbound_request)
+            return None
+
+    logger.debug("Command : `%s` not found - did not execute", inbound_identifier)
+    return None