From 12f0668b2f786800cf060fd5e22488df794963fb Mon Sep 17 00:00:00 2001
From: Wout Feys <bitterpanda@proton.me>
Date: Thu, 20 Mar 2025 14:24:46 +0100
Subject: [PATCH 01/25] 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 3021f26a013436f14ecec3e1e92315b2108eadaf Mon Sep 17 00:00:00 2001
From: BitterPanda63 <bitterpanda@proton.me>
Date: Fri, 21 Mar 2025 10:06:41 +0100
Subject: [PATCH 02/25] 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 808ee67c17b0d29c25ba141838a67f79b72e2ec7 Mon Sep 17 00:00:00 2001
From: BitterPanda63 <bitterpanda@proton.me>
Date: Tue, 15 Apr 2025 16:43:54 +0200
Subject: [PATCH 03/25] Add extra tests for psycopg2 and fix issue for
 immutables

---
 aikido_zen/sinks/tests/psycopg2_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/aikido_zen/sinks/tests/psycopg2_test.py b/aikido_zen/sinks/tests/psycopg2_test.py
index a888a4fef..f53844876 100644
--- a/aikido_zen/sinks/tests/psycopg2_test.py
+++ b/aikido_zen/sinks/tests/psycopg2_test.py
@@ -59,6 +59,7 @@ def test_cursor_execute(database_conn):
         "aikido_zen.vulnerabilities.run_vulnerability_scan"
     ) as mock_run_vulnerability_scan:
         cursor = database_conn.cursor()
+        print(cursor)
         query = "SELECT * FROM dogs"
         cursor.execute(query)
 

From 4a7526c2f1959f40cb272937ec470c983c8ab1bf Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Fri, 9 May 2025 17:21:45 +0200
Subject: [PATCH 04/25] Remove package tracking for gunicorn and uwsgi (should
 be standard)

---
 aikido_zen/__init__.py         |  3 ---
 aikido_zen/sources/gunicorn.py | 11 -----------
 aikido_zen/sources/uwsgi.py    | 11 -----------
 3 files changed, 25 deletions(-)
 delete mode 100644 aikido_zen/sources/gunicorn.py
 delete mode 100644 aikido_zen/sources/uwsgi.py

diff --git a/aikido_zen/__init__.py b/aikido_zen/__init__.py
index 6c9814055..72ce58c5d 100644
--- a/aikido_zen/__init__.py
+++ b/aikido_zen/__init__.py
@@ -50,9 +50,6 @@ def protect(mode="daemon"):
     import aikido_zen.sources.xml_sources.xml
     import aikido_zen.sources.xml_sources.lxml
 
-    import aikido_zen.sources.gunicorn
-    import aikido_zen.sources.uwsgi
-
     # Import DB Sinks
     import aikido_zen.sinks.pymysql
     import aikido_zen.sinks.mysqlclient
diff --git a/aikido_zen/sources/gunicorn.py b/aikido_zen/sources/gunicorn.py
deleted file mode 100644
index caa704e80..000000000
--- a/aikido_zen/sources/gunicorn.py
+++ /dev/null
@@ -1,11 +0,0 @@
-"""Gunicorn Module, report if module was found"""
-
-import aikido_zen.importhook as importhook
-from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
-
-
-@importhook.on_import("gunicorn")
-def on_gunicorn_import(gunicorn):
-    """Report to the core when gunicorn gets imported"""
-    is_package_compatible("gunicorn", required_version=ANY_VERSION)
-    return gunicorn
diff --git a/aikido_zen/sources/uwsgi.py b/aikido_zen/sources/uwsgi.py
deleted file mode 100644
index f2a61f4ee..000000000
--- a/aikido_zen/sources/uwsgi.py
+++ /dev/null
@@ -1,11 +0,0 @@
-"""UWSGI Module, report if module was found"""
-
-import aikido_zen.importhook as importhook
-from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
-
-
-@importhook.on_import("uwsgi")
-def on_uwsgi_import(uwsgi):
-    """Report to the core when uwsgi gets imported"""
-    is_package_compatible("uwsgi", required_version=ANY_VERSION)
-    return uwsgi

From 2479d003aa291072b1e609dd69d48fe5858c3588 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Fri, 9 May 2025 17:31:53 +0200
Subject: [PATCH 05/25] Update django wrapper to use wrapt

---
 aikido_zen/sources/django/__init__.py | 47 +++++++++++++--------------
 1 file changed, 22 insertions(+), 25 deletions(-)

diff --git a/aikido_zen/sources/django/__init__.py b/aikido_zen/sources/django/__init__.py
index 9c16ed579..273fc7b71 100644
--- a/aikido_zen/sources/django/__init__.py
+++ b/aikido_zen/sources/django/__init__.py
@@ -7,36 +7,33 @@
 from ..functions.request_handler import request_handler
 from .run_init_stage import run_init_stage
 from .pre_response_middleware import pre_response_middleware
+from ...helpers.get_argument import get_argument
+from ...sinks import on_import, patch_function, before, after
 
 
-@importhook.on_import("django.core.handlers.base")
-def on_django_gunicorn_import(django):
-    """
-    Hook 'n wrap on `django.core.handlers.base`
-    Our goal is to wrap the `_get_response` function
-    # https://github.com/django/django/blob/5865ff5adcf64da03d306dc32b36e87ae6927c85/django/core/handlers/base.py#L174
-    Returns : Modified django.core.handlers.base object
-    """
-    if not is_package_compatible("django", required_version=ANY_VERSION):
-        return django
-    modified_django = importhook.copy_module(django)
+@before
+def _get_response_before(func, instance, args, kwargs):
+    request = get_argument(args, kwargs, 0, "request")
 
-    former__get_response = copy.deepcopy(django.BaseHandler._get_response)
+    run_init_stage(request)
 
-    def aikido__get_response(self, request):  # Synchronous (WSGI)
-        run_init_stage(request)  # We do some initial request handling
+    if pre_response_middleware not in instance._view_middleware:
+        # The rate limiting middleware needs to be last in the chain.
+        instance._view_middleware += [pre_response_middleware]
 
-        if pre_response_middleware not in self._view_middleware:
-            # The rate limiting middleware needs to be last in the chain.
-            self._view_middleware += [pre_response_middleware]
 
-        res = former__get_response(self, request)
-        if hasattr(res, "status_code"):
-            request_handler(stage="post_response", status_code=res.status_code)
-        return res
+@after
+def _get_response_after(func, instance, args, kwargs, return_value):
+    if hasattr(return_value, "status_code"):
+        request_handler(stage="post_response", status_code=return_value.status_code)
 
-    # pylint: disable=no-member
-    setattr(modified_django.BaseHandler, "_get_response", aikido__get_response)
-    setattr(django.BaseHandler, "_get_response", aikido__get_response)
 
-    return modified_django
+@on_import("django.core.handlers.base")
+def patch(m):
+    """
+    Patch for _get_response
+    - before: Parse body, create context & add middleware to run before a response
+    # https://github.com/django/django/blob/5865ff5adcf64da03d306dc32b36e87ae6927c85/django/core/handlers/base.py#L174
+    """
+    patch_function(m, "BaseHandler._get_response", _get_response_before)
+    patch_function(m, "BaseHandler._get_response", _get_response_after)

From 793a1fe8cde983cd91f27e16913f4d98e4d02198 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 09:48:47 +0200
Subject: [PATCH 06/25] Django cleanup and expand comments

---
 aikido_zen/sources/django/__init__.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/aikido_zen/sources/django/__init__.py b/aikido_zen/sources/django/__init__.py
index 273fc7b71..912960646 100644
--- a/aikido_zen/sources/django/__init__.py
+++ b/aikido_zen/sources/django/__init__.py
@@ -1,9 +1,5 @@
 """`Django` source module"""
 
-import copy
-import aikido_zen.importhook as importhook
-from aikido_zen.helpers.logging import logger
-from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
 from ..functions.request_handler import request_handler
 from .run_init_stage import run_init_stage
 from .pre_response_middleware import pre_response_middleware
@@ -31,8 +27,9 @@ def _get_response_after(func, instance, args, kwargs, return_value):
 @on_import("django.core.handlers.base")
 def patch(m):
     """
-    Patch for _get_response
+    Patch for _get_response (Synchronous/WSGI)
     - before: Parse body, create context & add middleware to run before a response
+    - after: Check respone code to see if route should be analyzed
     # https://github.com/django/django/blob/5865ff5adcf64da03d306dc32b36e87ae6927c85/django/core/handlers/base.py#L174
     """
     patch_function(m, "BaseHandler._get_response", _get_response_before)

From 0ff814021d59a272248f744d28ca360dd7471b49 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:00:59 +0200
Subject: [PATCH 07/25] Create a before_async decorator for patching

---
 aikido_zen/sinks/__init__.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/aikido_zen/sinks/__init__.py b/aikido_zen/sinks/__init__.py
index a551ef8af..373381a16 100644
--- a/aikido_zen/sinks/__init__.py
+++ b/aikido_zen/sinks/__init__.py
@@ -50,8 +50,28 @@ def decorator(func, instance, args, kwargs):
             logger.debug(
                 "%s:%s wrapping-before error: %s", func.__module__, func.__name__, e
             )
+        finally:
+            return func(*args, **kwargs)  # Call the original function
 
-        return func(*args, **kwargs)  # Call the original function
+    return decorator
+
+
+def before_async(wrapper):
+    """
+    Surrounds a patch with try-except and calls the original function at the end (async)
+    """
+
+    async def decorator(func, instance, args, kwargs):
+        try:
+            await wrapper(func, instance, args, kwargs)  # Call the patch
+        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
+            )
+        finally:
+            return await func(*args, **kwargs)  # Call the original function
 
     return decorator
 

From 011963cc04a17da675189efef42ffe0e6791c98b Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:01:13 +0200
Subject: [PATCH 08/25] Add "django" as a module name

---
 aikido_zen/sources/django/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aikido_zen/sources/django/__init__.py b/aikido_zen/sources/django/__init__.py
index 912960646..91c92d29d 100644
--- a/aikido_zen/sources/django/__init__.py
+++ b/aikido_zen/sources/django/__init__.py
@@ -24,7 +24,7 @@ def _get_response_after(func, instance, args, kwargs, return_value):
         request_handler(stage="post_response", status_code=return_value.status_code)
 
 
-@on_import("django.core.handlers.base")
+@on_import("django.core.handlers.base", "django")
 def patch(m):
     """
     Patch for _get_response (Synchronous/WSGI)

From bb37d041e33aa17de320d55644ca498be0fd1321 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:01:30 +0200
Subject: [PATCH 09/25] Cleanup starlette's __init__.py code

---
 aikido_zen/sources/starlette/__init__.py | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/aikido_zen/sources/starlette/__init__.py b/aikido_zen/sources/starlette/__init__.py
index e22cb1891..c00bdd55e 100644
--- a/aikido_zen/sources/starlette/__init__.py
+++ b/aikido_zen/sources/starlette/__init__.py
@@ -1,6 +1,4 @@
 """
-Init.py file for starlette module
----
 Starlette wrapping is subdivided in two parts :
 - starlette.applications : Wraps __call__ on Starlette class to run "init" stage.
 - starlette.routing : request_response function : Run pre_response code and 
@@ -10,19 +8,3 @@
 - extract_data_from_request, which will extract the data from a request object safely,
     e.g. body, json, form. This also saves it inside the current context.
 """
-
-import aikido_zen.importhook as importhook
-from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
-
-
-@importhook.on_import("starlette")
-def on_starlette_import(starlette):
-    """
-    This checks for the package version of starlette so you don't have to do it twice,
-    once in starlette_applications and once in starlette_applications.
-    """
-    if not is_package_compatible("starlette", required_version=ANY_VERSION):
-        return starlette
-    # Package is compatible, start wrapping :
-    import aikido_zen.sources.starlette.starlette_applications
-    import aikido_zen.sources.starlette.starlette_routing

From 929b4f27d5a6133f61bc73634a01e694fadbf1ae Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:01:41 +0200
Subject: [PATCH 10/25] Refactor starlette.applications wrapping

---
 .../starlette/starlette_applications.py       | 44 +++++++------------
 1 file changed, 17 insertions(+), 27 deletions(-)

diff --git a/aikido_zen/sources/starlette/starlette_applications.py b/aikido_zen/sources/starlette/starlette_applications.py
index ce5ca8f66..8a3f5b68b 100644
--- a/aikido_zen/sources/starlette/starlette_applications.py
+++ b/aikido_zen/sources/starlette/starlette_applications.py
@@ -1,36 +1,26 @@
 """Wraps starlette.applications for initial request_handler"""
 
-import copy
-import aikido_zen.importhook as importhook
-from aikido_zen.helpers.logging import logger
 from aikido_zen.context import Context
 from ..functions.request_handler import request_handler
+from ...helpers.get_argument import get_argument
+from ...sinks import on_import, patch_function, before_async
 
 
-@importhook.on_import("starlette.applications")
-def on_starlette_import(starlette):
-    """
-    Hook 'n wrap on `starlette.applications`
-    Our goal is to wrap the __call__ function of the Starlette class
-    """
-    modified_starlette = importhook.copy_module(starlette)
-    former_call = copy.deepcopy(starlette.Starlette.__call__)
-
-    async def aikido___call__(app, scope, receive=None, send=None):
-        return await aik_call_wrapper(former_call, app, scope, receive, send)
+@before_async
+async def _call(func, instance, args, kwargs):
+    scope = get_argument(args, kwargs, 1, "scope")
+    if scope.get("type") != "http":
+        return
 
-    setattr(modified_starlette.Starlette, "__call__", aikido___call__)
-    return modified_starlette
+    new_context = Context(req=scope, source="starlette")
+    new_context.set_as_current_context()
+    request_handler(stage="init")
 
 
-async def aik_call_wrapper(former_call, app, scope, receive, send):
-    """Aikido's __call__ wrapper"""
-    try:
-        if scope["type"] != "http":
-            return await former_call(app, scope, receive, send)
-        context1 = Context(req=scope, source="starlette")
-        context1.set_as_current_context()
-        request_handler(stage="init")
-    except Exception as e:
-        logger.debug("Exception on aikido __call__ function : %s", e)
-    return await former_call(app, scope, receive, send)
+@on_import("starlette.applications", "starlette")
+def on_starlette_import(starlette):
+    """
+    patches `starlette.applications`
+    - patching Starlette.__call__
+    """
+    patch_function("Starlette", "__call__", _call)

From 10cb78a2783a8443a6023ba69a1bc393e33829b0 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:13:33 +0200
Subject: [PATCH 11/25] Fix previous mistakes with module not being ref'd in
 starlette_app's.py

---
 aikido_zen/sources/starlette/starlette_applications.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/aikido_zen/sources/starlette/starlette_applications.py b/aikido_zen/sources/starlette/starlette_applications.py
index 8a3f5b68b..84dea4911 100644
--- a/aikido_zen/sources/starlette/starlette_applications.py
+++ b/aikido_zen/sources/starlette/starlette_applications.py
@@ -18,9 +18,9 @@ async def _call(func, instance, args, kwargs):
 
 
 @on_import("starlette.applications", "starlette")
-def on_starlette_import(starlette):
+def patch(m):
     """
-    patches `starlette.applications`
-    - patching Starlette.__call__
+    patching module starlette.applications
+    - patches: Starlette.__call__
     """
-    patch_function("Starlette", "__call__", _call)
+    patch_function(m, "Starlette.__call__", _call)

From af7a8923cf9fc7bc867862d50abe2f685449408d Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:13:48 +0200
Subject: [PATCH 12/25] Port the starlette_routing wrapper

---
 .../sources/starlette/starlette_routing.py    | 39 +++++++------------
 1 file changed, 15 insertions(+), 24 deletions(-)

diff --git a/aikido_zen/sources/starlette/starlette_routing.py b/aikido_zen/sources/starlette/starlette_routing.py
index d65dc63e0..3c79351c2 100644
--- a/aikido_zen/sources/starlette/starlette_routing.py
+++ b/aikido_zen/sources/starlette/starlette_routing.py
@@ -1,37 +1,28 @@
-"""
-Wraps starlette.applications for initial request_handler
-Attention: We will be using rr to refer to request_response. It's used a lot and
-readability would be impaired if we did not abbreviate this
-"""
-
-import copy
-import aikido_zen.importhook as importhook
 from aikido_zen.helpers.logging import logger
 from .extract_data_from_request import extract_data_from_request
 from ..functions.request_handler import request_handler
+from ...sinks import on_import, patch_function, before
 
 
-@importhook.on_import("starlette.routing")
-def on_starlette_import(routing):
-    """
-    Hook 'n wrap on `starlette.routing`
-    Wraps the request_response function so we can wrap the function given to request_response
-    """
-    modified_routing = importhook.copy_module(routing)
-    former_rr_func = copy.deepcopy(routing.request_response)
+@before
+def _request_response(func, instance, args, kwargs):
+    if kwargs and "func" in kwargs:
+        kwargs["func"] = aik_route_func_wrapper(kwargs["func"])
+    elif args and args[0]:
+        args[0] = aik_route_func_wrapper(args[0])
 
-    def aikido_rr_func(func):
-        wrapped_route_func = aik_route_func_wrapper(func)
-        return former_rr_func(wrapped_route_func)
 
-    setattr(routing, "request_response", aikido_rr_func)
-    setattr(modified_routing, "request_response", aikido_rr_func)
-    return modified_routing
+@on_import("starlette.routing")
+def patch(m):
+    """
+    patching module starlette.routing (for initial request_handler)
+    - patches: request_response
+    (github src: https://github.com/encode/starlette/blob/4acf1d1ca3e8aa767567cb4e6e12f093f066553b/starlette/routing.py#L58)
+    """
+    patch_function(m, "request_response", _request_response)
 
 
 def aik_route_func_wrapper(func):
-    """Aikido's __call__ wrapper"""
-
     async def aikido_route_func(*args, **kwargs):
         # Code before response (pre_response stage)
         try:

From 635c7a9f756b0df1de3efb97cb8ef0ce21a636c8 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:16:43 +0200
Subject: [PATCH 13/25] Also mention starlette package in starlette_routing

---
 aikido_zen/sources/starlette/starlette_routing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aikido_zen/sources/starlette/starlette_routing.py b/aikido_zen/sources/starlette/starlette_routing.py
index 3c79351c2..48bebc218 100644
--- a/aikido_zen/sources/starlette/starlette_routing.py
+++ b/aikido_zen/sources/starlette/starlette_routing.py
@@ -12,7 +12,7 @@ def _request_response(func, instance, args, kwargs):
         args[0] = aik_route_func_wrapper(args[0])
 
 
-@on_import("starlette.routing")
+@on_import("starlette.routing", "starlette")
 def patch(m):
     """
     patching module starlette.routing (for initial request_handler)

From c82f1bf2b8973b2ef33f0b459a6511666ca1b2bd Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 10:55:32 +0200
Subject: [PATCH 14/25] 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 373381a16..d3978a334 100644
--- a/aikido_zen/sinks/__init__.py
+++ b/aikido_zen/sinks/__init__.py
@@ -56,6 +56,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 a patch with try-except and calls the original function at the end (async)

From 44f6d5378c4fbe6ed84898b9c9b00dbf4b676b27 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 11:00:53 +0200
Subject: [PATCH 15/25] Port flask.py to the new wrapt way

---
 aikido_zen/sources/flask.py | 101 ++++++++++++++++--------------------
 1 file changed, 44 insertions(+), 57 deletions(-)

diff --git a/aikido_zen/sources/flask.py b/aikido_zen/sources/flask.py
index 6fc862749..6f354afdf 100644
--- a/aikido_zen/sources/flask.py
+++ b/aikido_zen/sources/flask.py
@@ -2,47 +2,50 @@
 Flask source module, intercepts flask import and adds Aikido middleware
 """
 
-import copy
 import aikido_zen.importhook as importhook
+from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.helpers.logging import logger
 from aikido_zen.context import Context
 from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
 from aikido_zen.context import get_current_context
 import aikido_zen.sources.functions.request_handler as funcs
+from aikido_zen.sinks import (
+    on_import,
+    patch_function,
+    after,
+    before_modify_return,
+    before,
+)
 
 
-def aik_full_dispatch_request(*args, former_full_dispatch_request=None, **kwargs):
-    """
-    Creates a new full_dispatch_request function :
-    https://github.com/pallets/flask/blob/2fec0b206c6e83ea813ab26597e15c96fab08be7/src/flask/app.py#L884
-    This function gets called in the wsgi_app. So this function onlygets called after all the
-    middleware. This is important since we want to be able to access users. This also means the
-    request in request_ctx is available and we can extract data from it This function also
-    returns a response, so we can send status codes and error messages.
-    """
-    # pylint:disable=import-outside-toplevel # We don't want to install this by default
-    try:
-        from flask.globals import request_ctx
-        from flask import Response
-    except ImportError:
-        logger.info("Flask not properly installed.")
-        return former_full_dispatch_request(*args, **kwargs)
+@before_modify_return
+def _full_dispatch_request_before(func, instance, args, kwargs):
+    from flask.globals import request_ctx
+    from flask import Response
 
     req = request_ctx.request
-
     extract_cookies_from_flask_request_and_save_data(req)
     extract_form_data_from_flask_request_and_save_data(req)
     extract_view_args_from_flask_request_and_save_data(req)
 
     pre_response = funcs.request_handler(stage="pre_response")
-    if pre_response:
-        # This happens when a route is rate limited, a user blocked, etc...
-        return Response(pre_response[0], status=pre_response[1], mimetype="text/plain")
-    res = former_full_dispatch_request(*args, **kwargs)
-    funcs.request_handler(stage="post_response", status_code=res.status_code)
+    if not pre_response:
+        return None
+    # This happens when a route is rate limited, a user blocked, etc...
+    res = Response(pre_response[0], status=pre_response[1], mimetype="text/plain")
+    res.is_aikido_response = True  # Use this in the @after
     return res
 
 
+@after
+def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
+    if hasattr(return_value, "is_aikido_response") and return_value.is_aikido_response:
+        return  # We wrote this response, so the status code is not reflective of the route
+    if not hasattr(return_value, "status_code"):
+        return
+    funcs.request_handler(stage="post_response", status_code=return_value.status_code)
+
+
 def extract_view_args_from_flask_request_and_save_data(req):
     """Extract view args from flask request"""
     context = get_current_context()
@@ -79,41 +82,25 @@ def extract_cookies_from_flask_request_and_save_data(req):
         logger.debug("Exception occurred whilst extracting flask cookie data: %s", e)
 
 
-def aikido___call__(flask_app, environ, start_response):
-    """Aikido's __call__ wrapper"""
-    # We don't want to install werkzeug :
-    # pylint: disable=import-outside-toplevel
-    try:
-        context1 = Context(req=environ, source="flask")
-        context1.set_as_current_context()
-        funcs.request_handler(stage="init")
-    except Exception as e:
-        logger.debug("Exception on aikido __call__ function : %s", e)
-    res = flask_app.wsgi_app(environ, start_response)
-    return res
+@before
+def _call(func, instance, args, kwargs):
+    environ = get_argument(args, kwargs, 1, "environ")
+    context1 = Context(req=environ, source="flask")
+    context1.set_as_current_context()
+    funcs.request_handler(stage="init")
 
 
-FLASK_REQUIRED_VERSION = "2.3.0"
-
-
-@importhook.on_import("flask.app")
-def on_flask_import(flask):
+@on_import("flask.app", "flask", version_requirement="2.3.0")
+def patch(m):
     """
-    Hook 'n wrap on `flask.app`. Flask class |-> App class |-> Scaffold class
-    @app.route |-> `add_url_rule` |-> self.view_functions. these get called via
-    full_dispatch_request, which we wrap. We also wrap __call__ to run our middleware.
+    patching module flask.appimport
+
+    - patches: Flask.__call__ (context parsing/initial stage)
+    - patches: Flask.full_dispatch_request
+    **Why?** full_dispatch_request gets called in the WSGI app. It gets called after all middleware. request_ctx is
+    available, so we can extract data from it. It returns a response, so we can send status codes and error messages.
+    (github src: https://github.com/pallets/flask/blob/bc143499cf1137a271a7cf75bdd3e16e43ede2f0/src/flask/app.py#L1529)
     """
-    if not is_package_compatible("flask", required_version=FLASK_REQUIRED_VERSION):
-        return flask
-    modified_flask = importhook.copy_module(flask)
-    former_fdr = copy.deepcopy(flask.Flask.full_dispatch_request)
-
-    def aikido_wrapper_fdr(*args, **kwargs):
-        return aik_full_dispatch_request(
-            *args, former_full_dispatch_request=former_fdr, **kwargs
-        )
-
-    # pylint:disable=no-member # Pylint has issues with the wrapping
-    setattr(modified_flask.Flask, "__call__", aikido___call__)
-    setattr(modified_flask.Flask, "full_dispatch_request", aikido_wrapper_fdr)
-    return modified_flask
+    patch_function(m, "Flask.__call__", _call)
+    patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_before)
+    patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_after)

From 1eb0e089f753a6f7d6dc088fc7f681f7b7b4411f Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 11:14:09 +0200
Subject: [PATCH 16/25] Cleanup flask.py imports

---
 aikido_zen/sources/flask.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/aikido_zen/sources/flask.py b/aikido_zen/sources/flask.py
index 6f354afdf..7086ee58d 100644
--- a/aikido_zen/sources/flask.py
+++ b/aikido_zen/sources/flask.py
@@ -1,19 +1,13 @@
-"""
-Flask source module, intercepts flask import and adds Aikido middleware
-"""
-
-import aikido_zen.importhook as importhook
 from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.helpers.logging import logger
 from aikido_zen.context import Context
-from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
 from aikido_zen.context import get_current_context
 import aikido_zen.sources.functions.request_handler as funcs
 from aikido_zen.sinks import (
     on_import,
     patch_function,
     after,
-    before_modify_return,
+    before_modify_return,   
     before,
 )
 

From 444a529a809443a8746c58bdc601b6684a31c950 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 11:34:34 +0200
Subject: [PATCH 17/25] 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 7086ee58d..43236c18f 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 b958f4caddb4d2350b6fccd4c09b95e9e4228519 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 12:02:23 +0200
Subject: [PATCH 18/25] remove try-finally for return, which was swallowing
 error

---
 aikido_zen/sinks/__init__.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/aikido_zen/sinks/__init__.py b/aikido_zen/sinks/__init__.py
index d3978a334..6c0386be5 100644
--- a/aikido_zen/sinks/__init__.py
+++ b/aikido_zen/sinks/__init__.py
@@ -50,8 +50,7 @@ def decorator(func, instance, args, kwargs):
             logger.debug(
                 "%s:%s wrapping-before error: %s", func.__module__, func.__name__, e
             )
-        finally:
-            return func(*args, **kwargs)  # Call the original function
+        return func(*args, **kwargs)  # Call the original function
 
     return decorator
 
@@ -91,8 +90,7 @@ async def decorator(func, instance, args, kwargs):
             logger.debug(
                 "%s:%s wrapping-before error: %s", func.__module__, func.__name__, e
             )
-        finally:
-            return await func(*args, **kwargs)  # Call the original function
+        return await func(*args, **kwargs)  # Call the original function
 
     return decorator
 

From f671c5fd4a5a2ae2a8d7e5645a42d4867545cdc7 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 13:05:28 +0200
Subject: [PATCH 19/25] Flask patch, get correct "environ" variable

---
 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 43236c18f..9bbc376ed 100644
--- a/aikido_zen/sources/flask.py
+++ b/aikido_zen/sources/flask.py
@@ -78,7 +78,7 @@ def extract_cookies_from_flask_request_and_save_data(req):
 
 @before
 def _call(func, instance, args, kwargs):
-    environ = get_argument(args, kwargs, 1, "environ")
+    environ = get_argument(args, kwargs, 0, "environ")
     context1 = Context(req=environ, source="flask")
     context1.set_as_current_context()
     funcs.request_handler(stage="init")

From 9e0c8b88ec97a2ff8329f075dc7d97c05051e824 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 15:10:11 +0200
Subject: [PATCH 20/25] Fix starlette bugs encountered after testing

---
 aikido_zen/sinks/__init__.py                  | 19 -------------------
 aikido_zen/sources/starlette/__init__.py      |  3 +++
 .../starlette/starlette_applications.py       | 10 +++++-----
 .../sources/starlette/starlette_routing.py    |  6 ++++--
 4 files changed, 12 insertions(+), 26 deletions(-)

diff --git a/aikido_zen/sinks/__init__.py b/aikido_zen/sinks/__init__.py
index 6c0386be5..14a1df841 100644
--- a/aikido_zen/sinks/__init__.py
+++ b/aikido_zen/sinks/__init__.py
@@ -76,25 +76,6 @@ def decorator(func, instance, args, kwargs):
     return decorator
 
 
-def before_async(wrapper):
-    """
-    Surrounds a patch with try-except and calls the original function at the end (async)
-    """
-
-    async def decorator(func, instance, args, kwargs):
-        try:
-            await wrapper(func, instance, args, kwargs)  # Call the patch
-        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 await func(*args, **kwargs)  # Call the original function
-
-    return decorator
-
-
 def after(wrapper):
     """
     Surrounds a patch with try-except, calls the original function and gives the return value to the patch
diff --git a/aikido_zen/sources/starlette/__init__.py b/aikido_zen/sources/starlette/__init__.py
index c00bdd55e..43ca74234 100644
--- a/aikido_zen/sources/starlette/__init__.py
+++ b/aikido_zen/sources/starlette/__init__.py
@@ -8,3 +8,6 @@
 - extract_data_from_request, which will extract the data from a request object safely,
     e.g. body, json, form. This also saves it inside the current context.
 """
+
+import aikido_zen.sources.starlette.starlette_applications
+import aikido_zen.sources.starlette.starlette_routing
diff --git a/aikido_zen/sources/starlette/starlette_applications.py b/aikido_zen/sources/starlette/starlette_applications.py
index 84dea4911..36d02899e 100644
--- a/aikido_zen/sources/starlette/starlette_applications.py
+++ b/aikido_zen/sources/starlette/starlette_applications.py
@@ -3,13 +3,13 @@
 from aikido_zen.context import Context
 from ..functions.request_handler import request_handler
 from ...helpers.get_argument import get_argument
-from ...sinks import on_import, patch_function, before_async
+from ...sinks import on_import, patch_function, before
 
 
-@before_async
-async def _call(func, instance, args, kwargs):
-    scope = get_argument(args, kwargs, 1, "scope")
-    if scope.get("type") != "http":
+@before
+def _call(func, instance, args, kwargs):
+    scope = get_argument(args, kwargs, 0, "scope")
+    if not hasattr(scope, "get") or scope.get("type") != "http":
         return
 
     new_context = Context(req=scope, source="starlette")
diff --git a/aikido_zen/sources/starlette/starlette_routing.py b/aikido_zen/sources/starlette/starlette_routing.py
index 48bebc218..d38f3d4b7 100644
--- a/aikido_zen/sources/starlette/starlette_routing.py
+++ b/aikido_zen/sources/starlette/starlette_routing.py
@@ -4,12 +4,14 @@
 from ...sinks import on_import, patch_function, before
 
 
-@before
 def _request_response(func, instance, args, kwargs):
     if kwargs and "func" in kwargs:
         kwargs["func"] = aik_route_func_wrapper(kwargs["func"])
     elif args and args[0]:
-        args[0] = aik_route_func_wrapper(args[0])
+        # Modify first element of a tuple, tuples are immutable
+        args = (aik_route_func_wrapper(args[0]),) + args[1:]
+
+    return func(*args, **kwargs)  # Call the original function
 
 
 @on_import("starlette.routing", "starlette")

From 32cdce1529b5569db40436765e063741dee13c6e Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 21:04:27 +0200
Subject: [PATCH 21/25] Update quart patches to use new system

---
 aikido_zen/sinks/__init__.py |  19 +++++
 aikido_zen/sources/quart.py  | 142 +++++++++++++++--------------------
 2 files changed, 79 insertions(+), 82 deletions(-)

diff --git a/aikido_zen/sinks/__init__.py b/aikido_zen/sinks/__init__.py
index 14a1df841..5d3278b66 100644
--- a/aikido_zen/sinks/__init__.py
+++ b/aikido_zen/sinks/__init__.py
@@ -55,6 +55,25 @@ def decorator(func, instance, args, kwargs):
     return decorator
 
 
+def before_async(wrapper):
+    """
+    Surrounds an async patch with try-except and calls the original asynchronous function at the end
+    """
+
+    async def decorator(func, instance, args, kwargs):
+        try:
+            await wrapper(func, instance, args, kwargs)  # Call the patch
+        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 await func(*args, **kwargs)  # Call the original function
+
+    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.
diff --git a/aikido_zen/sources/quart.py b/aikido_zen/sources/quart.py
index a2057362f..9e0fa6f63 100644
--- a/aikido_zen/sources/quart.py
+++ b/aikido_zen/sources/quart.py
@@ -1,70 +1,71 @@
-"""
-Quart source module, intercepts quart import and adds Aikido middleware
-"""
-
-import copy
-import aikido_zen.importhook as importhook
-from aikido_zen.helpers.logging import logger
 from aikido_zen.context import Context, get_current_context
-from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
 from .functions.request_handler import request_handler
+from ..helpers.get_argument import get_argument
+from ..sinks import on_import, patch_function, before, before_async
 
 
-async def aikido___call___wrapper(former_call, quart_app, scope, receive, send):
-    """Aikido's __call__ wrapper"""
-    # We don't want to install werkzeug :
-    # pylint: disable=import-outside-toplevel
-    try:
-        if scope["type"] != "http":
-            return await former_call(quart_app, scope, receive, send)
-        context1 = Context(req=scope, source="quart")
-        context1.set_as_current_context()
+@before
+def _call(func, instance, args, kwargs):
+    scope = get_argument(args, kwargs, 0, "scope")
+    if not scope or scope.get("type") != "http":
+        return
 
-        request_handler(stage="init")
-    except Exception as e:
-        logger.debug("Exception on aikido __call__ function : %s", e)
-    return await former_call(quart_app, scope, receive, send)
+    new_context = Context(req=scope, source="quart")
+    new_context.set_as_current_context()
+    request_handler(stage="init")
 
 
-async def handle_request_wrapper(former_handle_request, quart_app, req):
-    """
-    https://github.com/pallets/quart/blob/2fc6d4fa6e3df017e8eef1411ec80b5a6dce25a5/src/quart/app.py#L1400
-    Wraps the handle_request function
-    """
-    # At this stage no middleware is called yet, running pre_response is
-    # not what we need to do now, but we can store the body inside context :
-    try:
-        context = get_current_context()
-        if context:
-            form = await req.form
-            if req.is_json:
-                context.set_body(await req.get_json())
-            elif form:
-                context.set_body(form)
-            else:
-                data = await req.data
-                context.set_body(data.decode("utf-8"))
-            context.cookies = req.cookies.to_dict()
-            context.set_as_current_context()
-    except Exception as e:
-        logger.debug("Exception in handle_request : %s", e)
-
-    # Fetch response and run post_response handler :
+@before_async
+async def _handle_request_before(func, instance, args, kwargs):
+    context = get_current_context()
+    if not context:
+        return
+
+    request = get_argument(args, kwargs, 0, "request")
+    if not request:
+        return
+
+    form = await request.form
+    if request.is_json:
+        context.set_body(await request.get_json())
+    elif form:
+        context.set_body(form)
+    else:
+        data = await request.data
+        context.set_body(data.decode("utf-8"))
+    context.cookies = request.cookies.to_dict()
+    context.set_as_current_context()
+
+
+async def _handle_request_after(func, instance, args, kwargs):
     # pylint:disable=import-outside-toplevel # We don't want to install this by default
     from werkzeug.exceptions import HTTPException
 
     try:
-        response = await former_handle_request(quart_app, req)
-        status_code = response.status_code
-        request_handler(stage="post_response", status_code=status_code)
+        response = await func(*args, **kwargs)
+        if hasattr(response, "status_code"):
+            request_handler(stage="post_response", status_code=response.status_code)
         return response
     except HTTPException as e:
         request_handler(stage="post_response", status_code=e.code)
         raise e
 
 
+async def _asgi_app(func, instance, args, kwargs):
+    scope = get_argument(args, kwargs, 0, "scope")
+    if not scope or scope.get("type") != "http":
+        return await func(*args, **kwargs)
+    send = get_argument(args, kwargs, 2, "send")
+    if not send:
+        return await func(*args, **kwargs)
+
+    pre_response = request_handler(stage="pre_response")
+    if pre_response:
+        return await send_status_code_and_text(send, pre_response)
+    return await func(*args, **kwargs)
+
+
 async def send_status_code_and_text(send, pre_response):
-    """Sends a status code and text"""
     await send(
         {
             "type": "http.response.start",
@@ -81,38 +82,15 @@ async def send_status_code_and_text(send, pre_response):
     )
 
 
-@importhook.on_import("quart.app")
-def on_quart_import(quart):
+@on_import("quart.app", "quart")
+def patch(m):
     """
-    Hook 'n wrap on `quart.app`
-    Our goal is to wrap the __call__, handle_request, asgi_app functios of the "Quart" class
+    patching module quart.app
+    - patches Quart.__call__ (creates Context)
+    - patches Quart.handle_request (Stores body/cookies, checks status code)
+    - patches Quart.asgi_app (Pre-response: puts in messages when request is blocked)
     """
-    if not is_package_compatible("quart", required_version=ANY_VERSION):
-        return quart
-    modified_quart = importhook.copy_module(quart)
-
-    former_handle_request = copy.deepcopy(quart.Quart.handle_request)
-    former_asgi_app = copy.deepcopy(quart.Quart.asgi_app)
-    former_call = copy.deepcopy(quart.Quart.__call__)
-
-    async def aikido___call__(quart_app, scope, receive=None, send=None):
-        return await aikido___call___wrapper(
-            former_call, quart_app, scope, receive, send
-        )
-
-    async def aikido_handle_request(quart_app, request):
-        return await handle_request_wrapper(former_handle_request, quart_app, request)
-
-    async def aikido_asgi_app(quart_app, scope, receive=None, send=None):
-        if scope["type"] == "http":
-            # Run pre_response code :
-            pre_response = request_handler(stage="pre_response")
-            if pre_response:
-                return await send_status_code_and_text(send, pre_response)
-        return await former_asgi_app(quart_app, scope, receive, send)
-
-    # pylint:disable=no-member # Pylint has issues with the wrapping
-    setattr(modified_quart.Quart, "__call__", aikido___call__)
-    setattr(modified_quart.Quart, "handle_request", aikido_handle_request)
-    setattr(modified_quart.Quart, "asgi_app", aikido_asgi_app)
-    return modified_quart
+    patch_function(m, "Quart.__call__", _call)
+    patch_function(m, "Quart.handle_request", _handle_request_before)
+    patch_function(m, "Quart.handle_request", _handle_request_after)
+    patch_function(m, "Quart.asgi_app", _asgi_app)

From a0a17dfe64c8f95e6728f23be2080b9c14cbdb83 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Mon, 12 May 2025 21:04:44 +0200
Subject: [PATCH 22/25] Remove checks for who wrote response (not necessary)

---
 aikido_zen/sources/flask.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/aikido_zen/sources/flask.py b/aikido_zen/sources/flask.py
index 9bbc376ed..cd02b3098 100644
--- a/aikido_zen/sources/flask.py
+++ b/aikido_zen/sources/flask.py
@@ -26,15 +26,11 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
     if not pre_response:
         return None
     # This happens when a route is rate limited, a user blocked, etc...
-    res = Response(pre_response[0], status=pre_response[1], mimetype="text/plain")
-    res.is_aikido_response = True  # Use this in the @after
-    return res
+    return Response(pre_response[0], status=pre_response[1], mimetype="text/plain")
 
 
 @after
 def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
-    if hasattr(return_value, "is_aikido_response") and return_value.is_aikido_response:
-        return  # We wrote this response, so the status code is not reflective of the route
     if not hasattr(return_value, "status_code"):
         return
     funcs.request_handler(stage="post_response", status_code=return_value.status_code)

From b2f833e52da4f6fb33dff71bbffe6c766d4d63ad Mon Sep 17 00:00:00 2001
From: BitterPanda63 <bitterpanda@proton.me>
Date: Tue, 13 May 2025 12:35:29 +0200
Subject: [PATCH 23/25] lxml.py use new wrapping system

---
 aikido_zen/sources/xml_sources/lxml.py | 60 ++++++++++----------------
 1 file changed, 22 insertions(+), 38 deletions(-)

diff --git a/aikido_zen/sources/xml_sources/lxml.py b/aikido_zen/sources/xml_sources/lxml.py
index 0519ce7da..4558959cc 100644
--- a/aikido_zen/sources/xml_sources/lxml.py
+++ b/aikido_zen/sources/xml_sources/lxml.py
@@ -1,47 +1,31 @@
-"""
-Sink module for `xml`, python's built-in function
-"""
-
-import copy
-import aikido_zen.importhook as importhook
 from aikido_zen.helpers.extract_data_from_xml_body import (
     extract_data_from_xml_body,
 )
-from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION
-
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import on_import, after, patch_function
 
-@importhook.on_import("lxml.etree")
-def on_lxml_import(eltree):
-    """
-    Hook 'n wrap on `lxml.etree`.
-    - Wrap on fromstring() function
-    - Wrap on
-    Returns : Modified `lxml.etree` object
-    """
-    if not is_package_compatible("lxml", required_version=ANY_VERSION):
-        return eltree
-    modified_eltree = importhook.copy_module(eltree)
 
-    former_fromstring = copy.deepcopy(eltree.fromstring)
+@after
+def _fromstring(func, instance, args, kwargs, return_value):
+    text = get_argument(args, kwargs, 0, "text")
+    if text:
+        extract_data_from_xml_body(user_input=text, root_element=return_value)
 
-    def aikido_fromstring(text, *args, **kwargs):
-        res = former_fromstring(text, *args, **kwargs)
-        extract_data_from_xml_body(user_input=text, root_element=res)
-        return res
 
-    former_fromstringlist = copy.deepcopy(eltree.fromstringlist)
+@after
+def _fromstringlist(func, instance, args, kwargs, return_value):
+    strings = get_argument(args, kwargs, 0, "strings")
+    for text in strings:
+        extract_data_from_xml_body(user_input=text, root_element=return_value)
 
-    def aikido_fromstringlist(strings, *args, **kwargs):
-        res = former_fromstringlist(strings, *args, **kwargs)
-        for string in strings:
-            extract_data_from_xml_body(user_input=string, root_element=res)
-        return res
 
-    # pylint: disable=no-member
-    setattr(eltree, "fromstring", aikido_fromstring)
-    setattr(modified_eltree, "fromstring", aikido_fromstring)
-
-    # pylint: disable=no-member
-    setattr(eltree, "fromstringlist", aikido_fromstringlist)
-    setattr(modified_eltree, "fromstringlist", aikido_fromstringlist)
-    return modified_eltree
+@on_import("lxml.etree", "lxml")
+def patch(m):
+    """
+    patching module lxml.etree
+    - patches function fromstring(text, ...)
+    - patches function fromstringlist(strings, ...)
+    (github src: https://github.com/lxml/lxml/blob/fe271a4b5a32e6e54d10983683f2f32b0647209a/src/lxml/etree.pyx#L3411)
+    """
+    patch_function(m, "fromstring", _fromstring)
+    patch_function(m, "fromstringlist", _fromstringlist)

From 3c427fb77c36189a84c3021a06eca759027edf54 Mon Sep 17 00:00:00 2001
From: BitterPanda63 <bitterpanda@proton.me>
Date: Tue, 13 May 2025 12:38:43 +0200
Subject: [PATCH 24/25] Update xml.py to use new wrapping system

---
 aikido_zen/sources/xml_sources/xml.py | 56 ++++++---------------------
 1 file changed, 12 insertions(+), 44 deletions(-)

diff --git a/aikido_zen/sources/xml_sources/xml.py b/aikido_zen/sources/xml_sources/xml.py
index 2062f04e2..373f1f960 100644
--- a/aikido_zen/sources/xml_sources/xml.py
+++ b/aikido_zen/sources/xml_sources/xml.py
@@ -1,51 +1,19 @@
-"""
-Sink module for `xml`, python's built-in function
-"""
-
-import copy
-import aikido_zen.importhook as importhook
-from aikido_zen.helpers.logging import logger
 from aikido_zen.helpers.extract_data_from_xml_body import (
     extract_data_from_xml_body,
 )
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import on_import, patch_function, after
 
 
-@importhook.on_import("xml.etree.ElementTree")
-def on_xml_import(eltree):
-    """
-    Hook 'n wrap on `xml.etree.ElementTree`, python's built-in xml lib
-    Our goal is to create a new and mutable aikido parser class
-    Returns : Modified ElementTree object
-    """
-    modified_eltree = importhook.copy_module(eltree)
-    copy_xml_parser = copy.deepcopy(eltree.XMLParser)
-
-    class MutableAikidoXMLParser:
-        """Mutable connection class used to instrument `xml` by Zen"""
-
-        def __init__(self, *args, **kwargs):
-            self._former_xml_parser = copy_xml_parser(*args, **kwargs)
-            self._feed_func_copy = copy.deepcopy(self._former_xml_parser.feed)
-
-        def __getattr__(self, name):
-            if name != "feed":
-                return getattr(self._former_xml_parser, name)
-
-            # Return aa function dynamically
-            def feed(data):
-                former_feed_result = self._feed_func_copy(data)
-
-                # Fetch the data, this should just return an internal attribute
-                # and not close a stream or something that is noticable by the end-user
-                parsed_xml = self.target.close()
-                extract_data_from_xml_body(user_input=data, root_element=parsed_xml)
-
-                return former_feed_result
-
-            return feed
+@after
+def _feed(func, instance, args, kwargs, return_value):
+    # Fetches XML data, this should just return an internal attribute
+    # and not close a stream or something that is noticable by the end-user
+    parsed_xml = instance.target.close()
+    data = get_argument(args, kwargs, 0, "data")
+    extract_data_from_xml_body(user_input=data, root_element=parsed_xml)
 
-    # pylint: disable=no-member
-    setattr(eltree, "XMLParser", MutableAikidoXMLParser)
-    setattr(modified_eltree, "XMLParser", MutableAikidoXMLParser)
 
-    return modified_eltree
+@on_import("xml.etree.ElementTree")
+def patch(m):
+    patch_function(m, "XMLParser.feed", _feed)

From 8a1484282f27ff8e95affba96101713e31ec4263 Mon Sep 17 00:00:00 2001
From: BitterPanda <bitterpanda@proton.me>
Date: Fri, 16 May 2025 13:11:45 +0200
Subject: [PATCH 25/25] Zen update xml.py to also patch
 fromstring/fromstringlist

---
 aikido_zen/sources/xml_sources/xml.py | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/aikido_zen/sources/xml_sources/xml.py b/aikido_zen/sources/xml_sources/xml.py
index 373f1f960..6f26d5ec9 100644
--- a/aikido_zen/sources/xml_sources/xml.py
+++ b/aikido_zen/sources/xml_sources/xml.py
@@ -6,14 +6,25 @@
 
 
 @after
-def _feed(func, instance, args, kwargs, return_value):
-    # Fetches XML data, this should just return an internal attribute
-    # and not close a stream or something that is noticable by the end-user
-    parsed_xml = instance.target.close()
-    data = get_argument(args, kwargs, 0, "data")
-    extract_data_from_xml_body(user_input=data, root_element=parsed_xml)
+def _fromstring(func, instance, args, kwargs, return_value):
+    text = get_argument(args, kwargs, 0, "text")
+    extract_data_from_xml_body(user_input=text, root_element=return_value)
+
+
+@after
+def _fromstringlist(func, instance, args, kwargs, return_value):
+    strings = get_argument(args, kwargs, 0, "sequence")
+    for text in strings:
+        extract_data_from_xml_body(user_input=text, root_element=return_value)
 
 
 @on_import("xml.etree.ElementTree")
 def patch(m):
-    patch_function(m, "XMLParser.feed", _feed)
+    """
+    patching module xml.etree.ElementTree
+    - patches function fromstring(text, ...)
+    - patches function fromstringlist(sequence, ...)
+    (src: https://github.com/python/cpython/blob/bc1a6ecfab02075acea79f8460a2dce70c61b2fd/Lib/xml/etree/ElementTree.py#L1370)
+    """
+    patch_function(m, "fromstring", _fromstring)
+    patch_function(m, "fromstringlist", _fromstringlist)