New Wrapping PR 5: Update all sources to new wrapping system

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

aikido_zen/sinks/__init__.py

@@ -50,7 +50,46 @@ def decorator(func, instance, args, kwargs):
             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
+    """
+
+    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.
+    """
+
+    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

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)
 

aikido_zen/sources/django/__init__.py

@@ -1,42 +1,36 @@
 """`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
+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", "django")
+def patch(m):
+    """
+    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)
+    patch_function(m, "BaseHandler._get_response", _get_response_after)

aikido_zen/sources/flask.py

@@ -1,46 +1,39 @@
-"""
-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)
-    return res
+    if not pre_response:
+        return None
+    # This happens when a route is rate limited, a user blocked, etc...
+    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 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):
@@ -79,41 +72,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
-
-
-FLASK_REQUIRED_VERSION = "2.3.0"
+@before
+def _call(func, instance, args, kwargs):
+    environ = get_argument(args, kwargs, 0, "environ")
+    context1 = Context(req=environ, source="flask")
+    context1.set_as_current_context()
+    funcs.request_handler(stage="init")
 
 
-@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)