diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index f011f5ca..6e36c7d1 100644 --- a/aikido_zen/background_process/__init__.py +++ b/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." diff --git a/aikido_zen/background_process/commands/__init__.py b/aikido_zen/background_process/commands/__init__.py index 5b36b758..77f7555e 100644 --- a/aikido_zen/background_process/commands/__init__.py +++ b/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 diff --git a/aikido_zen/background_process/commands/check_firewall_lists.py b/aikido_zen/background_process/commands/check_firewall_lists.py index acf32304..b23948e0 100644 --- a/aikido_zen/background_process/commands/check_firewall_lists.py +++ b/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) diff --git a/aikido_zen/background_process/commands/ping.py b/aikido_zen/background_process/commands/ping.py index 9363eab0..23742878 100644 --- a/aikido_zen/background_process/commands/ping.py +++ b/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, {})