From b807b097b98dbb5e53414fa2be8e83e0a77b8148 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Thu, 20 Mar 2025 14:24:46 +0100 Subject: [PATCH 1/7] Create new process worker script and run it on context creation --- aikido_zen/context/__init__.py | 2 ++ aikido_zen/process_worker/__init__.py | 38 +++++++++++++++++++++++++++ 2 files changed, 40 insertions(+) create mode 100644 aikido_zen/process_worker/__init__.py diff --git a/aikido_zen/context/__init__.py b/aikido_zen/context/__init__.py index e6f51f10a..dae224845 100644 --- a/aikido_zen/context/__init__.py +++ b/aikido_zen/context/__init__.py @@ -14,6 +14,7 @@ from .wsgi import set_wsgi_attributes_on_context from .asgi import set_asgi_attributes_on_context from .extract_route_params import extract_route_params +from .. import process_worker UINPUT_SOURCES = ["body", "cookies", "query", "headers", "xml", "route_params"] current_context = contextvars.ContextVar("current_context", default=None) @@ -37,6 +38,7 @@ class Context: """ def __init__(self, context_obj=None, body=None, req=None, source=None): + process_worker.start_worker() if context_obj: logger.debug("Creating Context instance based on dict object.") self.__dict__.update(context_obj) diff --git a/aikido_zen/process_worker/__init__.py b/aikido_zen/process_worker/__init__.py new file mode 100644 index 000000000..fe635abaf --- /dev/null +++ b/aikido_zen/process_worker/__init__.py @@ -0,0 +1,38 @@ +""" +process worker -> When a web server like gUnicorn makes new processes, and those have multiple threads, +Aikido's process worker is linked to those new processes, so in essence it's 1 extra thread. This thread +is responsible for syncing statistics, route data, ... +""" +import multiprocessing +import threading +import time + +from aikido_zen.helpers.logging import logger + + +def start_worker(): + # Find out the running process: + logger.info("[%s](%s) <-- `%s`", + multiprocessing.current_process().name, + multiprocessing.current_process().pid, + threading.current_thread().name) + + # The name is aikido-process-worker- + the current PID + thread_name = "aikido-process-worker-" + str(multiprocessing.current_process().pid) + if any([thread.name == thread_name for thread in threading.enumerate()]): + return # The thread already exists, returning. + + # Create a new daemon thread tht will handle communication to and from background agent + thread = threading.Thread(target=aikido_process_worker_thread, name=thread_name) + thread.daemon = True + thread.start() + + +def aikido_process_worker_thread(): + # Get the current process + current_process = multiprocessing.current_process() + + while True: + # Print information about the process + logger.info(f"Process ID: {current_process.pid}, Name: {current_process.name}") + time.sleep(5) From b6538681485a3b9a1769b82ee5fe6a66994fbde7 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 21 Mar 2025 10:06:41 +0100 Subject: [PATCH 2/7] Cleanup of process_worker logic --- aikido_zen/context/__init__.py | 2 -- aikido_zen/process_worker/__init__.py | 38 --------------------------- 2 files changed, 40 deletions(-) delete mode 100644 aikido_zen/process_worker/__init__.py diff --git a/aikido_zen/context/__init__.py b/aikido_zen/context/__init__.py index dae224845..e6f51f10a 100644 --- a/aikido_zen/context/__init__.py +++ b/aikido_zen/context/__init__.py @@ -14,7 +14,6 @@ from .wsgi import set_wsgi_attributes_on_context from .asgi import set_asgi_attributes_on_context from .extract_route_params import extract_route_params -from .. import process_worker UINPUT_SOURCES = ["body", "cookies", "query", "headers", "xml", "route_params"] current_context = contextvars.ContextVar("current_context", default=None) @@ -38,7 +37,6 @@ class Context: """ def __init__(self, context_obj=None, body=None, req=None, source=None): - process_worker.start_worker() if context_obj: logger.debug("Creating Context instance based on dict object.") self.__dict__.update(context_obj) diff --git a/aikido_zen/process_worker/__init__.py b/aikido_zen/process_worker/__init__.py deleted file mode 100644 index fe635abaf..000000000 --- a/aikido_zen/process_worker/__init__.py +++ /dev/null @@ -1,38 +0,0 @@ -""" -process worker -> When a web server like gUnicorn makes new processes, and those have multiple threads, -Aikido's process worker is linked to those new processes, so in essence it's 1 extra thread. This thread -is responsible for syncing statistics, route data, ... -""" -import multiprocessing -import threading -import time - -from aikido_zen.helpers.logging import logger - - -def start_worker(): - # Find out the running process: - logger.info("[%s](%s) <-- `%s`", - multiprocessing.current_process().name, - multiprocessing.current_process().pid, - threading.current_thread().name) - - # The name is aikido-process-worker- + the current PID - thread_name = "aikido-process-worker-" + str(multiprocessing.current_process().pid) - if any([thread.name == thread_name for thread in threading.enumerate()]): - return # The thread already exists, returning. - - # Create a new daemon thread tht will handle communication to and from background agent - thread = threading.Thread(target=aikido_process_worker_thread, name=thread_name) - thread.daemon = True - thread.start() - - -def aikido_process_worker_thread(): - # Get the current process - current_process = multiprocessing.current_process() - - while True: - # Print information about the process - logger.info(f"Process ID: {current_process.pid}, Name: {current_process.name}") - time.sleep(5) From e79c8ddfed3c37bfac5b04ab66bdabd9df6bf0c4 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Mon, 12 May 2025 10:55:32 +0200 Subject: [PATCH 3/7] Create new @before_modifies_return --- aikido_zen/sinks/__init__.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/aikido_zen/sinks/__init__.py b/aikido_zen/sinks/__init__.py index 5d3278b66..5ef020450 100644 --- a/aikido_zen/sinks/__init__.py +++ b/aikido_zen/sinks/__init__.py @@ -55,6 +55,27 @@ def decorator(func, instance, args, kwargs): return decorator +def before_modify_return(wrapper): + """ + Surrounds a patch with try-except and calls the original function at the end unless a return value is present. + """ + + def decorator(func, instance, args, kwargs): + try: + rv = wrapper(func, instance, args, kwargs) # Call the patch + if rv is not None: + return rv + except AikidoException as e: + raise e # Re-raise AikidoException + except Exception as e: + logger.debug( + "%s:%s wrapping-before error: %s", func.__module__, func.__name__, e + ) + return func(*args, **kwargs) # Call the original function + + return decorator + + def before_async(wrapper): """ Surrounds an async patch with try-except and calls the original asynchronous function at the end From 69383c0648e08a1fce25ded2701aaf4da362b409 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Mon, 12 May 2025 11:14:09 +0200 Subject: [PATCH 4/7] Cleanup flask.py imports --- aikido_zen/sources/flask.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/sources/flask.py b/aikido_zen/sources/flask.py index cd02b3098..f17d8faf4 100644 --- a/aikido_zen/sources/flask.py +++ b/aikido_zen/sources/flask.py @@ -7,7 +7,7 @@ on_import, patch_function, after, - before_modify_return, + before_modify_return, before, ) From 798187eaea172b3ea470dc993db7bd8784d6f44c Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Mon, 12 May 2025 11:34:34 +0200 Subject: [PATCH 5/7] linting --- aikido_zen/sources/flask.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/sources/flask.py b/aikido_zen/sources/flask.py index f17d8faf4..cd02b3098 100644 --- a/aikido_zen/sources/flask.py +++ b/aikido_zen/sources/flask.py @@ -7,7 +7,7 @@ on_import, patch_function, after, - before_modify_return, + before_modify_return, before, ) From 076ad421b4ea14b088a90da0848b1fa61a8a004d Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Mon, 12 May 2025 12:02:23 +0200 Subject: [PATCH 6/7] remove try-finally for return, which was swallowing error --- aikido_zen/sinks/__init__.py | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/aikido_zen/sinks/__init__.py b/aikido_zen/sinks/__init__.py index 5ef020450..5d3278b66 100644 --- a/aikido_zen/sinks/__init__.py +++ b/aikido_zen/sinks/__init__.py @@ -55,27 +55,6 @@ def decorator(func, instance, args, kwargs): return decorator -def before_modify_return(wrapper): - """ - Surrounds a patch with try-except and calls the original function at the end unless a return value is present. - """ - - def decorator(func, instance, args, kwargs): - try: - rv = wrapper(func, instance, args, kwargs) # Call the patch - if rv is not None: - return rv - except AikidoException as e: - raise e # Re-raise AikidoException - except Exception as e: - logger.debug( - "%s:%s wrapping-before error: %s", func.__module__, func.__name__, e - ) - return func(*args, **kwargs) # Call the original function - - return decorator - - def before_async(wrapper): """ Surrounds an async patch with try-except and calls the original asynchronous function at the end From 2c5bb590767561e0b9ba91bbfb9917c8ae4764b9 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Tue, 13 May 2025 12:41:52 +0200 Subject: [PATCH 7/7] delete importhook/ --- aikido_zen/importhook/__init__.py | 166 ---------------------------- aikido_zen/importhook/finder.py | 55 --------- aikido_zen/importhook/loader.py | 84 -------------- aikido_zen/importhook/meta_paths.py | 15 --- aikido_zen/importhook/registry.py | 41 ------- aikido_zen/importhook/utils.py | 5 - 6 files changed, 366 deletions(-) delete mode 100644 aikido_zen/importhook/__init__.py delete mode 100644 aikido_zen/importhook/finder.py delete mode 100644 aikido_zen/importhook/loader.py delete mode 100644 aikido_zen/importhook/meta_paths.py delete mode 100644 aikido_zen/importhook/registry.py delete mode 100644 aikido_zen/importhook/utils.py diff --git a/aikido_zen/importhook/__init__.py b/aikido_zen/importhook/__init__.py deleted file mode 100644 index a33145a66..000000000 --- a/aikido_zen/importhook/__init__.py +++ /dev/null @@ -1,166 +0,0 @@ -""" -importhook -========== - -Python module for registering hooks to call when certain modules are imported. - -.. code:: python - - import importhook - - # Configure a function to call when `socket` is imported - @importhook.on_import('socket') - def socket_import(socket): - print('Socket module imported') - - # Import the `socket` module - import socket -""" - -import functools -import importlib -import sys -import types - -from .meta_paths import HookMetaPaths -from .registry import registry -from .utils import get_module_name - -__all__ = [ - "ANY_MODULE", - "copy_module", - "get_module_name", - "on_import", - "registry", - "reload_module", - "reset_module", -] - -ANY_MODULE = None - -# Wrap existing (and future) system meta path finders -sys.meta_path = HookMetaPaths(sys.meta_path[:]) - - -def on_import(module_name, func=None): - """ - Helper function used to register a hook function for a given module - - .. code:: python - - import importhook - - @importhook.on_import('socket') - def on_socket_import(socket): - print('socket imported') - - - @importhook.on_import(importhook.ANY_MODULE) - def on_any_import(module): - print(f'{module.__spec__.name} imported') - - - def on_httplib_import(httplib): - print('httplib imported') - - - importhook.on_import('httplib', on_httplib_import) - """ - if func is None: - - @functools.wraps(func) - def decorator(func): - registry[module_name] = func - return func - - return decorator - else: - registry[module_name] = func - - -def reset_module(module_name): - """ - Helper function to reset a copied module. - - .. code:: python - - import socket - import importhook - - # Copy `socket` module - socket = importhook.copy_module(socket) - - # Reset copied `socket` module back to it's original version - socket = importhook.reset_module(socket) - """ - if not isinstance(module_name, str): - module_name = get_module_name(module_name) - - module = sys.modules.get(module_name) - if not module: - return None - - if not hasattr(module, "__original_module__"): - return module - - sys.modules[module_name] = module.__original_module__ - return module.__original_module__ - - -def copy_module(module, copy_attributes=True, copy_spec=True): - """ - Helper function for copying a python module - - .. code:: python - - import importhook - - @importhook.on_import('socket') - def on_socket_import(socket): - new_socket = importhook.copy_module(socket) - setattr(new_socket, 'get_hostname', lambda: 'hostname') - return new_socket - """ - name = get_module_name(module) - new_mod = types.ModuleType(name) - setattr(new_mod, "__original_module__", module) - setattr(new_mod, "__reset_module__", lambda: reset_module(name)) - - # Copy all module attributes - if copy_attributes: - for attr, value in module.__dict__.items(): - setattr(new_mod, attr, value) - - # Make a copy of the modules spec if one is present - if copy_spec and getattr(new_mod, "__spec__", None): - spec = type(new_mod.__spec__)(name=name, loader=new_mod.__spec__.loader) - for attr, value in new_mod.__spec__.__dict__.items(): - if attr not in ("name", "loader"): - setattr(spec, attr, value) - new_mod.__spec__ = spec - return new_mod - - -def reload_module(module_name): - """ - Helper function to reload the specified module - - .. code:: python - - import socket - import importhook - - # Reload the `socket` module by passing in module - socket = importhook.reload_module(socket) - - # Reload the `socket` module by passing in the name - socket = importhook.reload_module('socket') - """ - if not isinstance(module_name, str): - module_name = get_module_name(module_name) - - module = sys.modules.get(module_name) - if not module: - return None - - return importlib.reload(module) diff --git a/aikido_zen/importhook/finder.py b/aikido_zen/importhook/finder.py deleted file mode 100644 index 6ba0938e1..000000000 --- a/aikido_zen/importhook/finder.py +++ /dev/null @@ -1,55 +0,0 @@ -import functools - -from .loader import HookLoader - - -def hook_finder(finder): - """ - Helper function to create a new "hooked" subclass of the provided finder class - - This function replaces the `Finder.find_spec` function to ensure that any ModuleSpecs will - use an `importmod.HookLoader` - """ - # If this finder has already been 'hooked', then return as-is - if hasattr(finder, "__hooked__"): - return finder - - # Determine if we were given an instance or a class - if isinstance(finder, type): - finder_cls = finder - else: - finder_cls = finder.__class__ - - # Determine the class name of the finder - finder_name = finder_cls.__name__ - - def wrap_find_spec(find_spec): - @functools.wraps(find_spec) - def wrapper(fullname, path, target=None): - spec = find_spec(fullname, path, target=target) - if spec is not None and spec.loader is not None: - spec.loader = HookLoader(spec.loader) - return spec - - return wrapper - - def wrap_find_loader(find_loader): - @functools.wraps(find_loader) - def wrapper(fullname, path): - loader = find_loader(fullname, path) - if loader is None: - return None - else: - return HookLoader(loader) - - return wrapper - - # Override the functions we care about - if hasattr(finder, "find_spec"): - setattr(finder, "find_spec", wrap_find_spec(finder.find_spec)) - if hasattr(finder, "find_loader"): - setattr(finder, "find_loader", wrap_find_loader(finder.find_loader)) - - # Make this finder as being 'hooked' - setattr(finder, "__hooked__", True) - return finder diff --git a/aikido_zen/importhook/loader.py b/aikido_zen/importhook/loader.py deleted file mode 100644 index 048ed5b2d..000000000 --- a/aikido_zen/importhook/loader.py +++ /dev/null @@ -1,84 +0,0 @@ -from importlib.abc import Loader -import sys - -from .registry import registry -from .utils import get_module_name - - -def call_module_hooks(module): - name = get_module_name(module) - - # If we have a hook in the registry, then call it now - if name in registry: - mod = registry[name](module) - if mod is not None: - sys.modules[name] = mod - - # If we have a global hook in the registry, then call it now - if None in registry: - mod = registry[None](module) - if mod is not None: - sys.modules[name] = mod - - -class HookLoader(Loader): - """ - Custom `importlib.abc.Loader` which ensures we call any registered hooks when a module is loaded. - """ - - __slots__ = ["loader"] - - def __init__(self, loader): - self.loader = loader - - def __getattribute__(self, name): - # If they are requesting the "loader" attribute, return it right away - loader = super(HookLoader, self).__getattribute__("loader") - if name == "loader": - return loader - - # Pass through attributes/methods only if they exist on the underlying loader - if hasattr(loader, name): - try: - return super(HookLoader, self).__getattribute__(name) - except AttributeError: - return getattr(loader, name) - - raise AttributeError - - def create_module(self, *args, **kwargs): - if not hasattr(self.loader, "create_module"): - return None - - return self.loader.create_module(*args, **kwargs) - - def find_module(self, name, *args, **kwargs): - if not hasattr(self.loader, "find_module"): - return None - - module = self.loader.find_module(name=name, *args, **kwargs) - if module is None: - return None - call_module_hooks(module) - return module - - def load_module(self, name, *args, **kwargs): - if not hasattr(self.loader, "load_module"): - return None - - module = self.loader.load_module(name, *args, **kwargs) - if module is None: - return None - call_module_hooks(module) - return module - - def exec_module(self, module, *args, **kwargs): - if not hasattr(self.loader, "exec_module"): - return None - - mod = self.loader.exec_module(module, *args, **kwargs) - if mod is not None: - module = mod - - call_module_hooks(module) - return module diff --git a/aikido_zen/importhook/meta_paths.py b/aikido_zen/importhook/meta_paths.py deleted file mode 100644 index 102d093ec..000000000 --- a/aikido_zen/importhook/meta_paths.py +++ /dev/null @@ -1,15 +0,0 @@ -from .finder import hook_finder - - -class HookMetaPaths(list): - """ - Custom list that will ensure any items added are wrapped as a "hooked" finder - - This class is made to replace `sys.meta_paths` - """ - - def __init__(self, finders): - super(HookMetaPaths, self).__init__([hook_finder(f) for f in finders]) - - def __setitem__(self, key, val): - super(HookMetaPaths, self).__setitem__(key, hook_finder(val)) diff --git a/aikido_zen/importhook/registry.py b/aikido_zen/importhook/registry.py deleted file mode 100644 index 9c6a1b0f6..000000000 --- a/aikido_zen/importhook/registry.py +++ /dev/null @@ -1,41 +0,0 @@ -import sys - - -class Hooks(set): - """Custom set that is used to maintain and call a list of hooks""" - - def __call__(self, module): - for hook in self: - mod = hook(module) - # If they modified the module, then use that instead - if mod is not None: - module = mod - return module - - -class HookRegistry(dict): - """ - Class used as the registry for import hooks - """ - - def __setitem__(self, key, hook): - module_name = key or "ANY_MODULE" - - # Ensure we have a key for this module and it's value is a `Hooks` set - if key not in self: - super(HookRegistry, self).__setitem__(key, Hooks()) - - # Add our hook to the registry - self[key].add(hook) - - # Call hook for a module which has already been loaded - if key is not None and key in sys.modules: - module = sys.modules[key] - - module = hook(sys.modules[key]) - if module is not None: - sys.modules[key] = module - - -# Create our global registry -registry = HookRegistry() diff --git a/aikido_zen/importhook/utils.py b/aikido_zen/importhook/utils.py deleted file mode 100644 index b3d242804..000000000 --- a/aikido_zen/importhook/utils.py +++ /dev/null @@ -1,5 +0,0 @@ -def get_module_name(module): - """Helper function to get a module's name""" - if hasattr(module, "__spec__"): - return module.__spec__.name - return module.__name__