From a6ec4c687179f2619f36e829f924e7380f310f48 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 19 Jul 2024 16:57:38 +0200 Subject: [PATCH 1/6] Add a new /create page with a form --- .../django-mysql/sample_app/templates/app/create_dog.html | 6 ++++++ sample-apps/django-mysql/sample_app/urls.py | 5 +++-- sample-apps/django-mysql/sample_app/views.py | 6 +++++- 3 files changed, 14 insertions(+), 3 deletions(-) create mode 100644 sample-apps/django-mysql/sample_app/templates/app/create_dog.html diff --git a/sample-apps/django-mysql/sample_app/templates/app/create_dog.html b/sample-apps/django-mysql/sample_app/templates/app/create_dog.html new file mode 100644 index 000000000..b2375a804 --- /dev/null +++ b/sample-apps/django-mysql/sample_app/templates/app/create_dog.html @@ -0,0 +1,6 @@ +

Create a Dog

+
+ + + +
diff --git a/sample-apps/django-mysql/sample_app/urls.py b/sample-apps/django-mysql/sample_app/urls.py index f5fc991de..96f6509ea 100644 --- a/sample-apps/django-mysql/sample_app/urls.py +++ b/sample-apps/django-mysql/sample_app/urls.py @@ -5,5 +5,6 @@ urlpatterns = [ path("", views.index, name="index"), path("dogpage/", views.dog_page, name="dog_page"), - path("create/", views.create_dogpage, name="create"), -] \ No newline at end of file + path("create/", views.create_dogpage, name="create_old"), + path("create", views.create, name="create") +] diff --git a/sample-apps/django-mysql/sample_app/views.py b/sample-apps/django-mysql/sample_app/views.py index c634cc717..cd9fa50e7 100644 --- a/sample-apps/django-mysql/sample_app/views.py +++ b/sample-apps/django-mysql/sample_app/views.py @@ -28,4 +28,8 @@ def create_dogpage(request, dog_name): query = 'INSERT INTO sample_app_dogs (dog_name, dog_boss) VALUES ("%s", "N/A")' % dog_name print("QUERY : ", query) cursor.execute(query) - return HttpResponse("Dog page created") \ No newline at end of file + return HttpResponse("Dog page created") + +def create(request): + template = loader.get_template("app/create_dog.html") + return HttpResponse(template.render({}, request)) From 4e063b6272f14bbc57070df7b8fe2a6436272ed1 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 19 Jul 2024 16:58:06 +0200 Subject: [PATCH 2/6] if the module is django, don't import flask --- aikido_firewall/__init__.py | 6 ++++-- sample-apps/django-mysql/manage.py | 2 ++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 292c8aed9..aa4910279 100644 --- a/aikido_firewall/__init__.py +++ b/aikido_firewall/__init__.py @@ -15,11 +15,13 @@ load_dotenv() -def protect(): +def protect(module="any"): """Start Aikido agent""" # Import sources import aikido_firewall.sources.django - import aikido_firewall.sources.flask + + if module != "django": + import aikido_firewall.sources.flask # Import sinks import aikido_firewall.sinks.pymysql diff --git a/sample-apps/django-mysql/manage.py b/sample-apps/django-mysql/manage.py index a3de10f7d..f66fda4b7 100755 --- a/sample-apps/django-mysql/manage.py +++ b/sample-apps/django-mysql/manage.py @@ -1,6 +1,8 @@ #!/usr/bin/env python """Django's command-line utility for administrative tasks.""" import aikido_firewall # Aikido module +aikido_firewall.protect("django") + import os import sys From 00bad8ad12d70b8e93c4f40a2c56bcea851222e1 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 19 Jul 2024 16:59:07 +0200 Subject: [PATCH 3/6] Update context function --- aikido_firewall/context/__init__.py | 47 ++++++++++++++++++++++------ aikido_firewall/context/init_test.py | 6 ++-- aikido_firewall/middleware/django.py | 3 ++ aikido_firewall/sources/flask.py | 2 +- 4 files changed, 46 insertions(+), 12 deletions(-) diff --git a/aikido_firewall/context/__init__.py b/aikido_firewall/context/__init__.py index 723b0d1a7..ffd0e809f 100644 --- a/aikido_firewall/context/__init__.py +++ b/aikido_firewall/context/__init__.py @@ -3,7 +3,7 @@ """ import threading - +SUPPORTED_SOURCES = ["django", "flask"] local = threading.local() @@ -15,21 +15,50 @@ def get_current_context(): return None +def parse_headers(headers): + """Parse EnvironHeaders object into a dict""" + if isinstance(headers, dict): + return headers + return dict(zip(headers.keys(), headers.values())) + + class Context: """ A context object, it stores everything that is important for vulnerability detection """ - def __init__(self, req): + def __init__(self, req, source): + if not source in SUPPORTED_SOURCES: + raise ValueError(f"Source {source} not supported") + self.source = source + self.method = req.method - self.remote_address = req.remote_addr - self.url = req.url - self.body = req.form - self.headers = req.headers - self.query = req.args - self.cookies = req.cookies - self.source = "flask" + if source == "flask": + self.remote_address = req.remote_addr + elif source == "django": + self.remote_address = req.META.get("REMOTE_ADDR") + + if source == "flask": + self.url = req.url + elif source == "django": + self.url = req.build_absolute_uri() + + if source == "flask": + self.body = req.form.to_dict() + elif source == "django": + self.body = dict(req.POST) + + self.headers = parse_headers(req.headers) + if source == "flask": + self.query = req.args.to_dict() + elif source == "django": + self.query = dict(req.GET) + + if source == "flask": + self.cookies = req.cookies.to_dict() + elif source == "django": + self.cookies = req.COOKIES def __reduce__(self): return ( diff --git a/aikido_firewall/context/init_test.py b/aikido_firewall/context/init_test.py index 0c247666e..2feb2346c 100644 --- a/aikido_firewall/context/init_test.py +++ b/aikido_firewall/context/init_test.py @@ -25,13 +25,15 @@ def test_get_current_context_no_context(): def test_set_as_current_context(sample_request): # Test set_as_current_context() method - context = Context(sample_request) + sample_request = mocker.MagicMock() + context = Context(sample_request, "flask") context.set_as_current_context() assert get_current_context() == context def test_get_current_context_with_context(sample_request): # Test get_current_context() when a context is set - context = Context(sample_request) + sample_request = mocker.MagicMock() + context = Context(sample_request, "flask") context.set_as_current_context() assert get_current_context() == context diff --git a/aikido_firewall/middleware/django.py b/aikido_firewall/middleware/django.py index c5fe8bf2b..276387b36 100644 --- a/aikido_firewall/middleware/django.py +++ b/aikido_firewall/middleware/django.py @@ -4,6 +4,7 @@ """ from aikido_firewall.helpers.logging import logger +from aikido_firewall.context import Context class AikidoMiddleware: @@ -16,6 +17,8 @@ def __init__(self, get_response): def __call__(self, request, *args, **kwargs): logger.debug("Aikido middleware for `django` was called : __call__") + context = Context(request, "django") + context.set_as_current_context() return self.get_response(request) def process_exception(self, request, exception): diff --git a/aikido_firewall/sources/flask.py b/aikido_firewall/sources/flask.py index 99a477320..d5d92cced 100644 --- a/aikido_firewall/sources/flask.py +++ b/aikido_firewall/sources/flask.py @@ -21,7 +21,7 @@ def __init__(self): def dispatch(self, request, call_next): """Dispatch function""" logger.debug("Aikido middleware for `flask` was called") - context = Context(request) + context = Context(request, "flask") context.set_as_current_context() response = call_next(request) From ebc982178aa3ab9855ee7951c8dd2b2caccbb2b9 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 19 Jul 2024 17:00:44 +0200 Subject: [PATCH 4/6] Fix broken tests and linting --- aikido_firewall/context/__init__.py | 1 + aikido_firewall/context/init_test.py | 20 ++------------------ 2 files changed, 3 insertions(+), 18 deletions(-) diff --git a/aikido_firewall/context/__init__.py b/aikido_firewall/context/__init__.py index ffd0e809f..adfa12b1d 100644 --- a/aikido_firewall/context/__init__.py +++ b/aikido_firewall/context/__init__.py @@ -3,6 +3,7 @@ """ import threading + SUPPORTED_SOURCES = ["django", "flask"] local = threading.local() diff --git a/aikido_firewall/context/init_test.py b/aikido_firewall/context/init_test.py index 2feb2346c..d01d1680a 100644 --- a/aikido_firewall/context/init_test.py +++ b/aikido_firewall/context/init_test.py @@ -2,28 +2,12 @@ from aikido_firewall.context import Context, get_current_context -@pytest.fixture -def sample_request(): - # Mock a sample request object for testing - class Request: - def __init__(self): - self.method = "GET" - self.remote_addr = "127.0.0.1" - self.url = "/test" - self.form = {} - self.headers = {} - self.args = {} - self.cookies = {} - - return Request() - - def test_get_current_context_no_context(): # Test get_current_context() when no context is set assert get_current_context() is None -def test_set_as_current_context(sample_request): +def test_set_as_current_context(mocker): # Test set_as_current_context() method sample_request = mocker.MagicMock() context = Context(sample_request, "flask") @@ -31,7 +15,7 @@ def test_set_as_current_context(sample_request): assert get_current_context() == context -def test_get_current_context_with_context(sample_request): +def test_get_current_context_with_context(mocker): # Test get_current_context() when a context is set sample_request = mocker.MagicMock() context = Context(sample_request, "flask") From e2d21a90935a5a0b264a515757359b33993aa1e4 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 19 Jul 2024 17:05:29 +0200 Subject: [PATCH 5/6] Add testing to context shizzle --- aikido_firewall/context/init_test.py | 42 ++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/aikido_firewall/context/init_test.py b/aikido_firewall/context/init_test.py index d01d1680a..f73c06e33 100644 --- a/aikido_firewall/context/init_test.py +++ b/aikido_firewall/context/init_test.py @@ -21,3 +21,45 @@ def test_get_current_context_with_context(mocker): context = Context(sample_request, "flask") context.set_as_current_context() assert get_current_context() == context + + +def test_context_init_flask(mocker): + req = mocker.MagicMock() + req.method = "GET" + req.remote_addr = "127.0.0.1" + req.url = "http://example.com" + req.form.to_dict.return_value = {"key": "value"} + req.headers = {"Content-Type": "application/json"} + req.args.to_dict.return_value = {"key": "value"} + req.cookies.to_dict.return_value = {"cookie": "value"} + + context = Context(req, "flask") + assert context.source == "flask" + assert context.method == "GET" + assert context.remote_address == "127.0.0.1" + assert context.url == "http://example.com" + assert context.body == {"key": "value"} + assert context.headers == {"Content-Type": "application/json"} + assert context.query == {"key": "value"} + assert context.cookies == {"cookie": "value"} + + +def test_context_init_django(mocker): + req = mocker.MagicMock() + req.method = "POST" + req.META.get.return_value = "127.0.0.1" + req.build_absolute_uri.return_value = "http://example.com" + req.POST = {"key": "value"} + req.headers = {"Content-Type": "application/json"} + req.GET = {"key": "value"} + req.COOKIES = {"cookie": "value"} + + context = Context(req, "django") + assert context.source == "django" + assert context.method == "POST" + assert context.remote_address == "127.0.0.1" + assert context.url == "http://example.com" + assert context.body == {"key": "value"} + assert context.headers == {"Content-Type": "application/json"} + assert context.query == {"key": "value"} + assert context.cookies == {"cookie": "value"} From 757c30e6239512ed6ef21e82fc90cac14b604ee0 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 19 Jul 2024 18:22:12 +0200 Subject: [PATCH 6/6] Split django and flask settings in seperate functions --- aikido_firewall/context/__init__.py | 41 +++++++++++++---------------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/aikido_firewall/context/__init__.py b/aikido_firewall/context/__init__.py index adfa12b1d..53910f4ab 100644 --- a/aikido_firewall/context/__init__.py +++ b/aikido_firewall/context/__init__.py @@ -33,33 +33,28 @@ def __init__(self, req, source): if not source in SUPPORTED_SOURCES: raise ValueError(f"Source {source} not supported") self.source = source - self.method = req.method - if source == "flask": - self.remote_address = req.remote_addr - elif source == "django": - self.remote_address = req.META.get("REMOTE_ADDR") - - if source == "flask": - self.url = req.url - elif source == "django": - self.url = req.build_absolute_uri() - - if source == "flask": - self.body = req.form.to_dict() - elif source == "django": - self.body = dict(req.POST) - self.headers = parse_headers(req.headers) if source == "flask": - self.query = req.args.to_dict() - elif source == "django": - self.query = dict(req.GET) - - if source == "flask": - self.cookies = req.cookies.to_dict() + self.set_flask_attrs(req) elif source == "django": - self.cookies = req.COOKIES + self.set_django_attrs(req) + + def set_django_attrs(self, req): + """set properties that are specific to django""" + self.remote_address = req.META.get("REMOTE_ADDR") + self.url = req.build_absolute_uri() + self.body = dict(req.POST) + self.query = dict(req.GET) + self.cookies = req.COOKIES + + def set_flask_attrs(self, req): + """Set properties that are specific to flask""" + self.remote_address = req.remote_addr + self.url = req.url + self.body = req.form.to_dict() + self.query = req.args.to_dict() + self.cookies = req.cookies.to_dict() def __reduce__(self): return (