add uvicorn context-creation support when run with gUnicorn

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/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")))

aikido_zen/sources/quart.py

@@ -1,4 +1,5 @@
 from aikido_zen.context import Context, get_current_context
+from .functions.asgi_middleware import send_status_code_and_text
 from .functions.request_handler import request_handler
 from ..helpers.get_argument import get_argument
 from ..sinks import on_import, patch_function, before, before_async
@@ -65,23 +66,6 @@ async def _asgi_app(func, instance, args, kwargs):
     return await func(*args, **kwargs)
 
 
-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,
-        }
-    )
-
-
 @on_import("quart.app", "quart")
 def patch(m):
     """

aikido_zen/sources/servers/uvicorn.py

@@ -0,0 +1,21 @@
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import before, on_import, patch_function
+from aikido_zen.helpers.logging import logger
+from aikido_zen.sources.functions.asgi_middleware import InternalASGIMiddleware
+
+
+@before
+def _init(func, instance, args, kwargs):
+    config = get_argument(args, kwargs, 0, "config")
+    if not config.app:
+        return
+    logger.debug("Server detected as: Uvicorn, Wrapping ASGI app with Zen.")
+    config.app = InternalASGIMiddleware(config.app, "uvicorn")
+
+
+@on_import("uvicorn.server", "uvicorn")
+def patch(m):
+    """
+    We patch uvicorn.server, so that we can modify the app to go through our ASGI middleware first
+    """
+    patch_function(m, "Server.__init__", _init)

docs/django_asgi.md

@@ -0,0 +1,53 @@
+# Django ASGI
+
+## Support
+Currently Django ASGI is **only supported with gUnicorn** and a uvicorn worker.
+
+
+## Installation for gUnicorn with a Uvicorn Worker
+1. Install `aikido_zen` package with pip :
+```sh
+pip install aikido_zen
+```
+
+2. Use the following template for your `gunicorn_config.py` file :
+```python
+import aikido_zen.decorators.gunicorn as aik
+
+@aik.post_fork
+def post_fork(server, worker):
+    # If you already have a config file, replace pass with your own code.
+    pass
+```
+And make sure to include this config when starting gunicorn by adding the `-c gunicorn_config.py` flag.
+```sh
+gunicorn -c gunicorn_config.py --workers ...
+```
+
+3. Setting your environment variables :
+Make sure to set your token in order to communicate with Aikido's servers
+```env
+AIKIDO_TOKEN="AIK_RUNTIME_YOUR_TOKEN_HERE"
+```
+
+## Blocking mode
+
+By default, the firewall will run in non-blocking mode. When it detects an attack, the attack will be reported to Aikido and continue executing the call.
+
+You can enable blocking mode by setting the environment variable `AIKIDO_BLOCK` to `true`:
+
+```sh
+AIKIDO_BLOCK=true
+```
+
+It's recommended to enable this on your staging environment for a considerable amount of time before enabling it on your production environment (e.g. one week).
+
+## Debug mode
+
+If you need to debug the firewall, you can run your code with the environment variable `AIKIDO_DEBUG` set to `true`:
+
+```sh
+AIKIDO_DEBUG=true
+```
+
+This will output debug information to the console (e.g. no token was found, unsupported packages, extra information, ...).

sample-apps/README.md

@@ -50,3 +50,7 @@ Overview :
 - `fastapi-postgres-uvicorn/` is a Fastapi app using Postgres with uvicorn, runs async.
   - it runs **multi-threaded & async**
   - Runs on 8112. Without Aikido runs on 8113
+- `django-asgi-uvicorn/` is a Django ASGI app using Postgres with Gunicorn + Uvicorn workers.
+  - it runs 4 processes, called workers, (**multi-process**) which handles requests using ASGI protocol
+  - This application uses **ASGI** (Asynchronous Server Gateway Interface)
+  - Runs on 8114. Without Aikido runs on 8115

sample-apps/django-asgi-uvicorn/Makefile

@@ -0,0 +1,21 @@
+include ../common.mk
+
+PORT = 8114
+PORT_DISABLED = 8115
+
+.PHONY: run
+run: install
+	@echo "Running sample app django-asgi-uvicorn with Zen on port $(PORT)"
+	poetry run python manage.py migrate || true
+
+	OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES \
+	$(AIKIDO_ENV_COMMON) \
+	poetry run gunicorn -c gunicorn_config.py --workers 4 --log-level debug --access-logfile '-' --error-logfile '-' --bind 0.0.0.0:$(PORT) sample-django-asgi-uvicorn-app.asgi:application
+
+.PHONY: runZenDisabled
+runZenDisabled: install
+	@echo "Running sample app django-asgi-uvicorn without Zen on port $(PORT_DISABLED)"
+	poetry run python manage.py migrate || true
+	OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES \
+	$(AIKIDO_ENV_DISABLED) \
+	poetry run gunicorn -c gunicorn_config.py --workers 4 --log-level debug --access-logfile '-' --error-logfile '-' --bind 0.0.0.0:$(PORT_DISABLED) sample-django-asgi-uvicorn-app.asgi:application

sample-apps/django-asgi-uvicorn/README.md

@@ -0,0 +1,15 @@
+# Sample Django ASGI App with Gunicorn + Uvicorn
+This is a Django ASGI application running with Gunicorn using the UvicornWorker.
+
+## Starting the app
+
+Run:
+```bash
+make run # Runs app with zen
+make runZenDisabled # Runs app with zen disabled.
+```
+
+## URLS
+- Homepage: `http://localhost:8114/app`
+- Create a dog: `http://localhost:8114/app/create/<dog_name>`
+- SQL injection attack example: `Malicious dog", "Injected wrong boss name"); -- `

sample-apps/django-asgi-uvicorn/gunicorn_config.py

@@ -0,0 +1,7 @@
+import aikido_zen.decorators.gunicorn as aik
+
+@aik.post_fork
+def post_fork(server, worker): pass
+
+# Use a UvicornWorker
+worker_class = 'uvicorn.workers.UvicornWorker'

sample-apps/django-asgi-uvicorn/manage.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+"""Django's command-line utility for administrative tasks."""
+import os
+import sys
+
+
+def main():
+    """Run administrative tasks."""
+    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'sample-django-asgi-uvicorn-app.settings')
+    try:
+        from django.core.management import execute_from_command_line
+    except ImportError as exc:
+        raise ImportError(
+            "Couldn't import Django. Are you sure it's installed and "
+            "available on your PYTHONPATH environment variable? Did you "
+            "forget to activate a virtual environment?"
+        ) from exc
+    execute_from_command_line(sys.argv)
+
+
+if __name__ == '__main__':
+    main()