From ad1f2b509af6b4db02a0c9cbb5edd797933adc40 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 3 Sep 2024 11:11:33 +0200 Subject: [PATCH] Try-except and fetch cookies for flask --- aikido_firewall/sources/flask.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/aikido_firewall/sources/flask.py b/aikido_firewall/sources/flask.py index d1d845f6..822049bb 100644 --- a/aikido_firewall/sources/flask.py +++ b/aikido_firewall/sources/flask.py @@ -22,17 +22,21 @@ def aikido_view_func(*args, **kwargs): req = request_ctx.request # Set body : - context = get_current_context() - if context: - if req.is_json: - context.body = req.get_json() - elif req.form: - context.body = req.form - else: - context.body = req.data.decode("utf-8") - context.set_as_current_context() + try: + context = get_current_context() + if context: + if req.is_json: + context.body = req.get_json() + elif req.form: + context.body = req.form + else: + context.body = req.data.decode("utf-8") + context.cookies = req.cookies.to_dict() + context.set_as_current_context() - pre_response = request_handler(stage="pre_response") + pre_response = request_handler(stage="pre_response") + except Exception as e: + logger.debug("Exception in aikido's view function : %s", e) if pre_response: return pre_response[0], pre_response[1] try: