Remove parse_raw_body and let django do the parsing

aikido_firewall/context/__init__.py

@@ -12,7 +12,6 @@
 from .parse_cookies import parse_cookies
 from .extract_wsgi_headers import extract_wsgi_headers
 from .build_url_from_wsgi import build_url_from_wsgi
-from .parse_raw_body import parse_raw_body
 
 UINPUT_SOURCES = ["body", "cookies", "query", "headers", "xml"]
 local = threading.local()
@@ -32,7 +31,7 @@ class Context:
     for vulnerability detection
     """
 
-    def __init__(self, context_obj=None, req=None, raw_body=None, source=None):
+    def __init__(self, context_obj=None, body=None, req=None, source=None):
         if context_obj:
             logger.debug("Creating Context instance based on dict object.")
             self.__dict__.update(context_obj)
@@ -48,7 +47,7 @@ def __init__(self, context_obj=None, req=None, raw_body=None, source=None):
         self.url = build_url_from_wsgi(req)
         self.query = parse_qs(req["QUERY_STRING"])
         content_type = req.get("CONTENT_TYPE", None)
-        self.body = parse_raw_body(raw_body, content_type)
+        self.body = body
         self.route = build_route_from_url(self.url)
         self.subdomains = get_subdomains_from_url(self.url)
         self.user = None

aikido_firewall/context/init_test.py

@@ -23,8 +23,7 @@ def test_wsgi_context_1():
         "CONTENT_TYPE": "application/x-www-form-urlencoded",
         "REMOTE_ADDR": "198.51.100.23",
     }
-    wsgi_raw_body = "dog_name=Doggo 1&test=Test 1"
-    context = Context(req=wsgi_request, raw_body=wsgi_raw_body, source="django")
+    context = Context(req=wsgi_request, body=123, source="django")
     assert context.__dict__ == {
         "source": "django",
         "method": "POST",
@@ -37,7 +36,7 @@ def test_wsgi_context_1():
         "cookies": {"sessionId": "abc123xyz456"},
         "url": "https://example.com/hello",
         "query": {"user": ["JohnDoe"], "age": ["30", "35"]},
-        "body": {"dog_name": ["Doggo 1"], "test": ["Test 1"]},
+        "body": 123,
         "route": "/hello",
         "subdomains": [],
         "user": None,
@@ -61,8 +60,7 @@ def test_wsgi_context_2():
         "CONTENT_TYPE": "application/json",
         "REMOTE_ADDR": "198.51.100.23",
     }
-    wsgi_raw_body = '{"a": 23, "b": 45, "Hello": [1, 2, 3]}'
-    context = Context(req=wsgi_request, raw_body=wsgi_raw_body, source="flask")
+    context = Context(req=wsgi_request, body={"test": True}, source="flask")
     assert context.__dict__ == {
         "source": "flask",
         "method": "GET",
@@ -75,7 +73,7 @@ def test_wsgi_context_2():
         "cookies": {"sessionId": "abc123xyz456"},
         "url": "http://localhost:8080/hello",
         "query": {"user": ["JohnDoe"], "age": ["30", "35"]},
-        "body": {"a": 23, "b": 45, "Hello": [1, 2, 3]},
+        "body": {"test": True},
         "route": "/hello",
         "subdomains": [],
         "user": None,
@@ -100,8 +98,7 @@ def test_set_as_current_context(mocker):
         "CONTENT_TYPE": "application/json",
         "REMOTE_ADDR": "198.51.100.23",
     }
-    wsgi_raw_body = '{"a": 23, "b": 45, "Hello": [1, 2, 3]}'
-    context = Context(req=wsgi_request, raw_body=wsgi_raw_body, source="flask")
+    context = Context(req=wsgi_request, body=12, source="flask")
     context.set_as_current_context()
     assert get_current_context() == context
 
@@ -121,8 +118,7 @@ def test_get_current_context_with_context(mocker):
         "CONTENT_TYPE": "application/json",
         "REMOTE_ADDR": "198.51.100.23",
     }
-    wsgi_raw_body = '{"a": 23, "b": 45, "Hello": [1, 2, 3]}'
-    context = Context(req=wsgi_request, raw_body=wsgi_raw_body, source="flask")
+    context = Context(req=wsgi_request, body=456, source="flask")
     context.set_as_current_context()
     assert get_current_context() == context
 
@@ -141,15 +137,14 @@ def test_context_is_picklable(mocker):
         "CONTENT_TYPE": "application/json",
         "REMOTE_ADDR": "198.51.100.23",
     }
-    wsgi_raw_body = '{"a": 23, "b": 45, "Hello": [1, 2, 3]}'
-    context = Context(req=wsgi_request, raw_body=wsgi_raw_body, source="flask")
+    context = Context(req=wsgi_request, body=123, source="flask")
     pickled_obj = pickle.dumps(context)
     unpickled_obj = pickle.loads(pickled_obj)
     assert unpickled_obj.source == "flask"
     assert unpickled_obj.method == "GET"
     assert unpickled_obj.remote_address == "198.51.100.23"
     assert unpickled_obj.url == "http://localhost:8080/hello"
-    assert unpickled_obj.body == {"a": 23, "b": 45, "Hello": [1, 2, 3]}
+    assert unpickled_obj.body == 123
     assert unpickled_obj.headers == {
         "HEADER_1": "header 1 value",
         "HEADER_2": "Header 2 value",

aikido_firewall/sources/django.py

@@ -17,9 +17,18 @@ def gen_aikido_middleware_function(former__middleware_chain):
     """
 
     def aikido_middleware_function(request):
-        context = Context(
-            req=request.META, raw_body=request.body.decode("utf-8"), source="django"
-        )
+        # Get a parsed body from Django :
+        body = request.POST.dict()
+        if len(body) == 0 and request.content_type == "application/json":
+            try:
+                body = json.loads(request.body)
+            except Exception:
+                pass
+        if len(body) == 0:
+            # E.g. XML Data
+            body = request.body
+
+        context = Context(req=request.META, body=body, source="django")
         context.set_as_current_context()
         request_handler(stage="init")
 

aikido_firewall/sources/flask.py

@@ -56,7 +56,7 @@ def aikido___call__(flask_app, environ, start_response):
     # We don't want to install werkzeug :
     # pylint: disable=import-outside-toplevel
     try:
-        context1 = Context(req=environ, raw_body={}, source="flask")
+        context1 = Context(req=environ, source="flask")
         context1.set_as_current_context()
         request_handler(stage="init")
     except Exception as e: