From cecc67880cda948950bf914ac295cacbd0ca531e Mon Sep 17 00:00:00 2001 From: teta2k Date: Sun, 1 Nov 2026 20:49:08 +0100 Subject: [PATCH 1/2] Capture Flask request bodies after application reads --- aikido_zen/sources/flask/__init__.py | 24 ++- aikido_zen/sources/flask/flask_test.py | 254 +++++++++++++++++++++++++ 2 files changed, 275 insertions(+), 3 deletions(-) diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py index bc73df788..e907286b7 100644 --- a/aikido_zen/sources/flask/__init__.py +++ b/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") diff --git a/aikido_zen/sources/flask/flask_test.py b/aikido_zen/sources/flask/flask_test.py index 1aa58a4b0..d4c928464 100644 --- a/aikido_zen/sources/flask/flask_test.py +++ b/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 From 40976727f6d937d9fa0a172db32d2ae35b4d239c Mon Sep 17 00:00:00 2001 From: teta2k Date: Mon, 2 Nov 2026 11:06:04 +0100 Subject: [PATCH 2/2] Remove obsolete Flask body extraction code --- aikido_zen/sources/flask/__init__.py | 2 +- aikido_zen/sources/flask/extract_form_data.py | 16 ---------------- 2 files changed, 1 insertion(+), 17 deletions(-) delete mode 100644 aikido_zen/sources/flask/extract_form_data.py diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py index e907286b7..fac83671f 100644 --- a/aikido_zen/sources/flask/__init__.py +++ b/aikido_zen/sources/flask/__init__.py @@ -1,4 +1,4 @@ -from aikido_zen.context import Context, get_current_context +from aikido_zen.context import 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 diff --git a/aikido_zen/sources/flask/extract_form_data.py b/aikido_zen/sources/flask/extract_form_data.py deleted file mode 100644 index 991ed16cf..000000000 --- a/aikido_zen/sources/flask/extract_form_data.py +++ /dev/null @@ -1,16 +0,0 @@ -from aikido_zen.context import get_current_context -from aikido_zen.helpers.logging import logger - - -def extract_form_data_from_flask_request_and_save_data(req): - """Extract form data from flask request""" - context = get_current_context() - try: - if context: - if req.form: - context.set_body(req.form) - else: - context.set_body(req.data) - context.set_as_current_context() - except Exception as e: - logger.debug("Exception occurred whilst extracting flask body data: %s", e)