From 16245eedb686475e65950d2d9b5ebc23b2b5915b Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 30 Jul 2024 14:17:00 +0200 Subject: [PATCH 1/7] Create a boilerplate pymongo wrap --- aikido_firewall/__init__.py | 1 + aikido_firewall/sinks/pymongo.py | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+) create mode 100644 aikido_firewall/sinks/pymongo.py diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 18eafc779..a50fb476e 100644 --- a/aikido_firewall/__init__.py +++ b/aikido_firewall/__init__.py @@ -27,6 +27,7 @@ def protect(module="any"): # Import sinks import aikido_firewall.sinks.pymysql + import aikido_firewall.sinks.pymongo logger.info("Aikido python firewall started") start_background_process() diff --git a/aikido_firewall/sinks/pymongo.py b/aikido_firewall/sinks/pymongo.py new file mode 100644 index 000000000..9c0291f98 --- /dev/null +++ b/aikido_firewall/sinks/pymongo.py @@ -0,0 +1,20 @@ + +""" +Sink module for `pymongo` +""" +from importlib.metadata import version +import importhook +from aikido_firewall.helpers.logging import logger + +@importhook.on_import("pymongo") +def on_pymongo_import(pymongo): + """ + Hook 'n wrap on `pymongo` + Our goal is to ... + https://github.com/... + Returns : Modified pymongo object + """ + modified_pymongo = importhook.copy_module(pymongo) + logger.debug("Wrapper - `pymongo` version : %s", version("pymongo")) + + return modified_pymongo From d8f9d55daf2201d09f89f82da5e4e39361b383d8 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 30 Jul 2024 16:03:34 +0200 Subject: [PATCH 2/7] Wrap the operations with filter --- aikido_firewall/sinks/pymongo.py | 47 +++++++++++++++++++++++++++----- 1 file changed, 40 insertions(+), 7 deletions(-) diff --git a/aikido_firewall/sinks/pymongo.py b/aikido_firewall/sinks/pymongo.py index 9c0291f98..a16ea3c80 100644 --- a/aikido_firewall/sinks/pymongo.py +++ b/aikido_firewall/sinks/pymongo.py @@ -1,20 +1,53 @@ - """ Sink module for `pymongo` """ + from importlib.metadata import version +from copy import deepcopy import importhook from aikido_firewall.helpers.logging import logger -@importhook.on_import("pymongo") +OPERATIONS_WITH_FILTER = [ + "replace_one", # L1087 + "update_one", # L1189 + "update_many", # L1302 + "delete_one", # L1542 + "delete_many", # L1607 + "find_one", # L1672 + "count_documents", # L2020 + "find_one_and_delete", # L3207 + "find_one_and_replace", # L3296 + "find_one_and_update", # L3403 +] + +# ISSUE : Asynchronous +# ISSUE : `find` on L1707 and `find_raw_batches` on L1895 +# ISSUE : `aggregate` on L2847 and `aggregate_raw_batches` on L2942 +# ISSUE : `distinct` on L3054 + + +# Synchronous : +@importhook.on_import("pymongo.collection") def on_pymongo_import(pymongo): """ - Hook 'n wrap on `pymongo` - Our goal is to ... - https://github.com/... - Returns : Modified pymongo object + Hook 'n wrap on `pymongo.collection` + Our goal is to wrap the following functions in the Collection class : + https://github.com/mongodb/mongo-python-driver/blob/98658cfd1fea42680a178373333bf27f41153759/pymongo/synchronous/collection.py#L136 + Returns : Modified pymongo.collection.Collection object """ + modified_pymongo = importhook.copy_module(pymongo) - logger.debug("Wrapper - `pymongo` version : %s", version("pymongo")) + for operation in OPERATIONS_WITH_FILTER: + if not hasattr(pymongo.Collection, operation): + logger.warning("Operation `%s` not found on Collection object.", operation) + old_function = deepcopy(getattr(pymongo.Collection, operation)) + + def wrapped_operation_function(_self, filter, *args, **kwargs): + logger.debug(" `pymongo.collection.Collection.%s` wrapped", operation) + logger.debug(filter) + old_function(_self, filter, *args, **kwargs) + + setattr(modified_pymongo.Collection, operation, wrapped_operation_function) + logger.debug("Wrapped `pymongo` with version %s", version("pymongo")) return modified_pymongo From a96bd46942d7ab4fe5dc827fe77ec0455e178f38 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 30 Jul 2024 16:58:43 +0200 Subject: [PATCH 3/7] Fix bug in NoSQL algorithm --- aikido_firewall/vulnerabilities/nosql_injection/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aikido_firewall/vulnerabilities/nosql_injection/__init__.py b/aikido_firewall/vulnerabilities/nosql_injection/__init__.py index 0637fad11..810c0c208 100644 --- a/aikido_firewall/vulnerabilities/nosql_injection/__init__.py +++ b/aikido_firewall/vulnerabilities/nosql_injection/__init__.py @@ -100,8 +100,8 @@ def detect_nosql_injection(request, _filter): return {"injection": False} for source in UINPUT_SOURCES: - if request.get(source): - result = find_filter_part_with_operators(request[source], _filter) + if hasattr(request, source): + result = find_filter_part_with_operators(getattr(request, source), _filter) if result.get("found"): return { From 0314dd58d6212c56b265e07f9201a6d458b52aad Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 30 Jul 2024 16:59:04 +0200 Subject: [PATCH 4/7] Get working injection algorithm --- aikido_firewall/sinks/pymongo.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/aikido_firewall/sinks/pymongo.py b/aikido_firewall/sinks/pymongo.py index a16ea3c80..3e6930a22 100644 --- a/aikido_firewall/sinks/pymongo.py +++ b/aikido_firewall/sinks/pymongo.py @@ -35,17 +35,23 @@ def on_pymongo_import(pymongo): https://github.com/mongodb/mongo-python-driver/blob/98658cfd1fea42680a178373333bf27f41153759/pymongo/synchronous/collection.py#L136 Returns : Modified pymongo.collection.Collection object """ - modified_pymongo = importhook.copy_module(pymongo) for operation in OPERATIONS_WITH_FILTER: if not hasattr(pymongo.Collection, operation): logger.warning("Operation `%s` not found on Collection object.", operation) - old_function = deepcopy(getattr(pymongo.Collection, operation)) - def wrapped_operation_function(_self, filter, *args, **kwargs): - logger.debug(" `pymongo.collection.Collection.%s` wrapped", operation) - logger.debug(filter) - old_function(_self, filter, *args, **kwargs) + prev_func = deepcopy(getattr(pymongo.Collection, operation)) + + def wrapped_operation_function( + _self, _filter, *args, prev_func=prev_func, op=operation, **kwargs + ): + context = get_current_context() + logger.debug("Context : %s, Filter : %s", context, _filter) + injection_results = detect_nosql_injection(context, _filter) + logger.debug("Injection results : %s", injection_results) + if injection_results["injection"]: + raise Exception(contains_nosql_injection) + return prev_func(_self, _filter, *args, **kwargs) setattr(modified_pymongo.Collection, operation, wrapped_operation_function) From 51ae58139df7c95192bf57a4b97739147379a600 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 30 Jul 2024 17:02:37 +0200 Subject: [PATCH 5/7] Add missing imports and report to bg process --- aikido_firewall/sinks/pymongo.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/aikido_firewall/sinks/pymongo.py b/aikido_firewall/sinks/pymongo.py index 3e6930a22..52d7f9b59 100644 --- a/aikido_firewall/sinks/pymongo.py +++ b/aikido_firewall/sinks/pymongo.py @@ -6,6 +6,9 @@ from copy import deepcopy import importhook from aikido_firewall.helpers.logging import logger +from aikido_firewall.vulnerabilities.nosql_injection import detect_nosql_injection +from aikido_firewall.context import get_current_context +from aikido_firewall.background_process import get_comms OPERATIONS_WITH_FILTER = [ "replace_one", # L1087 @@ -46,11 +49,10 @@ def wrapped_operation_function( _self, _filter, *args, prev_func=prev_func, op=operation, **kwargs ): context = get_current_context() - logger.debug("Context : %s, Filter : %s", context, _filter) injection_results = detect_nosql_injection(context, _filter) - logger.debug("Injection results : %s", injection_results) if injection_results["injection"]: - raise Exception(contains_nosql_injection) + get_comms().send_data("ATTACK", injection_results) + raise Exception("NOSQL Injection [aikido_firewall]") return prev_func(_self, _filter, *args, **kwargs) setattr(modified_pymongo.Collection, operation, wrapped_operation_function) From 164b0d46fdcdcb4795f313366b740d32613d844f Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 30 Jul 2024 17:16:59 +0200 Subject: [PATCH 6/7] Update nosql testing to use an actual context class and not a dict --- .../nosql_injection/init_test.py | 32 +++++++++++-------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/aikido_firewall/vulnerabilities/nosql_injection/init_test.py b/aikido_firewall/vulnerabilities/nosql_injection/init_test.py index d3458e688..449d13bd6 100644 --- a/aikido_firewall/vulnerabilities/nosql_injection/init_test.py +++ b/aikido_firewall/vulnerabilities/nosql_injection/init_test.py @@ -1,5 +1,6 @@ import pytest from aikido_firewall.vulnerabilities.nosql_injection import detect_nosql_injection +from aikido_firewall.context import Context @pytest.fixture @@ -7,19 +8,24 @@ def create_context(): def _create_context( query=None, headers=None, body=None, cookies=None, route_params=None ): - context = { - "remote_address": "::1", - "method": "GET", - "url": "http://localhost:4000", - "query": query if query else {}, - "headers": headers if headers else {}, - "body": body, - "cookies": cookies if cookies else {}, - "route_params": route_params if route_params else {}, - "source": "express", - "route": "/posts/:id", - } - return context + class RequestContext: + remote_address = "::1" + method = "GET" + url = "http://localhost:4000" + query = {} + headers = {} + body = None + cookies = {} + route_params = {} + source = "express" + route = "/posts/:id" + + RequestContext.query = query if query else {} + RequestContext.headers = headers if headers else {} + RequestContext.body = body + RequestContext.cookies = cookies if cookies else {} + RequestContext.route_params = route_params if route_params else {} + return RequestContext() return _create_context From 7c48ea97f53b38f8b2ad4a19725c12cd75029da9 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 30 Jul 2024 17:27:22 +0200 Subject: [PATCH 7/7] Add some more basic testing --- aikido_firewall/sinks/pymongo.py | 2 +- aikido_firewall/sinks/pymongo_test.py | 61 +++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 aikido_firewall/sinks/pymongo_test.py diff --git a/aikido_firewall/sinks/pymongo.py b/aikido_firewall/sinks/pymongo.py index 52d7f9b59..9b41f1a83 100644 --- a/aikido_firewall/sinks/pymongo.py +++ b/aikido_firewall/sinks/pymongo.py @@ -57,5 +57,5 @@ def wrapped_operation_function( setattr(modified_pymongo.Collection, operation, wrapped_operation_function) - logger.debug("Wrapped `pymongo` with version %s", version("pymongo")) + # logger.debug("Wrapped `pymongo` with version %s", version("pymongo")) return modified_pymongo diff --git a/aikido_firewall/sinks/pymongo_test.py b/aikido_firewall/sinks/pymongo_test.py new file mode 100644 index 000000000..3116a53cd --- /dev/null +++ b/aikido_firewall/sinks/pymongo_test.py @@ -0,0 +1,61 @@ +import pytest +from unittest.mock import MagicMock +from aikido_firewall.sinks.pymongo import on_pymongo_import + + +@pytest.fixture +def mock_pymongo(): + mock_pymongo = MagicMock() + mock_collection = MagicMock() + mock_pymongo.Collection = mock_collection + return mock_pymongo + + +def test_on_pymongo_import(mocker, mock_pymongo): + mocker.patch("importhook.copy_module", return_value=mock_pymongo) + + for operation in [ + "replace_one", + "update_one", + "update_many", + "delete_one", + "delete_many", + "find_one", + "count_documents", + "find_one_and_delete", + "find_one_and_replace", + "find_one_and_update", + ]: + setattr( + mock_pymongo.Collection, + operation, + MagicMock(return_value="original_result"), + ) + + mocker.patch("aikido_firewall.helpers.logging.logger") + mocker.patch( + "aikido_firewall.vulnerabilities.nosql_injection.detect_nosql_injection", + return_value={"injection": False}, + ) + mocker.patch("aikido_firewall.context.get_current_context", return_value={}) + mocker.patch( + "aikido_firewall.background_process.get_comms", + return_value=MagicMock(send_data=MagicMock()), + ) + + modified_pymongo = on_pymongo_import(mock_pymongo) + + for operation in [ + "replace_one", + "update_one", + "update_many", + "delete_one", + "delete_many", + "find_one", + "count_documents", + "find_one_and_delete", + "find_one_and_replace", + "find_one_and_update", + ]: + wrapped_function = getattr(modified_pymongo.Collection, operation) + assert wrapped_function is not None