add uvicorn context-creation support when run with gUnicorn

.github/workflows/end2end.yml

@@ -29,6 +29,7 @@ jobs:
           - { name: django-mysql, testfile: end2end/django_mysql_test.py }
           - { name: django-mysql-gunicorn, testfile: end2end/django_mysql_gunicorn_test.py }
           - { name: django-postgres-gunicorn, testfile: end2end/django_postgres_gunicorn_test.py }
+          - { name: django-asgi-uvicorn, testfile: end2end/django_asgi_uvicorn_test.py }
           - { name: flask-mongo, testfile: end2end/flask_mongo_test.py }
           - { name: flask-mysql, testfile: end2end/flask_mysql_test.py }
           - { name: flask-mysql-uwsgi, testfile: end2end/flask_mysql_uwsgi_test.py }

README.md

@@ -43,6 +43,7 @@ Zen for Python 3 is compatible with:
 * ✅ [Quart](docs/quart.md)
 * ✅ [Starlette](docs/starlette.md)
 * ✅ [FastAPI](docs/fastapi.md)
+* 🚧 [Django ASGI](docs/django_asgi.md) (*limited support*)
 
 ### Database drivers
 * ✅ [`mysqlclient`](https://pypi.org/project/mysqlclient/) ^1.5

aikido_zen/__init__.py

@@ -56,15 +56,21 @@ def protect(mode="daemon", token=""):
     # pylint: disable=unused-import
     import aikido_zen.sinks.builtins_import
 
-    # Import sources
+    # 1. Import sources
+    ## 1a. Import frameworks
     import aikido_zen.sources.django
     import aikido_zen.sources.flask
     import aikido_zen.sources.quart
     import aikido_zen.sources.starlette
+
+    ## 1b. Import servers
+    import aikido_zen.sources.servers.uvicorn
+
+    ## 1c. Import xml sources
     import aikido_zen.sources.xml_sources.xml
     import aikido_zen.sources.xml_sources.lxml
 
-    # Import DB Sinks
+    # 2. Import database sinks
     import aikido_zen.sinks.pymysql
     import aikido_zen.sinks.mysqlclient
     import aikido_zen.sinks.pymongo
@@ -73,19 +79,22 @@ def protect(mode="daemon", token=""):
     import aikido_zen.sinks.asyncpg
     import aikido_zen.sinks.clickhouse_driver
 
+    # 3. Import fs sinks
     import aikido_zen.sinks.builtins
     import aikido_zen.sinks.os
     import aikido_zen.sinks.pathlib
     import aikido_zen.sinks.shutil
     import aikido_zen.sinks.io
+
+    # 4. Import sinks related to http requests
     import aikido_zen.sinks.http_client
     import aikido_zen.sinks.socket
 
-    # Import shell sinks
+    # 5. Import shell sinks
     import aikido_zen.sinks.os_system
     import aikido_zen.sinks.subprocess
 
-    # Import AI sinks
+    # 6. Import AI sinks
     import aikido_zen.sinks.openai
     import aikido_zen.sinks.anthropic
     import aikido_zen.sinks.mistralai

aikido_zen/context/__init__.py

@@ -21,7 +21,7 @@
 current_context = contextvars.ContextVar("current_context", default=None)
 
 WSGI_SOURCES = ["django", "flask"]
-ASGI_SOURCES = ["quart", "django_async", "starlette"]
+ASGI_SOURCES = ["quart", "django_async", "starlette", "uvicorn"]
 
 
 def get_current_context():

aikido_zen/sinks/psycopg.py

@@ -11,8 +11,7 @@
 @before
 def _copy(func, instance, args, kwargs):
     statement = get_argument(args, kwargs, 0, "statement")
-
-    op = "psycopg.Cursor.copy"
+    op = f"psycopg.{instance.__class__.__name__}.copy"
     register_call(op, "sql_op")
 
     vulns.run_vulnerability_scan(
@@ -23,7 +22,7 @@ def _copy(func, instance, args, kwargs):
 @before
 def _execute(func, instance, args, kwargs):
     query = get_argument(args, kwargs, 0, "query")
-    op = f"psycopg.Cursor.{func.__name__}"
+    op = f"psycopg.{instance.__class__.__name__}.{func.__name__}"
     vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=(query, "postgres"))
 
 
@@ -38,3 +37,13 @@ def patch(m):
     patch_function(m, "Cursor.copy", _copy)
     patch_function(m, "Cursor.execute", _execute)
     patch_function(m, "Cursor.executemany", _execute)
+
+
+@on_import("psycopg.cursor_async", "psycopg", version_requirement="3.1.0")
+def patch_async(m):
+    """
+    patching module psycopg.cursor_async (similar to normal patch)
+    """
+    patch_function(m, "AsyncCursor.copy", _copy)
+    patch_function(m, "AsyncCursor.execute", _execute)
+    patch_function(m, "AsyncCursor.executemany", _execute)

aikido_zen/sources/django/run_init_stage.py

@@ -15,11 +15,7 @@ def run_init_stage(request):
     # In a separate try-catch we set the context :
     try:
         context = None
-        if (
-            hasattr(request, "scope") and request.scope is not None
-        ):  # This request is an ASGI request
-            context = Context(req=request.scope, body=body, source="django_async")
-        elif hasattr(request, "META") and request.META is not None:  # WSGI request
+        if hasattr(request, "META") and request.META is not None:
             context = Context(req=request.META, body=body, source="django")
         else:
             return

aikido_zen/sources/django/run_init_stage_test.py

@@ -19,21 +19,6 @@
     "CONTENT_TYPE": "application/json",
     "REMOTE_ADDR": "198.51.100.23",
 }
-asgi_scope = {
-    "method": "PUT",
-    "headers": [
-        (b"COOKIE", b"a=b; c=d"),
-        (b"header1_test-2", b"testValue2198&"),
-        (b"USER-AGENT", b"testUserAgent"),
-        (b"USER-AGENT", b"testUserAgent2"),
-    ],
-    "query_string": b"a=b&b=d",
-    "client": ["1.1.1.1"],
-    "server": ["192.168.0.1", 443],
-    "scheme": "https",
-    "root_path": "192.168.0.1",
-    "path": "192.168.0.1/a/b/c/d",
-}
 
 
 @pytest.fixture
@@ -195,11 +180,3 @@ def test_uses_wsgi(mock_request):
     context: Context = get_current_context()
     assert "/hello" == context.route
 
-
-def test_uses_asgi_prio(mock_request):
-    mock_request.scope = asgi_scope
-    run_init_stage(mock_request)
-    # Assertions
-    context: Context = get_current_context()
-    assert "/a/b/c/d" == context.route
-    assert "testUserAgent2" == context.get_user_agent()

aikido_zen/sources/functions/asgi_middleware.py

@@ -0,0 +1,93 @@
+import inspect
+from aikido_zen.context import Context
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import before_async, patch_function
+from aikido_zen.sources.functions.request_handler import request_handler, post_response
+from aikido_zen.thread.thread_cache import get_cache
+
+
+class InternalASGIMiddleware:
+    def __init__(self, app, source: str):
+        self.client_app = app
+        self.source = source
+
+    async def __call__(self, scope, receive, send):
+        if not scope or scope.get("type") != "http":
+            # Zen checks requests coming into HTTP(S) server, ignore other requests (like ws)
+            return await self.continue_app(scope, receive, send)
+
+        context = Context(req=scope, source=self.source)
+
+        process_cache = get_cache()
+        if process_cache and process_cache.is_bypassed_ip(context.remote_address):
+            # IP address is bypassed, for simplicity we do not set a context,
+            # and we do not do any further handling of the request.
+            return await self.continue_app(scope, receive, send)
+
+        context.set_as_current_context()
+        if process_cache:
+            # Since this SHOULD be the highest level of the apps we wrap, this is the safest place
+            # to increment total hits.
+            process_cache.stats.increment_total_hits()
+
+        intercept_response = request_handler(stage="pre_response")
+        if intercept_response:
+            # The request has already been blocked (e.g. IP is on blocklist)
+            return await send_status_code_and_text(send, intercept_response)
+
+        return await self.run_with_intercepts(scope, receive, send)
+
+    async def run_with_intercepts(self, scope, receive, send):
+        # We use a skeleton class so we can use patch_function (and the logic already defined in @before_async)
+        class InterceptorSkeleton:
+            @staticmethod
+            async def send(*args, **kwargs):
+                return await send(*args, **kwargs)
+
+        patch_function(InterceptorSkeleton, "send", send_interceptor)
+
+        return await self.continue_app(scope, receive, InterceptorSkeleton.send)
+
+    async def continue_app(self, scope, receive, send):
+        client_app_parameters = len(inspect.signature(self.client_app).parameters)
+        if client_app_parameters == 2:
+            # This is possible if the app is still using ASGI v2.0
+            # See https://asgi.readthedocs.io/en/latest/specs/main.html#legacy-applications
+            # client_app = coroutine application_instance(receive, send)
+            await self.client_app(receive, send)
+        else:
+            # client_app = coroutine application(scope, receive, send)
+            await self.client_app(scope, receive, send)
+
+
+async def send_status_code_and_text(send, pre_response):
+    await send(
+        {
+            "type": "http.response.start",
+            "status": pre_response[1],
+            "headers": [(b"content-type", b"text/plain")],
+        }
+    )
+    await send(
+        {
+            "type": "http.response.body",
+            "body": pre_response[0].encode("utf-8"),
+            "more_body": False,
+        }
+    )
+
+
+@before_async
+async def send_interceptor(func, instance, args, kwargs):
+    # There is no name for the send() comment in the standard, it really depends (quart uses message)
+    event = get_argument(args, kwargs, 0, name="message")
+
+    # https://asgi.readthedocs.io/en/latest/specs/www.html#response-start-send-event
+    if not event or "http.response.start" not in event.get("type", ""):
+        # If the event is not of type http.response.start it won't contain the status code.
+        # And this event is required before sending over a body (so even 200 status codes are intercepted).
+        return
+
+    if "status" in event:
+        # Handle post response logic (attack waves, route reporting, ...)
+        post_response(status_code=int(event.get("status")))