diff --git a/aikido_zen/context/__init__.py b/aikido_zen/context/__init__.py index 82de90b61..de3930878 100644 --- a/aikido_zen/context/__init__.py +++ b/aikido_zen/context/__init__.py @@ -6,6 +6,7 @@ import json from json import JSONDecodeError from time import sleep +from threading import RLock from typing import Dict from urllib.parse import parse_qs @@ -108,6 +109,36 @@ def set_as_current_context(self): def set_cookies(self, cookies): self.cookies = cookies + @property + def body(self): + self.load_body() + return self.__dict__.get("body") + + @body.setter + def body(self, value): + self.__dict__.pop("_body_loader", None) + self.__dict__["body"] = value + + def set_body_loader(self, loader): + self.__dict__["body"] = None + self.__dict__["_body_loader"] = loader + self.__dict__["_body_lock"] = RLock() + + def load_body(self): + lock = self.__dict__.get("_body_lock") + if lock is None: + return + with lock: + # The loader can read through hooks that access this property again. + loader = self.__dict__.pop("_body_loader", None) + if loader is not None: + try: + self.set_body(loader()) + except Exception as e: + logger.debug("Exception occurred whilst loading body: %s", e) + finally: + self.__dict__.pop("_body_lock", None) + def set_body(self, body): try: self.set_body_internal(body) diff --git a/aikido_zen/context/lazy_body_test.py b/aikido_zen/context/lazy_body_test.py new file mode 100644 index 000000000..04172a6f4 --- /dev/null +++ b/aikido_zen/context/lazy_body_test.py @@ -0,0 +1,102 @@ +import pickle +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout +from threading import Event +from unittest.mock import Mock + +import pytest + +from . import Context + + +def test_body_is_loaded_once_on_first_access(): + context = Context() + loader = Mock(return_value=b'{"name":"input"}') + context.set_body_loader(loader) + loader.assert_not_called() + assert context.body == {"name": "input"} + assert context.body == {"name": "input"} + loader.assert_called_once_with() + + +def test_body_loader_is_not_reentered(): + context = Context() + loader = Mock(side_effect=lambda: context.body or b"body") + context.set_body_loader(loader) + assert context.body == "body" + loader.assert_called_once_with() + + +def test_concurrent_body_access_waits_for_the_complete_body(): + context = Context() + started = Event() + second_started = Event() + release = Event() + + def load(): + started.set() + assert release.wait(5) + return b'{"name":"input"}' + + loader = Mock(side_effect=load) + context.set_body_loader(loader) + + def read_second(): + second_started.set() + return context.body + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(lambda: context.body) + assert started.wait(5) + second = pool.submit(read_second) + try: + assert second_started.wait(5) + with pytest.raises(FutureTimeout): + second.result(0.1) + finally: + release.set() + assert first.result(5) == {"name": "input"} + assert second.result(5) == {"name": "input"} + loader.assert_called_once_with() + + +def test_explicit_body_replaces_pending_loader(): + context = Context() + loader = Mock(return_value=b"old") + context.set_body_loader(loader) + context.set_body(b"replacement") + assert context.body == "replacement" + loader.assert_not_called() + + +def test_failed_loader_is_not_retried_on_every_inspection(): + context = Context() + loader = Mock(side_effect=ValueError("body unavailable")) + context.set_body_loader(loader) + assert context.body is None + assert context.body is None + loader.assert_called_once_with() + + +def test_pickle_loads_body_without_serializing_loader(): + context = Context() + loader = Mock(return_value=b'{"name":"input"}') + context.set_body_loader(loader) + restored = pickle.loads(pickle.dumps(context)) + assert restored.body == {"name": "input"} + assert "_body_loader" not in restored.__dict__ + loader.assert_called_once_with() + + +def test_body_loader_is_used_by_schema_discovery_without_application_read(): + from aikido_zen.api_discovery.get_api_info import get_body_info + from aikido_zen.helpers.headers import Headers + + context = Context() + context.headers = Headers({"CONTENT_TYPE": ["application/json"]}) + loader = Mock(return_value=b'{"name":"input"}') + context.set_body_loader(loader) + assert get_body_info(context) == { + "type": "json", + "schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + } + loader.assert_called_once_with() diff --git a/aikido_zen/helpers/lazy_body_stream.py b/aikido_zen/helpers/lazy_body_stream.py new file mode 100644 index 000000000..d67e396ea --- /dev/null +++ b/aikido_zen/helpers/lazy_body_stream.py @@ -0,0 +1,47 @@ +# Proxy members are supplied by wrapt's C extension. +# pylint: disable=abstract-method,no-member + +from wrapt import ObjectProxy + + +class LazyBodyStream(ObjectProxy): + _READ_METHODS = ( + "read", + "readall", + "read1", + "readinto", + "readinto1", + "readline", + "readlines", + "peek", + ) + + def __init__(self, stream, before_read): + super().__init__(stream) + self._self_before_read = before_read + + def _prepare(self): + callback = self._self_before_read + # Body capture reads this stream too, so it must not trigger itself. + self._self_before_read = None + if callback is not None: + callback() + + def __getattr__(self, name): + value = getattr(self.__wrapped__, name) + if name not in self._READ_METHODS or self._self_before_read is None: + return value + + def read(*args, **kwargs): + self._prepare() + return getattr(self.__wrapped__, name)(*args, **kwargs) + + return read + + def __iter__(self): + self._prepare() + return iter(self.__wrapped__) + + def __next__(self): + self._prepare() + return next(self.__wrapped__) diff --git a/aikido_zen/helpers/lazy_body_stream_test.py b/aikido_zen/helpers/lazy_body_stream_test.py new file mode 100644 index 000000000..d9b956d00 --- /dev/null +++ b/aikido_zen/helpers/lazy_body_stream_test.py @@ -0,0 +1,90 @@ +from io import BytesIO +from unittest.mock import Mock + +import pytest + +from .lazy_body_stream import LazyBodyStream + + +@pytest.mark.parametrize( + "method,args", [("read", (2,)), ("readline", ()), ("readlines", ())] +) +def test_capture_runs_before_reading_and_only_once(method, args): + original = BytesIO(b"first\nsecond\n") + callback = Mock(side_effect=lambda: original.seek(0)) + stream = LazyBodyStream(original, callback) + assert original.tell() == 0 + expected = getattr(BytesIO(b"first\nsecond\n"), method)(*args) + assert getattr(stream, method)(*args) == expected + stream.read() + callback.assert_called_once_with() + + +def test_readinto_is_captured(): + callback = Mock() + stream = LazyBodyStream(BytesIO(b"body"), callback) + buffer = bytearray(4) + assert stream.readinto(buffer) == 4 + assert buffer == b"body" + callback.assert_called_once_with() + + +def test_iteration_is_captured(): + callback = Mock() + stream = LazyBodyStream(BytesIO(b"first\nsecond\n"), callback) + assert list(stream) == [b"first\n", b"second\n"] + callback.assert_called_once_with() + + +def test_next_is_captured(): + callback = Mock() + stream = LazyBodyStream(BytesIO(b"first\nsecond\n"), callback) + assert next(stream) == b"first\n" + assert next(stream) == b"second\n" + callback.assert_called_once_with() + + +def test_loader_can_read_the_same_stream_without_recursion(): + captured = [] + stream = LazyBodyStream(BytesIO(b"body"), lambda: captured.append(stream.read())) + assert stream.read() == b"" + assert captured == [b"body"] + + +def test_read_uses_a_stream_replaced_by_the_loader(): + def replace(): + stream.__wrapped__ = BytesIO(b"replacement") + + stream = LazyBodyStream(BytesIO(b"original"), replace) + assert stream.read(3) == b"rep" + assert stream.read() == b"lacement" + + +def test_does_not_add_unsupported_read_methods(): + class Input: + def read(self, size=-1): + return b"body" + + stream = LazyBodyStream(Input(), Mock()) + assert not hasattr(stream, "readinto") + assert stream.read() == b"body" + + +def test_non_read_operations_do_not_load_body(): + callback = Mock() + stream = LazyBodyStream(BytesIO(b"body"), callback) + assert stream.tell() == 0 + assert stream.readable() + stream.close() + assert stream.closed + callback.assert_not_called() + + +def test_original_read_error_is_propagated(): + original = BytesIO(b"body") + original.close() + callback = Mock() + stream = LazyBodyStream(original, callback) + with pytest.raises(ValueError): + stream.read() + callback.assert_called_once_with() diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py index bc73df788..9e34b34b1 100644 --- a/aikido_zen/sources/flask/__init__.py +++ b/aikido_zen/sources/flask/__init__.py @@ -1,8 +1,19 @@ -from aikido_zen.context import Context +from functools import partial +from contextvars import copy_context +from weakref import proxy, ref + +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 aikido_zen.helpers.lazy_body_stream import LazyBodyStream +from aikido_zen.sinks import ( + on_import, + patch_function, + before_modify_return, + before, + 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_form_data import extract_body_from_flask_request 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,18 +26,73 @@ 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) + context = get_current_context() + if context is None or not ( + req.content_length or req.environ.get("wsgi.input_terminated") + ): + return + context.set_body_loader(partial(extract_body_from_flask_request, proxy(req))) + req.environ["aikido_zen.body_context"] = ref(context) + req.environ["wsgi.input"] = LazyBodyStream( + req.environ["wsgi.input"], context.load_body + ) + req.__dict__["stream"] = LazyBodyStream(req.stream, context.load_body) @after def _full_dispatch_request_after(func, instance, args, kwargs, return_value): if not hasattr(return_value, "status_code"): return + context = get_current_context() + if context is not None and "_body_loader" in context.__dict__: + from flask import has_request_context + from flask.globals import request_ctx + from werkzeug.wsgi import ClosingIterator + + # Response generators can need these bindings after Flask pops them. + saved_context = copy_context() + request_ref = ref(request_ctx.request) + loader = context.__dict__["_body_loader"] + + def load_body(): + try: + if saved_context is not None: + return saved_context.run(loader) + return loader() + finally: + # Deferred parsing can open uploads after Flask's request cleanup. + if ( + not has_request_context() + or request_ctx.request is not request_ref() + ): + req = request_ref() + if req is not None: + req.close() + + def release_request_context(): + nonlocal saved_context + saved_context = None + + context.set_body_loader(load_body) + if return_value.direct_passthrough: + return_value.response = ClosingIterator( + return_value.response, release_request_context + ) + else: + return_value.call_on_close(release_request_context) funcs.request_handler(stage="post_response", status_code=return_value.status_code) +@before +def _read_body_before(func, instance, args, kwargs): + context_ref = instance.environ.get("aikido_zen.body_context") + context = context_ref() if context_ref is not None else None + if context is not None: + context.load_body() + + @before_modify_return def _call(func, instance, args, kwargs): environ = get_argument(args, kwargs, 0, "environ") @@ -61,3 +127,9 @@ def patch(m): patch_function(m, "Flask.__call__", _call) patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_before) patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_after) + + +@on_import("flask.wrappers", "flask", version_requirement="2.2.4") +def patch_request(m): + patch_function(m, "Request.get_data", _read_body_before) + patch_function(m, "Request._load_form_data", _read_body_before) diff --git a/aikido_zen/sources/flask/body_test.py b/aikido_zen/sources/flask/body_test.py new file mode 100644 index 000000000..5c46dbb0f --- /dev/null +++ b/aikido_zen/sources/flask/body_test.py @@ -0,0 +1,476 @@ +import gc +import json +from io import BytesIO +from weakref import ref + +import pytest +from flask import Flask, Response, copy_current_request_context, request, send_file +from werkzeug.datastructures import MultiDict + +import aikido_zen.sources.flask as source +from aikido_zen.api_discovery.get_api_info import get_api_info +from aikido_zen.context import current_context, get_current_context +from aikido_zen.vulnerabilities.nosql_injection import detect_nosql_injection +from aikido_zen.vulnerabilities.sql_injection.context_contains_sql_injection import ( + context_contains_sql_injection, +) + + +class CountingInput(BytesIO): + reads = 0 + + def read(self, *args, **kwargs): + self.reads += 1 + return super().read(*args, **kwargs) + + def readinto(self, *args, **kwargs): + self.reads += 1 + return super().readinto(*args, **kwargs) + + +@pytest.fixture +def app(monkeypatch): + monkeypatch.setattr(source, "should_block", lambda: None) + monkeypatch.setattr(source.funcs, "request_handler", lambda **kwargs: None) + current_context.set(None) + yield Flask(__name__) + current_context.set(None) + + +@pytest.mark.parametrize("status", [404, 405, 308]) +@pytest.mark.parametrize("hooks", [False, True]) +def test_unread_error_body_is_not_captured(app, status, hooks): + app.add_url_rule("/known", view_func=lambda: "ok") + app.add_url_rule( + "/redirect/", endpoint="redirect", view_func=lambda: "ok", methods=["POST"] + ) + if hooks: + app.before_request(lambda: None) + app.after_request(lambda response: response) + app.register_error_handler(404, lambda error: ("missing", 404)) + stream = CountingInput(b"x" * 24 * 1024 * 1024) + path = {404: "/missing", 405: "/known", 308: "/redirect"}[status] + response = app.test_client().post( + path, input_stream=stream, content_type="application/octet-stream" + ) + assert response.status_code == status + assert stream.reads == 0 + assert get_current_context().__dict__["body"] is None + response.close() + + +@pytest.mark.parametrize( + "reader", + [ + "data", + "get_data", + "get_data_text", + "get_data_no_cache", + "json", + "get_json", + "get_json_force", + ], +) +def test_json_capture_before_body_reader_returns(app, reader): + @app.post("/body") + def view(): + if reader == "data": + data = request.data + elif reader == "get_data_text": + data = request.get_data(as_text=True) + elif reader == "get_data_no_cache": + data = request.get_data(cache=False) + elif reader == "json": + data = request.json + elif reader == "get_json_force": + data = request.get_json(force=True) + else: + data = getattr(request, reader)() + assert data + assert get_current_context().body == {"name": "value"} + return "ok" + + content_type = "text/plain" if reader == "get_json_force" else "application/json" + response = app.test_client().post( + "/body", data=b'{"name":"value"}', content_type=content_type + ) + assert response.status_code == 200 + response.close() + + +@pytest.mark.parametrize( + "reader", + [ + "read", + "readline", + "readlines", + "readinto", + "iter", + "next", + "input_stream", + "wsgi_input", + ], +) +def test_raw_stream_reads_keep_body_available_for_inspection(app, reader): + @app.post("/body") + def view(): + if reader == "wsgi_input": + request.environ["wsgi.input"].read() + elif reader == "input_stream": + request.input_stream.read() + elif reader == "readinto": + request.stream.readinto(bytearray(3)) + elif reader == "iter": + list(request.stream) + elif reader == "next": + next(request.stream, None) + else: + getattr(request.stream, reader)() + assert get_current_context().body == {"name": "value"} + return "ok" + + response = app.test_client().post( + "/body", data=b'{"name":"value"}', content_type="application/json" + ) + assert response.status_code == 200 + response.close() + + +def test_inspection_can_load_body_without_application_read(app): + attack = "Malicious Pet', 'Gru from the Minions') --" + + @app.post("/body") + def view(): + result = context_contains_sql_injection( + sql=f"INSERT INTO pets(name) VALUES ('{attack}')", + operation="psycopg2.Cursor.execute", + context=get_current_context(), + dialect="postgres", + ) + assert result["source"] == "body" + assert result["payload"] == attack + return "ok" + + response = app.test_client().post("/body", json={"name": attack}) + assert response.status_code == 200 + response.close() + + +def test_nosql_inspection_can_load_body_without_application_read(app): + @app.post("/body") + def view(): + result = detect_nosql_injection(get_current_context(), {"title": {"$ne": None}}) + assert result["source"] == "body" + assert result["injection"] + return "ok" + + response = app.test_client().post("/body", json={"title": {"$ne": None}}) + assert response.status_code == 200 + response.close() + + +def test_schema_discovery_can_load_unused_body(app, monkeypatch): + schemas = [] + + def discover(stage, status_code=0): + if stage == "post_response": + schemas.append(get_api_info(get_current_context())) + + monkeypatch.setattr(source.funcs, "request_handler", discover) + app.add_url_rule("/body", methods=["POST"], view_func=lambda: "ok") + response = app.test_client().post("/body", json={"name": "value"}) + assert response.status_code == 200 + assert schemas[0]["body"] == { + "type": "json", + "schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + } + response.close() + + +def test_form_fields_and_duplicates_are_preserved(app): + @app.post("/body") + def view(): + assert request.form.getlist("name") == ["first", "second"] + assert get_current_context().body.getlist("name") == ["first", "second"] + return "ok" + + response = app.test_client().post( + "/body", data=MultiDict([("name", "first"), ("name", "second")]) + ) + assert response.status_code == 200 + response.close() + + +def test_large_file_is_available_and_not_added_to_inspected_form(app): + @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 + assert request.form["name"] == "upload" + return "ok" + + response = app.test_client().post( + "/body", + data={ + "name": "upload", + "file": (BytesIO(b"x" * 10 * 1024 * 1024), "large.bin"), + }, + ) + assert response.status_code == 200 + response.close() + + +def test_large_json_remains_supported(app): + @app.post("/body") + def view(): + data = request.get_json() + assert len(data["padding"]) == 10 * 1024 * 1024 + assert get_current_context().body == data + return "ok" + + response = app.test_client().post("/body", json={"padding": "x" * 10 * 1024 * 1024}) + assert response.status_code == 200 + response.close() + + +def test_chunked_body_without_content_length_is_captured(app): + @app.post("/body") + def view(): + assert request.get_json() == {"name": "value"} + assert get_current_context().body == {"name": "value"} + return "ok" + + response = app.test_client().post( + "/body", + data=b'{"name":"value"}', + content_type="application/json", + environ_overrides={"CONTENT_LENGTH": "", "wsgi.input_terminated": True}, + ) + assert response.status_code == 200 + response.close() + + +def test_application_json_mutations_do_not_change_captured_input(app): + @app.post("/body") + def view(): + request.get_json()["name"] = "changed" + assert get_current_context().body == {"name": "original"} + return "ok" + + response = app.test_client().post("/body", json={"name": "original"}) + assert response.status_code == 200 + response.close() + + +def test_direct_passthrough_response_keeps_body_available(app): + @app.post("/body") + def view(): + def generate(): + gc.collect() + yield json.dumps(get_current_context().body).encode() + + return Response(generate(), direct_passthrough=True) + + response = app.test_client().post("/body", json={"name": "value"}, buffered=True) + assert response.status_code == 200 + assert json.loads(response.data) == {"name": "value"} + response.close() + + +def test_file_error_leaves_unused_body_unread(app, tmp_path): + path = tmp_path / "missing.txt" + path.write_bytes(b"missing") + # File I/O during the request legitimately triggers body inspection. + file = path.open("rb") + + @app.errorhandler(404) + def missing(error): + return send_file(file, download_name=path.name, mimetype="text/plain"), 404 + + stream = CountingInput(b"x" * 1024 * 1024) + response = app.test_client().post("/missing", input_stream=stream) + assert response.status_code == 404 + assert response.data == b"missing" + response.close() + assert stream.reads == 0 + assert file.closed + + +def test_routing_exception_traceback_is_preserved_after_response_close(app): + errors = [] + + @app.errorhandler(404) + def missing(error): + errors.append(error) + return "missing", 404 + + response = app.test_client().post("/missing", json={"name": "value"}) + error = errors[0] + traceback, context, cause = error.__traceback__, error.__context__, error.__cause__ + assert traceback is not None + response.close() + assert error.__traceback__ is traceback + assert error.__context__ is context + assert error.__cause__ is cause + + +@pytest.mark.parametrize("inspect_first", [False, True]) +@pytest.mark.parametrize("intervening_request", [False, True]) +@pytest.mark.parametrize("encoding", ["json", "form", "multipart"]) +def test_copied_request_context_can_capture_body_after_response_closes( + app, inspect_first, intervening_request, encoding +): + callbacks = [] + uploads = [] + attack = "Malicious Pet', 'Gru from the Minions') --" + app.add_url_rule("/other", view_func=lambda: "ok", methods=["POST"]) + + @app.errorhandler(404) + def missing(error): + context = get_current_context() + + @copy_current_request_context + def later(): + if inspect_first: + result = context_contains_sql_injection( + sql=f"INSERT INTO pets(name) VALUES ('{attack}')", + operation="psycopg2.Cursor.execute", + context=context, + dialect="postgres", + ) + assert result["source"] == "body" + if encoding == "json": + assert request.get_json() == {"name": attack} + else: + assert request.form["name"] == attack + assert context.body["name"] == attack + if encoding == "multipart": + assert request.files["file"].read() == b"upload" + uploads.append(request.files["file"].stream) + + callbacks.append(later) + return "missing", 404 + + client = app.test_client() + if encoding == "json": + response = client.post("/missing", json={"name": attack}) + else: + data = {"name": attack} + if encoding == "multipart": + data["file"] = (BytesIO(b"upload"), "file.txt") + response = client.post("/missing", data=data) + assert response.status_code == 404 + response.close() + assert get_current_context().__dict__["body"] is None + if intervening_request: + other = client.post("/other", json={"other": "request"}) + other.close() + callbacks.pop()() + if intervening_request: + assert get_current_context().__dict__["body"] is None + assert all(upload.closed for upload in uploads) + + +def test_body_can_be_captured_by_application_response_close_callback(app): + captured = [] + + @app.post("/body") + def view(): + req = request._get_current_object() + context = get_current_context() + response = Response("ok") + + def on_close(): + captured.append((req.get_json(), context.body)) + + response.call_on_close(on_close) + return response + + response = app.test_client().post("/body", json={"name": "value"}) + response.close() + assert captured == [({"name": "value"}, {"name": "value"})] + + +def test_streamed_response_can_inspect_body_after_request_teardown(app): + @app.post("/body") + def view(): + def generate(): + gc.collect() + yield json.dumps(get_current_context().body) + + return generate() + + response = app.test_client().post("/body", json={"name": "value"}, buffered=True) + assert response.status_code == 200 + assert json.loads(response.data) == {"name": "value"} + response.close() + + +def test_uploads_opened_during_response_streaming_are_closed(app): + streams = [] + + @app.post("/body") + def view(): + req = request._get_current_object() + + def generate(): + assert get_current_context().body.getlist("name") == ["upload"] + streams.append(req.files["file"].stream) + yield b"ok" + + return generate() + + response = app.test_client().post( + "/body", + data={"name": "upload", "file": (BytesIO(b"file"), "file.bin")}, + buffered=True, + ) + assert response.status_code == 200 + response.close() + assert streams[0].closed + + +def test_unread_request_is_not_retained_by_body_loader(app): + requests = [] + + @app.before_request + def remember(): + requests.append(ref(request._get_current_object())) + + response = app.test_client().post("/missing", data=b"body") + response.close() + del response + gc.collect() + assert requests[0]() is None + assert get_current_context().__dict__["body"] is None + + +@pytest.mark.parametrize( + "data,status", [(b'{"invalid":', 400), (b'{"name":"\xff"}', 400)] +) +def test_invalid_json_errors_are_preserved(app, data, status): + @app.post("/body") + def view(): + request.get_json() + return "ok" + + response = app.test_client().post( + "/body", data=data, content_type="application/json" + ) + assert response.status_code == status + response.close() + + +def test_application_body_limit_is_preserved(app): + app.config["MAX_CONTENT_LENGTH"] = 8 + + @app.post("/body") + def view(): + assert get_current_context().query == {"search": ["value"]} + request.get_data() + return "ok" + + response = app.test_client().post("/body?search=value", data=b"x" * 32) + assert response.status_code == 413 + response.close() diff --git a/aikido_zen/sources/flask/extract_form_data.py b/aikido_zen/sources/flask/extract_form_data.py index 991ed16cf..5ed4c4571 100644 --- a/aikido_zen/sources/flask/extract_form_data.py +++ b/aikido_zen/sources/flask/extract_form_data.py @@ -1,16 +1,11 @@ -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() +def extract_body_from_flask_request(req): try: - if context: - if req.form: - context.set_body(req.form) - else: - context.set_body(req.data) - context.set_as_current_context() + if req.form: + return req.form + return req.data except Exception as e: logger.debug("Exception occurred whilst extracting flask body data: %s", e) + return None