Add support for Django ASGI

Author: tomaisthorpe

aikido_zen/sources/django/__init__.py

@@ -2,9 +2,12 @@
 
 from ..functions.request_handler import request_handler
 from .run_init_stage import run_init_stage
-from .pre_response_middleware import pre_response_middleware
+from .pre_response_middleware import (
+    pre_response_middleware,
+    pre_response_middleware_async,
+)
 from ...helpers.get_argument import get_argument
-from ...sinks import on_import, patch_function, before, after
+from ...sinks import on_import, patch_function, before, after, before_async, after_async
 
 
 @before
@@ -25,13 +28,32 @@ def _get_response_after(func, instance, args, kwargs, return_value):
         request_handler(stage="post_response", status_code=return_value.status_code)
 
 
+@before_async
+async def _get_response_async_before(func, instance, args, kwargs):
+    request = get_argument(args, kwargs, 0, "request")
+
+    run_init_stage(request)
+
+    if pre_response_middleware_async not in getattr(instance, "_view_middleware"):
+        # pylint:disable=protected-access
+        instance._view_middleware += [pre_response_middleware_async]
+
+
+@after_async
+async def _get_response_async_after(func, instance, args, kwargs, return_value):
+    if hasattr(return_value, "status_code"):
+        request_handler(stage="post_response", status_code=return_value.status_code)
+
+
 @on_import("django.core.handlers.base", "django")
 def patch(m):
     """
-    Patch for _get_response (Synchronous/WSGI)
+    Patch for _get_response (WSGI) and _get_response_async (ASGI)
     - before: Parse body, create context & add middleware to run before a response
-    - after: Check respone code to see if route should be analyzed
+    - after: Check response 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)
+    patch_function(m, "BaseHandler._get_response_async", _get_response_async_before)
+    patch_function(m, "BaseHandler._get_response_async", _get_response_async_after)

aikido_zen/sources/django/pre_response_middleware.py

@@ -1,4 +1,4 @@
-"""Exports pre_response_middleware function"""
+"""Exports pre_response_middleware and pre_response_middleware_async functions"""
 
 from ..functions.request_handler import request_handler
 
@@ -12,3 +12,14 @@ def pre_response_middleware(request, *args, **kwargs):
 
         return HttpResponse(response[0], status=response[1])
     return None
+
+
+async def pre_response_middleware_async(request, *args, **kwargs):
+    """Async variant of pre_response_middleware for Django ASGI"""
+    response = request_handler(stage="pre_response")
+    if response:
+        # pylint:disable=import-outside-toplevel # We don't want to install this by default
+        from django.http import HttpResponse
+
+        return HttpResponse(response[0], status=response[1])
+    return None

end2end/django_asgi_uvicorn_test.py

@@ -0,0 +1,61 @@
+import pytest
+import requests
+import time
+from .server.check_events_from_mock import fetch_events_from_mock, validate_started_event, filter_on_event_type, \
+    clear_events_from_mock
+
+# e2e tests for django-asgi-uvicorn sample app
+post_url_fw = "http://localhost:8114/create"
+post_url_nofw = "http://localhost:8115/create"
+
+
+def test_firewall_started_okay():
+    events = fetch_events_from_mock("http://localhost:5000")
+    started_events = filter_on_event_type(events, "started")
+    assert len(started_events) == 1
+    validate_started_event(started_events[0], ["gunicorn", "django", "psycopg"])
+
+
+def test_safe_response_with_firewall():
+    dog_name = "Bobby Tables"
+    res = requests.post(post_url_fw, data={"dog_name": dog_name})
+    assert res.status_code == 200
+
+
+def test_safe_response_without_firewall():
+    dog_name = "Bobby Tables"
+    res = requests.post(post_url_nofw, data={"dog_name": dog_name})
+    assert res.status_code == 200
+
+
+def test_dangerous_response_with_firewall():
+    clear_events_from_mock("http://localhost:5000")
+    dog_name = "Dangerous bobby', TRUE); -- "
+    res = requests.post(post_url_fw, data={"dog_name": dog_name})
+    assert res.status_code == 500
+
+    time.sleep(5)  # Wait for attack to be reported
+    events = fetch_events_from_mock("http://localhost:5000")
+    attacks = filter_on_event_type(events, "detected_attack")
+
+    assert len(attacks) == 1
+    del attacks[0]["attack"]["stack"]
+    assert attacks[0]["attack"] == {
+        "blocked": True,
+        "kind": "sql_injection",
+        "metadata": {
+            "dialect": "postgres",
+            "sql": "INSERT INTO sample_app_dogs (dog_name, is_admin) VALUES ('Dangerous bobby', TRUE); -- ', FALSE)",
+        },
+        "operation": "psycopg.AsyncCursor.execute",
+        "pathToPayload": ".dog_name.[0]",
+        "payload": "\"Dangerous bobby', TRUE); -- \"",
+        "source": "body",
+        "user": None,
+    }
+
+
+def test_dangerous_response_without_firewall():
+    dog_name = "Dangerous bobby', TRUE); -- "
+    res = requests.post(post_url_nofw, data={"dog_name": dog_name})
+    assert res.status_code == 200