diff --git a/aikido_zen/sources/flask_test.py b/aikido_zen/sources/flask_test.py index f4cf11ee..89bd184d 100644 --- a/aikido_zen/sources/flask_test.py +++ b/aikido_zen/sources/flask_test.py @@ -10,7 +10,7 @@ "HTTP_HEADER_1": "header 1 value", "HTTP_HEADER_2": "Header 2 value", "RANDOM_VALUE": "Random value", - "HTTP_COOKIE": "sessionklalkdlkjasdlkj@@as2@@d;a@@2sd=asd@@2;as@2d;'asd'aksksId=abc123xyz456;", + "HTTP_COOKIE": "sessionId=abc123xyz456;", "wsgi.url_scheme": "https", "HTTP_HOST": "example.com", "PATH_INFO": "/hello", @@ -35,6 +35,20 @@ "REMOTE_ADDR": "198.51.100.23", } +sample_environ_view_args_and_malformed_json = { + "REQUEST_METHOD": "POST", + "HTTP_COOKIE": "sessionId=abc123xyz456;", + "HTTP_HEADER_1": "header 1 value", + "HTTP_HEADER_2": "Header 2 value", + "HTTP_HOST": "example.com", + "CONTENT_TYPE": "application/json", + "PATH_INFO": "/hello/JohnDoe/30", + "QUERY_STRING": "", + "body": '{"invalid_json": true', + "REMOTE_ADDR": "198.51.100.23", + "wsgi.url_scheme": "https", +} + class TimeoutException(Exception): pass @@ -46,6 +60,48 @@ def timeout_handler(signum, frame): signal.signal(signal.SIGALRM, timeout_handler) +def test_flask_all_3_func_with_view_args_and_invalid_json_body(): + with patch("aikido_zen.sources.functions.request_handler.request_handler") as mock_request_handler: + reset_comms() + current_context.set(None) + mock_request_handler.return_value = None + + from flask import Flask + + app = Flask(__name__) + + @app.route("/hello//", methods=["POST"]) + def hello(user, age): + return f"User: {user}, Age: {age}" + + try: + signal.alarm(1) + + app(sample_environ_view_args_and_malformed_json, lambda x, y: x) + app.run() + + except TimeoutException: + pass + + assert get_current_context().method == "POST" + assert get_current_context().body is None + assert get_current_context().headers == { + "COOKIE": "sessionId=abc123xyz456;", + "HEADER_1": "header 1 value", + "HEADER_2": "Header 2 value", + "HOST": "example.com", + "CONTENT_TYPE": "application/json", + } + calls = mock_request_handler.call_args_list + assert len(calls) == 3 + assert calls[0][1]["stage"] == "init" + assert calls[1][1]["stage"] == "pre_response" + assert calls[2][1]["stage"] == "post_response" + assert calls[2][1]["status_code"] == 200 + + assert get_current_context().route_params["user"] == "JohnDoe" + assert get_current_context().route_params["age"] == "30" + def test_flask_all_3_func_with_invalid_body(): """When the flask body can not be parsed (because it contains invalid json for example), we should still parse the cookies of the endpoint""" with patch( @@ -142,7 +198,7 @@ def test_flask_all_3_func(): assert get_current_context().method == "POST" assert get_current_context().body == None assert get_current_context().headers == { - "COOKIE": "sessionklalkdlkjasdlkj@@as2@@d;a@@2sd=asd@@2;as@2d;'asd'aksksId=abc123xyz456;", + "COOKIE": "sessionId=abc123xyz456;", "HEADER_1": "header 1 value", "HEADER_2": "Header 2 value", "HOST": "example.com",