Add sink statistics

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)

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"))
 
 

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")

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")

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)
 
 

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")

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")
     )

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,))
 

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")

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")
     )
 
 

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"))
 
 

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,),
         )
 

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")
     )

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):

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")

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,),
     )