Stop consuming Flask request bodies before the application

aikido_zen/sources/flask/__init__.py

@@ -1,8 +1,7 @@
-from aikido_zen.context import Context
+from aikido_zen.context import Context, get_current_context
 from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.sinks import on_import, patch_function, before_modify_return, after
 from .extract_cookies import extract_cookies_from_flask_request_and_save_data
-from .extract_form_data import extract_form_data_from_flask_request_and_save_data
 from .extract_query import extract_query_from_flask_request_and_save_data
 from .extract_view_args import extract_view_args_from_flask_request_and_save_data
 import aikido_zen.sources.functions.request_handler as funcs
@@ -15,7 +14,6 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
 
     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)
     extract_query_from_flask_request_and_save_data(req)
 
@@ -27,12 +25,32 @@ def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
     funcs.request_handler(stage="post_response", status_code=return_value.status_code)
 
 
+@after
+def _body_read_after(func, instance, args, kwargs, return_value):
+    context = instance.environ.get("aikido_zen.context")
+    if context is None:
+        return
+    form = instance.__dict__.get("form")
+    if form and context.body is not form:
+        context.set_body(form)
+        context.parsed_userinput.pop("body", None)
+    elif context.body is None and return_value:
+        context.set_body(return_value)
+
+
 @before_modify_return
 def _call(func, instance, args, kwargs):
     environ = get_argument(args, kwargs, 0, "environ")
     start_response = get_argument(args, kwargs, 1, "start_response")
 
+    request_class = instance.request_class
+    if not request_class.__dict__.get("_aikido_body_hooks_installed", False):
+        patch_function(request_class, "get_data", _body_read_after)
+        patch_function(request_class, "_load_form_data", _body_read_after)
+        setattr(request_class, "_aikido_body_hooks_installed", True)
+
     context1 = Context(req=environ, source="flask")
+    environ["aikido_zen.context"] = context1
     context1.set_as_current_context()
     funcs.request_handler(stage="init")
 

aikido_zen/sources/flask/flask_test.py

@@ -231,3 +231,257 @@ def test_startup_flask():
         app.run()
     except TimeoutException:
         pass
+
+
+@pytest.mark.parametrize("inspect_raw_first", [False, True])
+def test_form_body_after_raw_read_is_inspected(monkeypatch, inspect_raw_first):
+    from flask import Flask, request
+    from aikido_zen.helpers.extract_strings_from_context import (
+        extract_strings_from_context,
+    )
+    from aikido_zen.vulnerabilities.sql_injection.context_contains_sql_injection import (
+        context_contains_sql_injection,
+    )
+
+    monkeypatch.setattr(aikido_zen.sources.flask, "should_block", lambda: None)
+    monkeypatch.setattr(
+        aikido_zen.sources.flask.funcs, "request_handler", lambda **kwargs: None
+    )
+    app = Flask(__name__)
+    attack = "Malicious Pet', 'Gru from the Minions') --"
+
+    @app.post("/body")
+    def view():
+        request.get_data()
+        context = get_current_context()
+        if inspect_raw_first:
+            extract_strings_from_context(context)
+        name = request.form["name"]
+        result = context_contains_sql_injection(
+            sql=f"INSERT INTO pets(name) VALUES ('{name}')",
+            operation="psycopg2.Cursor.execute",
+            context=context,
+            dialect="postgres",
+        )
+        assert result["source"] == "body"
+        assert result["payload"] == attack
+        return "ok"
+
+    try:
+        response = app.test_client().post("/body", data={"name": attack})
+        assert response.status_code == 200
+        response.close()
+    finally:
+        current_context.set(None)
+
+
+@pytest.fixture
+def body_app(monkeypatch):
+    from flask import Flask
+
+    monkeypatch.setattr(aikido_zen.sources.flask, "should_block", lambda: None)
+    monkeypatch.setattr(
+        aikido_zen.sources.flask.funcs, "request_handler", lambda **kwargs: None
+    )
+    current_context.set(None)
+    yield Flask(__name__)
+    current_context.set(None)
+
+
+@pytest.mark.parametrize("status", [200, 404, 405, 308])
+@pytest.mark.parametrize("hooks", [False, True])
+def test_unused_body_is_not_read(body_app, status, hooks):
+    from io import BytesIO
+
+    body_app.add_url_rule("/known", endpoint="known", view_func=lambda: "ok")
+    body_app.add_url_rule(
+        "/ignored", endpoint="ignored", view_func=lambda: "ok", methods=["POST"]
+    )
+    body_app.add_url_rule(
+        "/redirect/", endpoint="redirect", view_func=lambda: "ok", methods=["POST"]
+    )
+    if hooks:
+        body_app.before_request(lambda: None)
+        body_app.after_request(lambda response: response)
+        body_app.register_error_handler(404, lambda error: ("missing", 404))
+    stream = BytesIO(b"x" * 1024 * 1024)
+    path = {200: "/ignored", 404: "/missing", 405: "/known", 308: "/redirect"}[status]
+    with patch.object(stream, "read", wraps=stream.read) as read, patch.object(
+        stream, "readinto", wraps=stream.readinto
+    ) as readinto:
+        response = body_app.test_client().post(path, input_stream=stream)
+        assert response.status_code == status
+        read.assert_not_called()
+        readinto.assert_not_called()
+        assert get_current_context().body is None
+        response.close()
+
+
+@pytest.mark.parametrize(
+    "reader",
+    [
+        "stream",
+        "chunks",
+        "input_stream",
+        "wsgi_input",
+        "data",
+        "get_data",
+        "get_data_text",
+        "get_data_no_cache",
+        "get_json",
+        "json",
+    ],
+)
+def test_body_readers_preserve_application_data(body_app, reader):
+    import json
+    from flask import request
+
+    data = {"name": "original"}
+    payload = json.dumps(data).encode()
+
+    @body_app.post("/body")
+    def view():
+        if reader == "stream":
+            result = request.stream.read()
+        elif reader == "chunks":
+            result = request.stream.read(3) + request.stream.read()
+        elif reader == "input_stream":
+            result = request.input_stream.read()
+        elif reader == "wsgi_input":
+            result = request.environ["wsgi.input"].read()
+        elif reader == "get_data_text":
+            result = request.get_data(as_text=True)
+        elif reader == "get_data_no_cache":
+            result = request.get_data(cache=False)
+        elif reader in ("data", "json"):
+            result = getattr(request, reader)
+        else:
+            result = getattr(request, reader)()
+        if reader in ("get_json", "json"):
+            assert result == data
+        elif reader == "get_data_text":
+            assert result == payload.decode()
+        else:
+            assert result == payload
+        if reader in ("stream", "chunks", "input_stream", "wsgi_input"):
+            assert get_current_context().body is None
+        else:
+            assert get_current_context().body == data
+        return "ok"
+
+    response = body_app.test_client().post(
+        "/body", data=payload, content_type="application/json"
+    )
+    assert response.status_code == 200
+    response.close()
+
+
+def test_large_upload_remains_available_and_only_form_fields_are_captured(body_app):
+    from io import BytesIO
+    from flask import request
+
+    @body_app.post("/body")
+    def view():
+        assert len(request.files["file"].read()) == 10 * 1024 * 1024
+        assert get_current_context().body.getlist("name") == ["upload"]
+        assert "file" not in get_current_context().body
+        return "ok"
+
+    response = body_app.test_client().post(
+        "/body",
+        data={"name": "upload", "file": (BytesIO(b"x" * 10 * 1024 * 1024), "x.bin")},
+    )
+    assert response.status_code == 200
+    response.close()
+
+
+@pytest.mark.parametrize("reader", ["json", "form", "data"])
+def test_body_read_during_session_setup_is_inspected(body_app, reader):
+    import json
+    from flask import session
+    from flask.sessions import SecureCookieSession, SecureCookieSessionInterface
+    from aikido_zen.vulnerabilities.sql_injection.context_contains_sql_injection import (
+        context_contains_sql_injection,
+    )
+
+    attack = "Malicious Pet', 'Gru from the Minions') --"
+
+    class BodySession(SecureCookieSessionInterface):
+        def open_session(self, app, req):
+            if reader == "json":
+                data = req.get_json()
+            elif reader == "form":
+                data = req.form
+            else:
+                data = json.loads(req.get_data())
+            return SecureCookieSession({"name": data["name"]})
+
+    body_app.secret_key = "test-key"
+    body_app.session_interface = BodySession()
+
+    @body_app.post("/body")
+    def view():
+        name = session["name"]
+        result = context_contains_sql_injection(
+            sql=f"INSERT INTO pets(name) VALUES ('{name}')",
+            operation="psycopg2.Cursor.execute",
+            context=get_current_context(),
+            dialect="postgres",
+        )
+        assert result["source"] == "body"
+        assert result["payload"] == attack
+        return "ok"
+
+    data = (
+        {"data": {"name": attack}} if reader == "form" else {"json": {"name": attack}}
+    )
+    response = body_app.test_client().post("/body", **data)
+    assert response.status_code == 200
+    response.close()
+
+
+@pytest.mark.parametrize("reader", ["json", "data", "get_data"])
+def test_custom_request_body_reader_is_inspected(body_app, reader):
+    import json
+    from flask import Request, request
+    from werkzeug.wrappers import Request as WerkzeugRequest
+    from aikido_zen.vulnerabilities.sql_injection.context_contains_sql_injection import (
+        context_contains_sql_injection,
+    )
+
+    attack = "Malicious Pet', 'Gru from the Minions') --"
+
+    class CustomRequest(Request):
+        def get_data(self, *args, **kwargs):
+            return WerkzeugRequest.get_data(self, *args, **kwargs)
+
+    @body_app.post("/body")
+    def view():
+        if reader == "json":
+            name = request.get_json()["name"]
+        else:
+            data = request.data if reader == "data" else request.get_data()
+            name = json.loads(data)["name"]
+        result = context_contains_sql_injection(
+            sql=f"INSERT INTO pets(name) VALUES ('{name}')",
+            operation="psycopg2.Cursor.execute",
+            context=get_current_context(),
+            dialect="postgres",
+        )
+        assert result["source"] == "body"
+        assert result["payload"] == attack
+        return "ok"
+
+    client = body_app.test_client()
+    response = client.post("/body", json={"name": attack})
+    assert response.status_code == 200
+    response.close()
+    body_app.request_class = CustomRequest
+    response = client.post("/body", json={"name": attack})
+    assert response.status_code == 200
+    response.close()
+    wrapped = CustomRequest.__dict__["get_data"]
+    response = client.post("/body", json={"name": attack})
+    assert response.status_code == 200
+    response.close()
+    assert CustomRequest.__dict__["get_data"] is wrapped