From 50b4219e53464c94198d56a844cc50f5b39fb5d8 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 10:26:28 +0200 Subject: [PATCH 01/24] create helper function to register a call --- aikido_zen/helpers/register_call.py | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 aikido_zen/helpers/register_call.py diff --git a/aikido_zen/helpers/register_call.py b/aikido_zen/helpers/register_call.py new file mode 100644 index 00000000..76f6cf61 --- /dev/null +++ b/aikido_zen/helpers/register_call.py @@ -0,0 +1,7 @@ +from aikido_zen.thread.thread_cache import get_cache + + +def register_call(operation, kind): + cache = get_cache() + if cache: + cache.stats.operations.register_call(operation, kind) From a83481ad6170c020c8c167e1d5c47c7670ba03d7 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 10:26:42 +0200 Subject: [PATCH 02/24] Pass along operation to statistics --- aikido_zen/vulnerabilities/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/aikido_zen/vulnerabilities/__init__.py b/aikido_zen/vulnerabilities/__init__.py index 006a5724..fbc141aa 100644 --- a/aikido_zen/vulnerabilities/__init__.py +++ b/aikido_zen/vulnerabilities/__init__.py @@ -99,7 +99,8 @@ def run_vulnerability_scan(kind, op, args): logger.debug("Injection results : %s", serialize_to_json(injection_results)) blocked = is_blocking_enabled() - thread_cache.stats.on_detected_attack(blocked) + operation = injection_results["operation"] + thread_cache.stats.on_detected_attack(blocked, operation) stack = get_clean_stacktrace() From b87af0541b9b25cf492cd45cd1a7e403fd5b2c8c Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 10:26:50 +0200 Subject: [PATCH 03/24] Create new operations dict-like object --- aikido_zen/storage/statistics/operations.py | 31 +++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 aikido_zen/storage/statistics/operations.py diff --git a/aikido_zen/storage/statistics/operations.py b/aikido_zen/storage/statistics/operations.py new file mode 100644 index 00000000..8baab062 --- /dev/null +++ b/aikido_zen/storage/statistics/operations.py @@ -0,0 +1,31 @@ +SUPPORTED_KINDS = ["sql_op", "nosql_op", "outgoing_http_op", "fs_op", "exec_op"] + + +class Operations(dict): + def __init__(self): + super().__init__() + + def ensure_operation(self, operation, kind): + if not kind in SUPPORTED_KINDS: + raise Exception(f"Kind {kind} is not supported for operations.") + if not operation in self.keys(): + self[operation] = { + "kind": kind, + "total": 0, + "attacksDetected": { + "total": 0, + "blocked": 0, + }, + } + + def register_call(self, operation, kind): + self.ensure_operation(operation, kind) + self[operation]["total"] += 1 + + def on_detected_attack(self, blocked, operation): + if operation not in self.keys(): + return + + self[operation]["attacksDetected"]["total"] += 1 + if blocked: + self[operation]["attacksDetected"]["blocked"] += 1 From 8162818259b4e40fd5e395f18ea8dacaa83ddf9f Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 10:27:21 +0200 Subject: [PATCH 04/24] integrate operations into statistics object --- aikido_zen/storage/statistics/__init__.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/aikido_zen/storage/statistics/__init__.py b/aikido_zen/storage/statistics/__init__.py index 93a08e69..9ee68973 100644 --- a/aikido_zen/storage/statistics/__init__.py +++ b/aikido_zen/storage/statistics/__init__.py @@ -1,4 +1,5 @@ import aikido_zen.helpers.get_current_unixtime_ms as t +from aikido_zen.storage.statistics.operations import Operations class Statistics: @@ -12,20 +13,23 @@ def __init__(self): self.attacks_detected = 0 self.attacks_blocked = 0 self.started_at = t.get_unixtime_ms() + self.operations = Operations() def clear(self): self.total_hits = 0 self.attacks_detected = 0 self.attacks_blocked = 0 self.started_at = t.get_unixtime_ms() + self.operations.clear() def increment_total_hits(self): self.total_hits += 1 - def on_detected_attack(self, blocked): + def on_detected_attack(self, blocked, operation): self.attacks_detected += 1 if blocked: self.attacks_blocked += 1 + self.operations.on_detected_attack(blocked, operation) def get_record(self): current_time = t.get_unixtime_ms() @@ -40,6 +44,7 @@ def get_record(self): "blocked": self.attacks_blocked, }, }, + "operations": self.operations, } def import_from_record(self, record): @@ -47,10 +52,13 @@ def import_from_record(self, record): self.total_hits += record.get("requests", {}).get("total", 0) self.attacks_detected += attacks_detected.get("total", 0) self.attacks_blocked += attacks_detected.get("blocked", 0) + self.operations.update(record.get("operations", {})) def empty(self): if self.total_hits > 0: return False if self.attacks_detected > 0: return False + if len(self.operations) > 0: + return False return True From b80fb225a5e1a0f4ced0bbf797e5acffdb21d6b3 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 10:27:39 +0200 Subject: [PATCH 05/24] Add register_call to pymysql --- aikido_zen/sinks/pymysql.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aikido_zen/sinks/pymysql.py b/aikido_zen/sinks/pymysql.py index c100039b..9405473e 100644 --- a/aikido_zen/sinks/pymysql.py +++ b/aikido_zen/sinks/pymysql.py @@ -4,6 +4,7 @@ import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import patch_function, on_import, before @@ -14,6 +15,7 @@ def _execute(func, instance, args, kwargs): # If query is type bytearray, it will be picked up by our wrapping of executemany return + register_call("pymysql.Cursor.execute", "sql_op") vulns.run_vulnerability_scan( kind="sql_injection", op="pymysql.Cursor.execute", args=(query, "mysql") ) @@ -23,6 +25,7 @@ def _execute(func, instance, args, kwargs): def _executemany(func, instance, args, kwargs): query = get_argument(args, kwargs, 0, "query") + register_call("pymysql.Cursor.executemany", "sql_op") vulns.run_vulnerability_scan( kind="sql_injection", op="pymysql.Cursor.executemany", args=(query, "mysql") ) From 37cf792e3783dfec0ed3a2d09b8a36b109557824 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 10:27:53 +0200 Subject: [PATCH 06/24] Add register_call to pymongo --- aikido_zen/sinks/pymongo.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/aikido_zen/sinks/pymongo.py b/aikido_zen/sinks/pymongo.py index 692ff87e..c9535c38 100644 --- a/aikido_zen/sinks/pymongo.py +++ b/aikido_zen/sinks/pymongo.py @@ -5,6 +5,7 @@ from aikido_zen.helpers.get_argument import get_argument import aikido_zen.vulnerabilities as vulns from . import patch_function, on_import, before +from ..helpers.register_call import register_call @before @@ -14,9 +15,12 @@ def _func_filter_first(func, instance, args, kwargs): if not nosql_filter: return + operation = f"pymongo.collection.Collection.{func.__name__}" + register_call(operation, "nosql_op") + vulns.run_vulnerability_scan( kind="nosql_injection", - op=f"pymongo.collection.Collection.{func.__name__}", + op=operation, args=(nosql_filter,), ) @@ -28,9 +32,12 @@ def _func_filter_second(func, instance, args, kwargs): if not nosql_filter: return + operation = f"pymongo.collection.Collection.{func.__name__}" + register_call(operation, "nosql_op") + vulns.run_vulnerability_scan( kind="nosql_injection", - op=f"pymongo.collection.Collection.{func.__name__}", + op=operation, args=(nosql_filter,), ) @@ -42,9 +49,12 @@ def _func_pipeline(func, instance, args, kwargs): if not nosql_pipeline: return + operation = f"pymongo.collection.Collection.{func.__name__}" + register_call(operation, "nosql_op") + vulns.run_vulnerability_scan( kind="nosql_injection", - op=f"pymongo.collection.Collection.{func.__name__}", + op=operation, args=(nosql_pipeline,), ) @@ -53,12 +63,15 @@ def _func_pipeline(func, instance, args, kwargs): def _bulk_write(func, instance, args, kwargs): requests = get_argument(args, kwargs, 0, "requests") + operation = "pymongo.collection.Collection.bulk_write" + register_call(operation, "nosql_op") + # Filter requests that contain "_filter" requests_with_filter = [req for req in requests if hasattr(req, "_filter")] for request in requests_with_filter: vulns.run_vulnerability_scan( kind="nosql_injection", - op="pymongo.collection.Collection.bulk_write", + op=operation, args=(request._filter,), ) From 21455a80cc98b029d06dbf4f9d3eb6f0819a11e6 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 10:28:02 +0200 Subject: [PATCH 07/24] add register_call to mysqlclient --- aikido_zen/sinks/mysqlclient.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aikido_zen/sinks/mysqlclient.py b/aikido_zen/sinks/mysqlclient.py index 86278a1f..c8a44565 100644 --- a/aikido_zen/sinks/mysqlclient.py +++ b/aikido_zen/sinks/mysqlclient.py @@ -4,6 +4,7 @@ from aikido_zen.helpers.get_argument import get_argument import aikido_zen.vulnerabilities as vulns +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import patch_function, on_import, before @@ -14,6 +15,7 @@ def _execute(func, instance, args, kwargs): # If query is type bytearray, it will be picked up by our wrapping of executemany return + register_call("MySQLdb.Cursor.execute", "sql_op") vulns.run_vulnerability_scan( kind="sql_injection", op="MySQLdb.Cursor.execute", args=(query, "mysql") ) @@ -23,6 +25,7 @@ def _execute(func, instance, args, kwargs): def _executemany(func, instance, args, kwargs): query = get_argument(args, kwargs, 0, "query") + register_call("MySQLdb.Cursor.executemany", "sql_op") vulns.run_vulnerability_scan( kind="sql_injection", op="MySQLdb.Cursor.executemany", args=(query, "mysql") ) From 99968fbbc47606a3e7df931d1cde6fe53080b58d Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:10:49 +0200 Subject: [PATCH 08/24] asyncpg add register_call --- aikido_zen/sinks/asyncpg.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aikido_zen/sinks/asyncpg.py b/aikido_zen/sinks/asyncpg.py index f81e5d70..b5a7a034 100644 --- a/aikido_zen/sinks/asyncpg.py +++ b/aikido_zen/sinks/asyncpg.py @@ -4,6 +4,7 @@ import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import patch_function, before, on_import @@ -12,6 +13,8 @@ def _execute(func, instance, args, kwargs): query = get_argument(args, kwargs, 0, "query") op = f"asyncpg.connection.Connection.{func.__name__}" + register_call(op, "sql_op") + vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=(query, "postgres")) From 2c91970999a3c377357dba8271c369537b12adbb Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:00 +0200 Subject: [PATCH 09/24] builtins add register_call --- aikido_zen/sinks/builtins.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/aikido_zen/sinks/builtins.py b/aikido_zen/sinks/builtins.py index c1dd54da..6b02dce2 100644 --- a/aikido_zen/sinks/builtins.py +++ b/aikido_zen/sinks/builtins.py @@ -5,6 +5,7 @@ from pathlib import PurePath import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import patch_function, on_import, before @@ -14,9 +15,10 @@ def _open(func, instance, args, kwargs): if not isinstance(filename, (str, bytes, PurePath)): return - vulns.run_vulnerability_scan( - kind="path_traversal", op="builtins.open", args=(filename,) - ) + op = "builtins.open" + register_call(op, "fs_op") + + vulns.run_vulnerability_scan(kind="path_traversal", op=op, args=(filename,)) @on_import("builtins") From 4f08157520620bfe6934628349f41e05c1bfb68f Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:12 +0200 Subject: [PATCH 10/24] add sink stats to clickhouse_driver --- aikido_zen/sinks/clickhouse_driver.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/aikido_zen/sinks/clickhouse_driver.py b/aikido_zen/sinks/clickhouse_driver.py index 53da3c83..d577acb7 100644 --- a/aikido_zen/sinks/clickhouse_driver.py +++ b/aikido_zen/sinks/clickhouse_driver.py @@ -1,14 +1,17 @@ from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import before, on_import, patch_function from aikido_zen.vulnerabilities import run_vulnerability_scan @before def _execute(func, instance, args, kwargs): - kind = "sql_injection" - op = "clickhouse_driver.Client.execute" query = get_argument(args, kwargs, 0, "query") - run_vulnerability_scan(kind, op, args=(query, "clickhouse")) + + op = "clickhouse_driver.Client.execute" + register_call(op, "sql_op") + + run_vulnerability_scan("sql_injection", op, args=(query, "clickhouse")) @on_import("clickhouse_driver", package="clickhouse_driver") From 9b0b0d7a02c84debf0afa22c8a98d8774ce27e5c Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:21 +0200 Subject: [PATCH 11/24] add sink stats to http client --- aikido_zen/sinks/http_client.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/aikido_zen/sinks/http_client.py b/aikido_zen/sinks/http_client.py index d8e495ed..d01b0094 100644 --- a/aikido_zen/sinks/http_client.py +++ b/aikido_zen/sinks/http_client.py @@ -3,6 +3,7 @@ """ from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import before, after, patch_function, on_import from aikido_zen.vulnerabilities.ssrf.handle_http_response import ( handle_http_response, @@ -22,6 +23,9 @@ def _putrequest(func, instance, args, kwargs): def _getresponse(func, instance, args, kwargs, return_value): path = getattr(instance, "_aikido_var_path") source_url = try_parse_url(f"http://{instance.host}:{instance.port}{path}") + + register_call("HTTPConnection.getresponse", "outgoing_http_op") + handle_http_response(http_response=return_value, source=source_url) From 30836f7d1e1ab0f8afe647b2603bc5d8436f1825 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:27 +0200 Subject: [PATCH 12/24] Add sink stats to io --- aikido_zen/sinks/io.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/aikido_zen/sinks/io.py b/aikido_zen/sinks/io.py index 48d54308..90f62997 100644 --- a/aikido_zen/sinks/io.py +++ b/aikido_zen/sinks/io.py @@ -4,6 +4,7 @@ import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import patch_function, before, on_import @@ -13,7 +14,10 @@ def _open(func, instance, args, kwargs): if not file: return - vulns.run_vulnerability_scan(kind="path_traversal", op="io.open", args=(file,)) + op = "io.open" + register_call(op, "fs_op") + + vulns.run_vulnerability_scan(kind="path_traversal", op=op, args=(file,)) @before @@ -22,7 +26,10 @@ def _open_code(func, instance, args, kwargs): if not path: return - vulns.run_vulnerability_scan(kind="path_traversal", op="io.open_code", args=(path,)) + op = "io.open_code" + register_call(op, "fs_op") + + vulns.run_vulnerability_scan(kind="path_traversal", op=op, args=(path,)) @on_import("io") From 93d97b70070a4574691302bfb5ae67866257bc35 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:37 +0200 Subject: [PATCH 13/24] Add sink stats to os_system --- aikido_zen/sinks/os_system.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/aikido_zen/sinks/os_system.py b/aikido_zen/sinks/os_system.py index 52d7887b..ea56f89e 100644 --- a/aikido_zen/sinks/os_system.py +++ b/aikido_zen/sinks/os_system.py @@ -4,6 +4,7 @@ import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import patch_function, before, on_import @@ -13,9 +14,10 @@ def _system(func, instance, args, kwargs): if not isinstance(command, str): return - vulns.run_vulnerability_scan( - kind="shell_injection", op="os.system", args=(command,) - ) + op = "os.system" + register_call(op, "exec_op") + + vulns.run_vulnerability_scan(kind="shell_injection", op=op, args=(command,)) @on_import("os") From 71cdf7550a5a774fb0b3f78d47138be8fce32bbc Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:43 +0200 Subject: [PATCH 14/24] Add sink stats to os --- aikido_zen/sinks/os.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aikido_zen/sinks/os.py b/aikido_zen/sinks/os.py index 8d2d9647..5755811b 100644 --- a/aikido_zen/sinks/os.py +++ b/aikido_zen/sinks/os.py @@ -4,6 +4,7 @@ from pathlib import PurePath import aikido_zen.vulnerabilities as vulns +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import before, patch_function, on_import @@ -17,6 +18,7 @@ def _os_patch(func, instance, args, kwargs): op = f"os.{func.__name__}" if func.__name__ in ("getsize", "join", "expanduser", "expandvars", "realpath"): op = f"os.path.{func.__name__}" + register_call(op, "fs_op") vulns.run_vulnerability_scan(kind="path_traversal", op=op, args=(path,)) From 9846ada729e043ec7eb4ad3f696a4b94a3ce9411 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:51 +0200 Subject: [PATCH 15/24] Add sink stats to psycopg --- aikido_zen/sinks/psycopg.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/aikido_zen/sinks/psycopg.py b/aikido_zen/sinks/psycopg.py index 45ae3fd1..fd2f436c 100644 --- a/aikido_zen/sinks/psycopg.py +++ b/aikido_zen/sinks/psycopg.py @@ -4,14 +4,19 @@ import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import patch_function, on_import, before @before def _copy(func, instance, args, kwargs): statement = get_argument(args, kwargs, 0, "statement") + + op = "psycopg.Cursor.copy" + register_call(op, "sql_op") + vulns.run_vulnerability_scan( - kind="sql_injection", op="psycopg.Cursor.copy", args=(statement, "postgres") + kind="sql_injection", op=op, args=(statement, "postgres") ) From 09a20c9d46d0b4c886c6144643543a40f6b21dcc Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:11:58 +0200 Subject: [PATCH 16/24] Add sink stats to psycopg2 --- aikido_zen/sinks/psycopg2.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/aikido_zen/sinks/psycopg2.py b/aikido_zen/sinks/psycopg2.py index 5b294552..c7e7e2ac 100644 --- a/aikido_zen/sinks/psycopg2.py +++ b/aikido_zen/sinks/psycopg2.py @@ -5,6 +5,7 @@ from aikido_zen.background_process.packages import is_package_compatible import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import on_import, before, patch_function, after @@ -35,7 +36,10 @@ def _connect(func, instance, _args, _kwargs, rv): @before def psycopg2_patch(func, instance, args, kwargs): query = get_argument(args, kwargs, 0, "query") + op = f"psycopg2.Connection.Cursor.{func.__name__}" + register_call(op, "sql_op") + vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=(query, "postgres")) From 8c799a8c0438723369ae724f372b4df8ec09decf Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:12:04 +0200 Subject: [PATCH 17/24] Add sink stats to shutil --- aikido_zen/sinks/shutil.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/aikido_zen/sinks/shutil.py b/aikido_zen/sinks/shutil.py index 0d8805d6..e743339e 100644 --- a/aikido_zen/sinks/shutil.py +++ b/aikido_zen/sinks/shutil.py @@ -4,6 +4,7 @@ import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import on_import, patch_function, before @@ -12,8 +13,10 @@ def _shutil_func(func, instance, args, kwargs): source = get_argument(args, kwargs, 0, "src") destination = get_argument(args, kwargs, 1, "dst") - kind = "path_traversal" op = f"shutil.{func.__name__}" + register_call(op, "fs_op") + + kind = "path_traversal" if isinstance(source, str): vulns.run_vulnerability_scan(kind, op, args=(source,)) if isinstance(destination, str): From 040e9cb27e15e4762d71215303c170f889728ab0 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:12:09 +0200 Subject: [PATCH 18/24] Add sink stats to socket --- aikido_zen/sinks/socket.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/aikido_zen/sinks/socket.py b/aikido_zen/sinks/socket.py index fc613db1..cf200da1 100644 --- a/aikido_zen/sinks/socket.py +++ b/aikido_zen/sinks/socket.py @@ -3,6 +3,7 @@ """ from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import on_import, patch_function, after from aikido_zen.vulnerabilities import run_vulnerability_scan @@ -11,8 +12,12 @@ def _getaddrinfo(func, instance, args, kwargs, return_value): host = get_argument(args, kwargs, 0, "host") port = get_argument(args, kwargs, 1, "port") + + op = "socket.getaddrinfo" + register_call(op, "outgoing_http_op") + arguments = (return_value, host, port) # return_value = dns response - run_vulnerability_scan(kind="ssrf", op="socket.getaddrinfo", args=arguments) + run_vulnerability_scan(kind="ssrf", op=op, args=arguments) @on_import("socket") From 63915c634b2cba8d8694cfff0c64ed54beb62a9b Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:12:18 +0200 Subject: [PATCH 19/24] Add sink stats to subprocess --- aikido_zen/sinks/subprocess.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/aikido_zen/sinks/subprocess.py b/aikido_zen/sinks/subprocess.py index 2e20fbfb..ca6169b8 100644 --- a/aikido_zen/sinks/subprocess.py +++ b/aikido_zen/sinks/subprocess.py @@ -4,6 +4,7 @@ import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import on_import, patch_function, before @@ -26,9 +27,13 @@ def _subprocess_init(func, instance, args, kwargs): command = shell_arguments if not command: return + + op = "subprocess.Popen" + register_call(op, "exec_op") + vulns.run_vulnerability_scan( kind="shell_injection", - op=f"subprocess.Popen", + op=op, args=(command,), ) From d53c18d96b3dfdb6efa5e5098c870acd1db307dd Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 11:16:20 +0200 Subject: [PATCH 20/24] Add test cases for new Operations calss --- .../storage/statistics/operations_test.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 aikido_zen/storage/statistics/operations_test.py diff --git a/aikido_zen/storage/statistics/operations_test.py b/aikido_zen/storage/statistics/operations_test.py new file mode 100644 index 00000000..f44f1453 --- /dev/null +++ b/aikido_zen/storage/statistics/operations_test.py @@ -0,0 +1,67 @@ +import pytest +from .operations import Operations + + +@pytest.fixture +def operations(): + return Operations() + + +def test_ensure_operation(operations): + operation = "test_op" + kind = "sql_op" + operations.ensure_operation(operation, kind) + assert operation in operations + assert operations[operation]["kind"] == kind + assert operations[operation]["total"] == 0 + assert operations[operation]["attacksDetected"]["total"] == 0 + assert operations[operation]["attacksDetected"]["blocked"] == 0 + + +def test_ensure_operation_unsupported_kind(operations): + operation = "test_op" + kind = "unsupported_kind" + with pytest.raises(Exception): + operations.ensure_operation(operation, kind) + + +def test_register_call(operations): + operation = "test_op" + kind = "sql_op" + operations.register_call(operation, kind) + assert operation in operations + assert operations[operation]["total"] == 1 + + +def test_on_detected_attack(operations): + operation = "test_op" + kind = "sql_op" + operations.ensure_operation(operation, kind) + operations.on_detected_attack(blocked=True, operation=operation) + assert operations[operation]["attacksDetected"]["total"] == 1 + assert operations[operation]["attacksDetected"]["blocked"] == 1 + + +def test_on_detected_attack_not_blocked(operations): + operation = "test_op" + kind = "sql_op" + operations.ensure_operation(operation, kind) + operations.on_detected_attack(blocked=False, operation=operation) + assert operations[operation]["attacksDetected"]["total"] == 1 + assert operations[operation]["attacksDetected"]["blocked"] == 0 + + +def test_on_detected_attack_unknown_operation(operations): + operation = "unknown_op" + operations.on_detected_attack(blocked=True, operation=operation) + assert operation not in operations + + +def test_register_call_and_on_detected_attack(operations): + operation = "test_op" + kind = "sql_op" + operations.register_call(operation, kind) + operations.on_detected_attack(blocked=True, operation=operation) + assert operations[operation]["total"] == 1 + assert operations[operation]["attacksDetected"]["total"] == 1 + assert operations[operation]["attacksDetected"]["blocked"] == 1 From ea220807adb95da512d72f51510a66c94b8f971e Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 12:41:37 +0200 Subject: [PATCH 21/24] Update statistics/__init__.py test cases for operations --- aikido_zen/storage/statistics/init_test.py | 121 +++++++++++++++++---- 1 file changed, 98 insertions(+), 23 deletions(-) diff --git a/aikido_zen/storage/statistics/init_test.py b/aikido_zen/storage/statistics/init_test.py index 3ea80366..5ca481d0 100644 --- a/aikido_zen/storage/statistics/init_test.py +++ b/aikido_zen/storage/statistics/init_test.py @@ -1,5 +1,10 @@ import pytest -from . import Statistics +from . import Statistics, Operations + + +@pytest.fixture +def stats(): + return Statistics() def test_initialization(monkeypatch): @@ -14,6 +19,7 @@ def test_initialization(monkeypatch): assert stats.attacks_detected == 0 assert stats.attacks_blocked == 0 assert stats.started_at == mock_time + assert isinstance(stats.operations, Operations) def test_clear(monkeypatch): @@ -27,12 +33,14 @@ def test_clear(monkeypatch): stats.total_hits = 10 stats.attacks_detected = 5 stats.attacks_blocked = 3 + stats.operations.register_call("test", "sql_op") stats.clear() assert stats.total_hits == 0 assert stats.attacks_detected == 0 assert stats.attacks_blocked == 0 assert stats.started_at == mock_time + assert stats.operations == {} def test_increment_total_hits(): @@ -41,13 +49,12 @@ def test_increment_total_hits(): assert stats.total_hits == 1 -def test_on_detected_attack(): - stats = Statistics() - stats.on_detected_attack(blocked=True) +def test_on_detected_attack(stats): + stats.on_detected_attack(blocked=True, operation="test_op") assert stats.attacks_detected == 1 assert stats.attacks_blocked == 1 - stats.on_detected_attack(blocked=False) + stats.on_detected_attack(blocked=False, operation="test_op") assert stats.attacks_detected == 2 assert stats.attacks_blocked == 1 @@ -61,6 +68,8 @@ def test_get_record(monkeypatch): stats = Statistics() stats.total_hits = 10 + stats.operations.register_call("test.test", "nosql_op") + stats.on_detected_attack(blocked=True, operation="test.test") stats.attacks_detected = 5 stats.attacks_blocked = 3 @@ -71,10 +80,20 @@ def test_get_record(monkeypatch): assert record["requests"]["aborted"] == 0 assert record["requests"]["attacksDetected"]["total"] == 5 assert record["requests"]["attacksDetected"]["blocked"] == 3 + assert record["operations"] == { + "test.test": { + "attacksDetected": {"blocked": 1, "total": 1}, + "kind": "nosql_op", + "total": 1, + } + } def test_import_from_record(): stats = Statistics() + stats.operations.register_call("test.test", "nosql_op") + stats.operations.register_call("test.test", "nosql_op") + stats.operations.register_call("test.test", "nosql_op") record = { "requests": { "total": 10, @@ -82,16 +101,39 @@ def test_import_from_record(): "total": 5, "blocked": 3, }, - } + }, + "operations": { + "test.test": { + "attacksDetected": {"blocked": 1, "total": 1}, + "kind": "nosql_op", + "total": 1, + }, + "test.test2": { + "kind": "sql_op", + "total": 5, + "attacksDetected": {"blocked": 5, "total": 200}, + }, + }, } stats.import_from_record(record) assert stats.total_hits == 10 assert stats.attacks_detected == 5 assert stats.attacks_blocked == 3 + assert stats.operations == { + "test.test": { + "attacksDetected": {"blocked": 1, "total": 1}, + "kind": "nosql_op", + "total": 4, + }, + "test.test2": { + "kind": "sql_op", + "total": 5, + "attacksDetected": {"blocked": 5, "total": 200}, + }, + } -def test_empty(): - stats = Statistics() +def test_empty(stats): assert stats.empty() == True stats.total_hits = 1 @@ -101,9 +143,12 @@ def test_empty(): stats.attacks_detected = 1 assert stats.empty() == False + stats.attacks_detected = 0 + stats.operations = {"test_op": {"total": 1}} + assert stats.empty() == False -def test_multiple_imports(): - stats = Statistics() + +def test_multiple_imports(stats): record1 = { "requests": { "total": 10, @@ -111,7 +156,14 @@ def test_multiple_imports(): "total": 5, "blocked": 3, }, - } + }, + "operations": { + "test_op1": { + "total": 1, + "kind": "fs_op", + "attacksDetected": {"total": 0, "blocked": 0}, + } + }, } record2 = { "requests": { @@ -120,50 +172,73 @@ def test_multiple_imports(): "total": 10, "blocked": 7, }, - } + }, + "operations": { + "test_op2": { + "total": 1, + "kind": "fs_op", + "attacksDetected": {"total": 0, "blocked": 0}, + } + }, } stats.import_from_record(record1) stats.import_from_record(record2) assert stats.total_hits == 30 assert stats.attacks_detected == 15 assert stats.attacks_blocked == 10 + assert stats.operations == { + "test_op1": { + "attacksDetected": {"blocked": 0, "total": 0}, + "kind": "fs_op", + "total": 1, + }, + "test_op2": { + "attacksDetected": {"blocked": 0, "total": 0}, + "kind": "fs_op", + "total": 1, + }, + } -def test_import_empty_record(): - stats = Statistics() +def test_import_empty_record(stats): record = {"requests": {}} stats.import_from_record(record) assert stats.total_hits == 0 assert stats.attacks_detected == 0 assert stats.attacks_blocked == 0 + assert stats.operations == {} -def test_import_partial_record(): - stats = Statistics() +def test_import_partial_record(stats): record = {"requests": {"total": 10}} stats.import_from_record(record) assert stats.total_hits == 10 assert stats.attacks_detected == 0 assert stats.attacks_blocked == 0 + assert stats.operations == {} -def test_increment_and_detect(): - stats = Statistics() +def test_increment_and_detect(stats): stats.increment_total_hits() - stats.on_detected_attack(blocked=True) + stats.on_detected_attack(blocked=True, operation="test_op") assert stats.total_hits == 1 assert stats.attacks_detected == 1 assert stats.attacks_blocked == 1 -def test_multiple_increments_and_detects(): - stats = Statistics() +def test_multiple_increments_and_detects(stats): + stats.operations.register_call("test_op", "sql_op") for _ in range(10): stats.increment_total_hits() for _ in range(5): - stats.on_detected_attack(blocked=True) + stats.on_detected_attack(blocked=True, operation="test_op") for _ in range(5): - stats.on_detected_attack(blocked=False) + stats.on_detected_attack(blocked=False, operation="test_op") assert stats.total_hits == 10 assert stats.attacks_detected == 10 assert stats.attacks_blocked == 5 + assert stats.operations.get("test_op") == { + "attacksDetected": {"blocked": 5, "total": 10}, + "kind": "sql_op", + "total": 1, + } From 388afc8f07d257673a9d9db09da9b6ee501104a2 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 12:42:02 +0200 Subject: [PATCH 22/24] Add modified update function to operations.py, in order to merge different counts --- aikido_zen/storage/statistics/operations.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/aikido_zen/storage/statistics/operations.py b/aikido_zen/storage/statistics/operations.py index 8baab062..07113708 100644 --- a/aikido_zen/storage/statistics/operations.py +++ b/aikido_zen/storage/statistics/operations.py @@ -29,3 +29,15 @@ def on_detected_attack(self, blocked, operation): self[operation]["attacksDetected"]["total"] += 1 if blocked: self[operation]["attacksDetected"]["blocked"] += 1 + + def update(self, m, /, **kwargs): + for operation in m.keys(): + self.ensure_operation(operation, kind=m[operation]["kind"]) + + self[operation]["total"] += m[operation]["total"] + + imported_attacks_total = m[operation]["attacksDetected"]["total"] + self[operation]["attacksDetected"]["total"] += imported_attacks_total + + imported_attacks_blocked = m[operation]["attacksDetected"]["blocked"] + self[operation]["attacksDetected"]["blocked"] += imported_attacks_blocked From e237bc750e3b300fc250840332f783750b42e0e9 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 13:15:39 +0200 Subject: [PATCH 23/24] Fix bug with thread_cache sync and update thread_cache test cases --- aikido_zen/storage/statistics/__init__.py | 2 +- aikido_zen/thread/thread_cache_test.py | 25 +++++++++++++++++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/aikido_zen/storage/statistics/__init__.py b/aikido_zen/storage/statistics/__init__.py index 9ee68973..bedf42ca 100644 --- a/aikido_zen/storage/statistics/__init__.py +++ b/aikido_zen/storage/statistics/__init__.py @@ -44,7 +44,7 @@ def get_record(self): "blocked": self.attacks_blocked, }, }, - "operations": self.operations, + "operations": dict(self.operations), } def import_from_record(self, record): diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index 08846ad9..cfb725bd 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -68,7 +68,7 @@ def test_reset(thread_cache: ThreadCache): thread_cache.config.bypassed_ips.add("192.168.1.1") thread_cache.config.blocked_uids.add("user123") thread_cache.stats.increment_total_hits() - thread_cache.stats.on_detected_attack(blocked=True) + thread_cache.stats.on_detected_attack(blocked=True, operation="test") thread_cache.reset() @@ -236,9 +236,11 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach # Setup initial state thread_cache.stats.increment_total_hits() thread_cache.stats.increment_total_hits() - thread_cache.stats.on_detected_attack(blocked=True) - thread_cache.stats.on_detected_attack(blocked=False) - thread_cache.stats.on_detected_attack(blocked=False) + thread_cache.stats.operations.register_call("op1", "sql_op") + thread_cache.stats.operations.register_call("op2", "sql_op") + thread_cache.stats.on_detected_attack(blocked=True, operation="op1") + thread_cache.stats.on_detected_attack(blocked=False, operation="op1") + thread_cache.stats.on_detected_attack(blocked=False, operation="op2") thread_cache.routes.initialize_route({"method": "GET", "route": "/test"}) thread_cache.routes.increment_route({"method": "GET", "route": "/test"}) @@ -269,6 +271,18 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach "aborted": 0, "attacksDetected": {"blocked": 1, "total": 3}, }, + "operations": { + "op1": { + "attacksDetected": {"blocked": 1, "total": 2}, + "kind": "sql_op", + "total": 1, + }, + "op2": { + "attacksDetected": {"blocked": 0, "total": 1}, + "kind": "sql_op", + "total": 1, + }, + }, }, "middleware_installed": False, "hostnames": [], @@ -313,6 +327,7 @@ def test_sync_data_for_users(mock_get_comms, thread_cache: ThreadCache): "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, }, + "operations": {}, }, "middleware_installed": False, "hostnames": [], @@ -362,6 +377,7 @@ def test_renew_called_with_empty_routes(mock_get_comms, thread_cache: ThreadCach "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, }, + "operations": {}, }, "middleware_installed": False, "hostnames": [], @@ -399,6 +415,7 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache "aborted": 0, "attacksDetected": {"total": 0, "blocked": 0}, }, + "operations": {}, }, "middleware_installed": False, "hostnames": [], From e96e4e8e9dcbb276d10b535cf1cfaf26d48290d9 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Thu, 5 Jun 2025 17:13:03 +0200 Subject: [PATCH 24/24] Add deserialize_op (xml/lxml) --- aikido_zen/sources/xml_sources/lxml.py | 5 +++++ aikido_zen/sources/xml_sources/xml.py | 5 +++++ aikido_zen/storage/statistics/operations.py | 9 ++++++++- 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/aikido_zen/sources/xml_sources/lxml.py b/aikido_zen/sources/xml_sources/lxml.py index 4558959c..c27ca44f 100644 --- a/aikido_zen/sources/xml_sources/lxml.py +++ b/aikido_zen/sources/xml_sources/lxml.py @@ -2,12 +2,15 @@ extract_data_from_xml_body, ) from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import on_import, after, patch_function @after def _fromstring(func, instance, args, kwargs, return_value): text = get_argument(args, kwargs, 0, "text") + register_call("lxml.etree.fromstring", "deserialize_op") + if text: extract_data_from_xml_body(user_input=text, root_element=return_value) @@ -15,6 +18,8 @@ def _fromstring(func, instance, args, kwargs, return_value): @after def _fromstringlist(func, instance, args, kwargs, return_value): strings = get_argument(args, kwargs, 0, "strings") + register_call("lxml.etree.fromstringlist", "deserialize_op") + for text in strings: extract_data_from_xml_body(user_input=text, root_element=return_value) diff --git a/aikido_zen/sources/xml_sources/xml.py b/aikido_zen/sources/xml_sources/xml.py index 6f26d5ec..26930904 100644 --- a/aikido_zen/sources/xml_sources/xml.py +++ b/aikido_zen/sources/xml_sources/xml.py @@ -2,18 +2,23 @@ extract_data_from_xml_body, ) from aikido_zen.helpers.get_argument import get_argument +from aikido_zen.helpers.register_call import register_call from aikido_zen.sinks import on_import, patch_function, after @after def _fromstring(func, instance, args, kwargs, return_value): text = get_argument(args, kwargs, 0, "text") + register_call("xml.etree.ElementTree.fromstring", "deserialize_op") + extract_data_from_xml_body(user_input=text, root_element=return_value) @after def _fromstringlist(func, instance, args, kwargs, return_value): strings = get_argument(args, kwargs, 0, "sequence") + register_call("xml.etree.ElementTree.fromstringlist", "deserialize_op") + for text in strings: extract_data_from_xml_body(user_input=text, root_element=return_value) diff --git a/aikido_zen/storage/statistics/operations.py b/aikido_zen/storage/statistics/operations.py index 07113708..d1d15eea 100644 --- a/aikido_zen/storage/statistics/operations.py +++ b/aikido_zen/storage/statistics/operations.py @@ -1,4 +1,11 @@ -SUPPORTED_KINDS = ["sql_op", "nosql_op", "outgoing_http_op", "fs_op", "exec_op"] +SUPPORTED_KINDS = [ + "sql_op", + "nosql_op", + "outgoing_http_op", + "fs_op", + "exec_op", + "deserialize_op", +] class Operations(dict):