Update quart patches to use new system

Author: bitterpanda63

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.

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)