From 4ac53767b49831cd8168de4fee2259022d01771b Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Tue, 2 Sep 2025 12:54:57 +0200 Subject: [PATCH 01/11] Create a lambda decorator --- aikido_zen/aws_lambda/__init__.py | 6 ---- aikido_zen/lambda/__init__.py | 47 +++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 6 deletions(-) delete mode 100644 aikido_zen/aws_lambda/__init__.py create mode 100644 aikido_zen/lambda/__init__.py diff --git a/aikido_zen/aws_lambda/__init__.py b/aikido_zen/aws_lambda/__init__.py deleted file mode 100644 index 35ee53f98..000000000 --- a/aikido_zen/aws_lambda/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Lambda init.py file""" - - -def protect(handler): - """Aikido protect function for the lambda""" - return handler diff --git a/aikido_zen/lambda/__init__.py b/aikido_zen/lambda/__init__.py new file mode 100644 index 000000000..770781bce --- /dev/null +++ b/aikido_zen/lambda/__init__.py @@ -0,0 +1,47 @@ +from aikido_zen.background_process import get_comms +from aikido_zen.thread.thread_cache import renew as process_cache_renew +from aikido_zen.helpers.logging import logger + +LAMBDA_SEND_HEARTBEAT_TIMEOUT = 500 / 1000 # 500ms + + +def protect_lambda(handler): + """Aikido protect function for the lambda""" + + def wrapper(*args, **kwargs): + + # Before + + rv = handler(*args, **kwargs) + + # After + lambda_post_handler() + + return rv + + return wrapper + + +def lambda_post_handler(): + """ + Lambda post handler, after the lambda is finished, we want to flush the data to core + """ + ipc_client = get_comms() + if not ipc_client: + logger.warning("Lambda: Failed to flush data, no IPC Client available.") + return + + # Flush data from this process + process_cache_renew() + + # Flush data from background process + res = ipc_client.send_data_to_bg_process( + action="SEND_HEARTBEAT", + obj=None, + receive=False, + timeout_in_sec=LAMBDA_SEND_HEARTBEAT_TIMEOUT, + ) + if res["success"]: + logger.info("Lambda: successfully flushed data.") + else: + logger.warning("Lambda: Failed to flush data, error %s.", res["error"]) From 35bd911374f85bcd679277db17b9720887fc392b Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Tue, 2 Sep 2025 12:55:09 +0200 Subject: [PATCH 02/11] Create an IPC SEND_HEARTBEAT call --- aikido_zen/background_process/commands/__init__.py | 2 ++ aikido_zen/background_process/commands/send_heartbeat.py | 5 +++++ 2 files changed, 7 insertions(+) create mode 100644 aikido_zen/background_process/commands/send_heartbeat.py diff --git a/aikido_zen/background_process/commands/__init__.py b/aikido_zen/background_process/commands/__init__.py index 151d25621..62ea69c4e 100644 --- a/aikido_zen/background_process/commands/__init__.py +++ b/aikido_zen/background_process/commands/__init__.py @@ -4,6 +4,7 @@ from .attack import process_attack from .check_firewall_lists import process_check_firewall_lists from .read_property import process_read_property +from .send_heartbeat import process_send_heartbeat from .should_ratelimit import process_should_ratelimit from .ping import process_ping from .sync_data import process_sync_data @@ -12,6 +13,7 @@ # This maps to a tuple : (function, returns_data?) # Commands that don't return data : "ATTACK": (process_attack, False), + "SEND_HEARTBEAT": (process_send_heartbeat, False), # Commands that return data : "SYNC_DATA": (process_sync_data, True), "READ_PROPERTY": (process_read_property, True), diff --git a/aikido_zen/background_process/commands/send_heartbeat.py b/aikido_zen/background_process/commands/send_heartbeat.py new file mode 100644 index 000000000..b7e3c6f5b --- /dev/null +++ b/aikido_zen/background_process/commands/send_heartbeat.py @@ -0,0 +1,5 @@ +def process_send_heartbeat(connection_manager, data, queue): + """ + SEND_HEARTBEAT: Used by lambdas to flush data. + """ + connection_manager.send_heartbeat() From 501f06f5b1a39df3cc98697177edb0ca4a994eee Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Tue, 2 Sep 2025 14:49:29 +0200 Subject: [PATCH 03/11] re-export and add a receive=True --- aikido_zen/__init__.py | 1 + aikido_zen/background_process/commands/send_heartbeat.py | 1 + aikido_zen/{lambda => lambda_helper}/__init__.py | 2 +- 3 files changed, 3 insertions(+), 1 deletion(-) rename aikido_zen/{lambda => lambda_helper}/__init__.py (98%) diff --git a/aikido_zen/__init__.py b/aikido_zen/__init__.py index 37b5ef16e..407888f75 100644 --- a/aikido_zen/__init__.py +++ b/aikido_zen/__init__.py @@ -7,6 +7,7 @@ from aikido_zen.background_process.test_uds_file_access import test_uds_file_access # Re-export functions : +from aikido_zen.lambda_helper import protect_lambda from aikido_zen.context.users import set_user from aikido_zen.middleware import should_block_request diff --git a/aikido_zen/background_process/commands/send_heartbeat.py b/aikido_zen/background_process/commands/send_heartbeat.py index b7e3c6f5b..f78505eda 100644 --- a/aikido_zen/background_process/commands/send_heartbeat.py +++ b/aikido_zen/background_process/commands/send_heartbeat.py @@ -3,3 +3,4 @@ def process_send_heartbeat(connection_manager, data, queue): SEND_HEARTBEAT: Used by lambdas to flush data. """ connection_manager.send_heartbeat() + return {"status": "Heartbeat sent"} diff --git a/aikido_zen/lambda/__init__.py b/aikido_zen/lambda_helper/__init__.py similarity index 98% rename from aikido_zen/lambda/__init__.py rename to aikido_zen/lambda_helper/__init__.py index 770781bce..205931128 100644 --- a/aikido_zen/lambda/__init__.py +++ b/aikido_zen/lambda_helper/__init__.py @@ -38,7 +38,7 @@ def lambda_post_handler(): res = ipc_client.send_data_to_bg_process( action="SEND_HEARTBEAT", obj=None, - receive=False, + receive=True, timeout_in_sec=LAMBDA_SEND_HEARTBEAT_TIMEOUT, ) if res["success"]: From 55e6cf02bf5983e1737f1361f6af232e0920bc8e Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 3 Sep 2025 14:23:25 +0200 Subject: [PATCH 04/11] Add log & debug options for aws lambda --- aikido_zen/background_process/__init__.py | 2 ++ aikido_zen/background_process/aikido_background_process.py | 1 + aikido_zen/background_process/commands/send_heartbeat.py | 4 ++++ aikido_zen/background_process/comms.py | 1 + 4 files changed, 8 insertions(+) diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 71f199e39..b9c16467d 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -37,6 +37,8 @@ def start_background_process(): if platform.system() == "Windows": # Python does not support Windows UDS just yet, so we have to rely on INET address = ("127.0.0.1", 49156) + if os.getenv("AIKIDO_INET_ONLY") is "1": + address = ("127.0.0.1", 49156) comms = AikidoIPCCommunications(address, secret_key_bytes) diff --git a/aikido_zen/background_process/aikido_background_process.py b/aikido_zen/background_process/aikido_background_process.py index dc0162110..6883d43fc 100644 --- a/aikido_zen/background_process/aikido_background_process.py +++ b/aikido_zen/background_process/aikido_background_process.py @@ -50,6 +50,7 @@ def __init__(self, address, key): conn = listener.accept() while True: try: + logger.debug("Waiting for data...") data = conn.recv() # because of this no sleep needed in thread logger.debug("Incoming data : %s", data) process_incoming_command( diff --git a/aikido_zen/background_process/commands/send_heartbeat.py b/aikido_zen/background_process/commands/send_heartbeat.py index f78505eda..a23d4fce9 100644 --- a/aikido_zen/background_process/commands/send_heartbeat.py +++ b/aikido_zen/background_process/commands/send_heartbeat.py @@ -1,6 +1,10 @@ +from aikido_zen.helpers.logging import logger + def process_send_heartbeat(connection_manager, data, queue): """ SEND_HEARTBEAT: Used by lambdas to flush data. """ + logger.debug("SEND_HEARTBEAT start [->]") connection_manager.send_heartbeat() + logger.debug("SEND_HEARTBEAT END [<-]") return {"status": "Heartbeat sent"} diff --git a/aikido_zen/background_process/comms.py b/aikido_zen/background_process/comms.py index 48b7e402e..a2b907988 100644 --- a/aikido_zen/background_process/comms.py +++ b/aikido_zen/background_process/comms.py @@ -77,6 +77,7 @@ def target(address, key, receive, data, result_obj): result_obj[0] = True # Connection ended gracefully except Exception as e: logger.debug("Exception occurred in thread : %s", e) + logger.debug("Address: %s", address) # Create a shared result object between the thread and this process : result_obj = [False, None] # Needs to be an array so we can make a ref. From 3456facb9d9d1883788e5d3526a67c70d63b251f Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 3 Sep 2025 14:36:07 +0200 Subject: [PATCH 05/11] lint --- aikido_zen/background_process/commands/send_heartbeat.py | 1 + 1 file changed, 1 insertion(+) diff --git a/aikido_zen/background_process/commands/send_heartbeat.py b/aikido_zen/background_process/commands/send_heartbeat.py index a23d4fce9..b2560bcb0 100644 --- a/aikido_zen/background_process/commands/send_heartbeat.py +++ b/aikido_zen/background_process/commands/send_heartbeat.py @@ -1,5 +1,6 @@ from aikido_zen.helpers.logging import logger + def process_send_heartbeat(connection_manager, data, queue): """ SEND_HEARTBEAT: Used by lambdas to flush data. From cdf83915f4df5bbf8956f5ae73fb109504bb2fa1 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Tue, 14 Oct 2025 15:41:56 +0200 Subject: [PATCH 06/11] use contextvars for circular import check --- aikido_zen/sinks/builtins_import.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/aikido_zen/sinks/builtins_import.py b/aikido_zen/sinks/builtins_import.py index ba646f295..0a2585ee0 100644 --- a/aikido_zen/sinks/builtins_import.py +++ b/aikido_zen/sinks/builtins_import.py @@ -1,12 +1,18 @@ +from aikido_zen.sinks import on_import, patch_function, after +import contextvars import importlib.metadata from importlib.metadata import PackageNotFoundError - from aikido_zen.background_process.packages import PackagesStore -from aikido_zen.sinks import on_import, patch_function, after + +running_import_scan = contextvars.ContextVar("running_import_scan", default=False) @after def _import(func, instance, args, kwargs, return_value): + if running_import_scan.get(): + return + running_import_scan.set(True) + if not hasattr(return_value, "__file__"): return # Would be built-in into the interpreter (system package) @@ -18,9 +24,6 @@ def _import(func, instance, args, kwargs, return_value): # Make sure the name exists return name = name.split(".")[0] # Remove submodules - if name == "importlib" or name == "importlib_metadata": - # Avoid circular dependencies - return if PackagesStore.get_package(name): return @@ -33,6 +36,8 @@ def _import(func, instance, args, kwargs, return_value): if version: PackagesStore.add_package(name, version) + running_import_scan.set(False) + @on_import("builtins") def patch(m): From f1698cfa7d2e46080ec1fc7f660665669af0ee27 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Tue, 14 Oct 2025 15:45:54 +0200 Subject: [PATCH 07/11] add previous recursion check back in --- aikido_zen/sinks/builtins_import.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aikido_zen/sinks/builtins_import.py b/aikido_zen/sinks/builtins_import.py index 0a2585ee0..42162e874 100644 --- a/aikido_zen/sinks/builtins_import.py +++ b/aikido_zen/sinks/builtins_import.py @@ -24,6 +24,9 @@ def _import(func, instance, args, kwargs, return_value): # Make sure the name exists return name = name.split(".")[0] # Remove submodules + if name == "importlib" or name == "importlib_metadata": + # Avoid circular dependencies, this is a double safety-check for if contextvar check fails. + return if PackagesStore.get_package(name): return From 61c4f8497cb2e332d4d9d8ee72d203bf899486a2 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 15 Oct 2025 11:59:43 +0200 Subject: [PATCH 08/11] Add a try-finally for builtins_import --- aikido_zen/sinks/builtins_import.py | 61 +++++++++++++++-------------- 1 file changed, 31 insertions(+), 30 deletions(-) diff --git a/aikido_zen/sinks/builtins_import.py b/aikido_zen/sinks/builtins_import.py index 42162e874..34dc9d854 100644 --- a/aikido_zen/sinks/builtins_import.py +++ b/aikido_zen/sinks/builtins_import.py @@ -9,37 +9,38 @@ @after def _import(func, instance, args, kwargs, return_value): - if running_import_scan.get(): - return - running_import_scan.set(True) - - if not hasattr(return_value, "__file__"): - return # Would be built-in into the interpreter (system package) - - if not hasattr(return_value, "__package__"): - return - name = getattr(return_value, "__package__") - - if not name: - # Make sure the name exists - return - name = name.split(".")[0] # Remove submodules - if name == "importlib" or name == "importlib_metadata": - # Avoid circular dependencies, this is a double safety-check for if contextvar check fails. - return - - if PackagesStore.get_package(name): - return - - version = None try: - version = importlib.metadata.version(name) - except PackageNotFoundError: - pass - if version: - PackagesStore.add_package(name, version) - - running_import_scan.set(False) + if running_import_scan.get(): + return + running_import_scan.set(True) + + if not hasattr(return_value, "__file__"): + return # Would be built-in into the interpreter (system package) + + if not hasattr(return_value, "__package__"): + return + name = getattr(return_value, "__package__") + + if not name: + # Make sure the name exists + return + name = name.split(".")[0] # Remove submodules + if name == "importlib" or name == "importlib_metadata": + # Avoid circular dependencies, this is a double safety-check for if contextvar check fails. + return + + if PackagesStore.get_package(name): + return + + version = None + try: + version = importlib.metadata.version(name) + except PackageNotFoundError: + pass + if version: + PackagesStore.add_package(name, version) + finally: + running_import_scan.set(False) @on_import("builtins") From af204b5a8f23bdaf546732639573e7f4469aa335 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 15 Oct 2025 12:05:12 +0200 Subject: [PATCH 09/11] add test cases for builtins import --- .../sinks/tests/builtins_import_test.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 aikido_zen/sinks/tests/builtins_import_test.py diff --git a/aikido_zen/sinks/tests/builtins_import_test.py b/aikido_zen/sinks/tests/builtins_import_test.py new file mode 100644 index 000000000..ca7cfd3b5 --- /dev/null +++ b/aikido_zen/sinks/tests/builtins_import_test.py @@ -0,0 +1,24 @@ +import pytest + +import aikido_zen + +aikido_zen.protect() + +from aikido_zen.background_process.packages import PackagesStore + + +@pytest.fixture(autouse=True) +def run_around_tests(): + PackagesStore.clear() + + +def test_flask_import(): + import flask + + assert PackagesStore.get_package("flask")["version"] == "3.0.3" + + +def test_django_import(): + import django + + assert PackagesStore.get_package("django")["version"] == "4.0" From 1eb37871695e3df1f732140fd7a9d9c1cffdb795 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 29 Oct 2025 19:30:12 +0100 Subject: [PATCH 10/11] Fix merge conflict --- aikido_zen/background_process/__init__.py | 2 +- aikido_zen/sinks/builtins_import.py | 8 -------- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index b0eb9167b..a3f0d38d4 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -37,7 +37,7 @@ def start_background_process(): if platform.system() == "Windows": # Python does not support Windows UDS just yet, so we have to rely on INET address = ("127.0.0.1", 49156) - if os.getenv("AIKIDO_INET_ONLY") is "1": + if os.getenv("AIKIDO_INET_ONLY") == "1": address = ("127.0.0.1", 49156) comms = AikidoIPCCommunications(address, secret_key_bytes) diff --git a/aikido_zen/sinks/builtins_import.py b/aikido_zen/sinks/builtins_import.py index f6c0dd138..5a752e80b 100644 --- a/aikido_zen/sinks/builtins_import.py +++ b/aikido_zen/sinks/builtins_import.py @@ -9,14 +9,6 @@ @after def _import(func, instance, args, kwargs, return_value): - try: - if running_import_scan.get(): - return - running_import_scan.set(True) - - if not hasattr(return_value, "__package__"): - return - try: if running_import_scan.get(): return From 2c017d5278e23326a4127f16d29c8c35d2514dff Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 29 Oct 2025 19:33:06 +0100 Subject: [PATCH 11/11] Fix merge conflict --- aikido_zen/sinks/builtins_import.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/aikido_zen/sinks/builtins_import.py b/aikido_zen/sinks/builtins_import.py index 5a752e80b..fe6a00a2c 100644 --- a/aikido_zen/sinks/builtins_import.py +++ b/aikido_zen/sinks/builtins_import.py @@ -9,6 +9,12 @@ @after def _import(func, instance, args, kwargs, return_value): + if not hasattr(return_value, "__file__"): + return # Would be built-in into the interpreter (system package) + + if not hasattr(return_value, "__package__"): + return + try: if running_import_scan.get(): return