Convert all commands to the modern equivalent

aikido_zen/background_process/__init__.py

@@ -18,6 +18,8 @@
     reset_comms,
 )
 from .aikido_background_process import AikidoBackgroundProcess
+from .commands import PingCommand
+from ..helpers.ipc.send_payload import send_payload
 
 
 def start_background_process():
@@ -64,8 +66,8 @@ def get_uds_filename():
 
 
 def background_process_already_active(comms):
-    res = comms.send_data_to_bg_process(action="PING", obj=tuple(), receive=True)
-    if res["success"] and res["data"] == "Received":
+    res = send_payload(comms, PingCommand.generate())
+    if res["success"] and res["data"] == "recv":
         # Ping is active, return.
         logger.debug(
             "A background agent is already running, not starting a new background agent."

aikido_zen/background_process/commands/__init__.py

@@ -1,21 +1,19 @@
 from aikido_zen.helpers.logging import logger
 from aikido_zen.helpers.ipc.command_types import CommandContext
+from .check_firewall_lists import CheckFirewallListsCommand
+from .ping import PingCommand
 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 = {
     "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]
+modern_commands = [PutEventCommand, PingCommand, CheckFirewallListsCommand]
 
 
 def process_incoming_command(connection_manager, obj, conn, queue):
@@ -25,10 +23,16 @@ def process_incoming_command(connection_manager, obj, conn, queue):
         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
+    command = None
+    for command_option in modern_commands:
+        if command_option.identifier() == inbound_identifier:
+            command = command_option
+    if command is None:
+        logger.debug("Command : `%s` not found - did not execute", inbound_identifier)
+        return None
+
+    res = command.run(CommandContext(connection_manager, queue, conn), inbound_request)
+    if command.returns_data():
+        conn.send(res)
 
-    logger.debug("Command : `%s` not found - did not execute", inbound_identifier)
     return None

aikido_zen/background_process/commands/check_firewall_lists.py

@@ -1,36 +1,46 @@
 """Exports process_check_firewall_lists"""
 
+from aikido_zen.helpers.ipc.command_types import Command, Payload, CommandContext
+
+
+class CheckFirewallListsCommand(Command):
+
+    @classmethod
+    def identifier(cls) -> str:
+        return "cfl"  # [C]heck [F]irewall [L]ists
+
+    @classmethod
+    def returns_data(cls) -> bool:
+        return True
+
+    @classmethod
+    def run(cls, ctx: CommandContext, request):
+        ip = request["ip"]
+        if ip is not None and isinstance(ip, str):
+            # Global IP Allowlist (e.g. for geofencing)
+            if not ctx.connection_manager.firewall_lists.is_allowed_ip(ip):
+                return {"blocked": True, "type": "allowlist"}
+
+            # Global IP Blocklist (e.g. blocking known threat actors)
+            reason = ctx.connection_manager.firewall_lists.is_blocked_ip(ip)
+            if reason:
+                return {
+                    "blocked": True,
+                    "type": "blocklist",
+                    "reason": reason,
+                }
 
-def process_check_firewall_lists(connection_manager, data):
-    """
-    Checks whether an IP is blocked
-    data: {"ip": string, "user-agent": string}
-    returns -> {"blocked": boolean, "type": string, "reason": string}
-    """
-    ip = data["ip"]
-    if ip is not None and isinstance(ip, str):
-        # Global IP Allowlist (e.g. for geofencing)
-        if not connection_manager.firewall_lists.is_allowed_ip(ip):
-            return {"blocked": True, "type": "allowlist"}
-
-        # Global IP Blocklist (e.g. blocking known threat actors)
-        reason = connection_manager.firewall_lists.is_blocked_ip(ip)
-        if reason:
-            return {
-                "blocked": True,
-                "type": "blocklist",
-                "reason": reason,
-            }
-
-    user_agent = data["user-agent"]
-    if user_agent is not None and isinstance(user_agent, str):
         # User agent blocking (e.g. blocking AI scrapers)
-        if connection_manager.firewall_lists.is_user_agent_blocked(user_agent):
-            return {
-                "blocked": True,
-                "type": "bot-blocking",
-            }
-
-    return {
-        "blocked": False,
-    }
+        user_agent = request["user-agent"]
+        if user_agent is not None and isinstance(user_agent, str):
+            if ctx.connection_manager.firewall_lists.is_user_agent_blocked(user_agent):
+                return {
+                    "blocked": True,
+                    "type": "bot-blocking",
+                }
+
+        return {"blocked": False}
+
+    @classmethod
+    def generate(cls, request) -> Payload:
+        return Payload(cls, request)

aikido_zen/background_process/commands/ping.py

@@ -1,6 +1,19 @@
-"""exports `process_ping`"""
+from aikido_zen.helpers.ipc.command_types import Command, Payload, CommandContext
 
 
-def process_ping(connection_manager, data):
-    """when main process quits , or during testing etc"""
-    return "Received"
+class PingCommand(Command):
+    @classmethod
+    def identifier(cls) -> str:
+        return "ping"
+
+    @classmethod
+    def returns_data(cls) -> bool:
+        return True
+
+    @classmethod
+    def run(cls, context: CommandContext, request):
+        return "recv"
+
+    @classmethod
+    def generate(cls) -> Payload:
+        return Payload(cls, {})