Add datadog compat

Makefile

@@ -24,15 +24,15 @@ dev_install:
 
 .PHONY: test
 test:
-	poetry run pytest aikido_zen/
+	poetry run ddtrace-run pytest aikido_zen/
 
 .PHONY: end2end
 end2end:
-	poetry run pytest end2end/ 	
+	poetry run pytest end2end/
 
 .PHONY: cov
 cov:
-	poetry run pytest aikido_zen/ --cov=aikido_zen --cov-report=xml
+	poetry run ddtrace-run pytest aikido_zen/ --cov=aikido_zen --cov-report=xml
 
 .PHONY: benchmark
 benchmark:

aikido_zen/sinks/builtins.py

@@ -3,11 +3,24 @@
 """
 
 from pathlib import PurePath
-import copy
 import aikido_zen.importhook as importhook
 import aikido_zen.vulnerabilities as vulns
 
 
+def aikido_open_decorator(func):
+    """Decorator for open(...)"""
+
+    def wrapper(*args, **kwargs):
+        #  args[0] is thefunc_name filename
+        if len(args) > 0 and isinstance(args[0], (str, bytes, PurePath)):
+            vulns.run_vulnerability_scan(
+                kind="path_traversal", op="builtins.open", args=(args[0],)
+            )
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 @importhook.on_import("builtins")
 def on_builtins_import(builtins):
     """
@@ -17,17 +30,7 @@ def on_builtins_import(builtins):
     """
     modified_builtins = importhook.copy_module(builtins)
 
-    former_open = copy.deepcopy(builtins.open)
-
-    def aikido_new_open(*args, **kwargs):
-        #  args[0] is the filename
-        if len(args) > 0 and isinstance(args[0], (str, bytes, PurePath)):
-            vulns.run_vulnerability_scan(
-                kind="path_traversal", op="builtins.open", args=(args[0],)
-            )
-        return former_open(*args, **kwargs)
-
     # pylint: disable=no-member
-    setattr(builtins, "open", aikido_new_open)
-    setattr(modified_builtins, "open", aikido_new_open)
+    setattr(builtins, "open", aikido_open_decorator(builtins.open))
+    setattr(modified_builtins, "open", aikido_open_decorator(builtins.open))
     return modified_builtins

aikido_zen/sinks/os_system.py

@@ -2,7 +2,6 @@
 Sink module for `os`, wrapping os.system
 """
 
-import copy
 import aikido_zen.importhook as importhook
 import aikido_zen.vulnerabilities as vulns
 
@@ -19,18 +18,12 @@ def on_os_import(os):
     """
     modified_os = importhook.copy_module(os)
 
-    former_system_func = copy.deepcopy(os.system)
-
-    def aikido_new_system(
-        command, *args, former_system_func=former_system_func, **kwargs
-    ):
+    def aikido_new_system(command, *args, **kwargs):
         if isinstance(command, str):
             vulns.run_vulnerability_scan(
                 kind="shell_injection", op="os.system", args=(command,)
             )
-        return former_system_func(command, *args, **kwargs)
+        return os.system(command, *args, **kwargs)
 
-    setattr(os, "system", aikido_new_system)
     setattr(modified_os, "system", aikido_new_system)
-
     return modified_os

aikido_zen/sinks/psycopg2.py

@@ -25,7 +25,7 @@ def execute(self, *args, **kwargs):
             )
             if former_cursor_factory and hasattr(former_cursor_factory, "execute"):
                 return former_cursor_factory.execute(self, *args, **kwargs)
-            return super().execute(*args, **kwargs)
+            return ext.cursor.execute(self, *args, **kwargs)
 
         def executemany(self, *args, **kwargs):
             """Aikido's wrapped executemany function"""
@@ -37,7 +37,7 @@ def executemany(self, *args, **kwargs):
             )
             if former_cursor_factory and hasattr(former_cursor_factory, "executemany"):
                 return former_cursor_factory.executemany(self, *args, **kwargs)
-            return super().executemany(*args, **kwargs)
+            return ext.cursor.executemany(self, *args, **kwargs)
 
     return AikidoWrappedCursor
 

aikido_zen/sinks/shutil.py

@@ -18,24 +18,22 @@
 # shutil.copy2(src, dst, *, **) => builtins.open
 
 
-def generate_aikido_function(op, former_func):
+def generate_aikido_function(aik_op, func):
     """
     Returns a generated aikido function given an operation
     and the previous function
     """
 
-    def aikido_new_func(
-        src, dst, *args, aik_op=op, aik_former_func=former_func, **kwargs
-    ):
+    def wrapper(src, dst, *args, **kwargs):
         kind = "path_traversal"
         op = f"shutil.{aik_op}"
         if src:
             vulns.run_vulnerability_scan(kind, op, args=(src,))
         if dst:
             vulns.run_vulnerability_scan(kind, op, args=(dst,))
-        return former_func(src, dst, *args, **kwargs)
+        return func(src, dst, *args, **kwargs)
 
-    return aikido_new_func
+    return wrapper
 
 
 @importhook.on_import("shutil")
@@ -48,9 +46,7 @@ def on_shutil_import(shutil):
     modified_shutil = importhook.copy_module(shutil)
     for op in SHUTIL_SRC_DST_FUNCTIONS:
         # Wrap shutil. functions
-        former_func = copy.deepcopy(getattr(shutil, op))
-        aikido_new_func = generate_aikido_function(op, former_func)
-        setattr(shutil, op, aikido_new_func)
+        aikido_new_func = generate_aikido_function(op, getattr(shutil, op))
         setattr(modified_shutil, op, aikido_new_func)
 
     return modified_shutil

aikido_zen/sinks/tests/asyncpg_test.py

@@ -107,7 +107,7 @@ async def test_conn_execute_parameterized(database_conn):
             args = call[1]["args"]
             if args[0] == query:
                 counter += 1
-        assert counter == 2
+        assert counter == 2 * 2
 
         await conn.close()
 

aikido_zen/sinks/tests/motor_test.py

@@ -110,7 +110,6 @@ async def test_find_one(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.find"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -127,7 +126,6 @@ async def test_count_documents(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.count_documents"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -144,7 +142,6 @@ async def test_find_one_and_delete(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.find_one_and_delete"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -159,11 +156,10 @@ async def test_find_one_and_replace(db):
             filter=_filter, replacement={"dog_name": "test2"}
         )
 
-        called_with = mock_run_vulnerability_scan.call_args[1]
+        called_with = mock_run_vulnerability_scan.call_args_list[0][1]
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.find_one_and_replace"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -180,7 +176,6 @@ async def test_find_one_and_update(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.find_one_and_update"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -208,7 +203,6 @@ async def test_find_not_empty(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.find"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -225,7 +219,6 @@ async def test_find_raw_batches(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.find_raw_batches"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -242,7 +235,6 @@ async def test_distinct(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.distinct"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -259,7 +251,6 @@ async def test_distinct_kwargs(db):
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.distinct"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -297,7 +288,6 @@ async def test_aggregate(db):
         assert called_with["args"][0] == pipeline
         assert called_with["op"] == "pymongo.collection.Collection.aggregate"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -323,7 +313,6 @@ async def test_aggregate_key(db):
         assert called_with["args"][0] == pipeline
         assert called_with["op"] == "pymongo.collection.Collection.aggregate"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -345,13 +334,12 @@ async def test_aggregate_raw_batches_key(db):
         async for doc in dogs.aggregate_raw_batches(pipeline=pipeline):
             pass
 
-        called_with = mock_run_vulnerability_scan.call_args[1]
+        called_with = mock_run_vulnerability_scan.call_args_list[0][1]
         assert called_with["args"][0] == pipeline
         assert (
             called_with["op"] == "pymongo.collection.Collection.aggregate_raw_batches"
         )
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -379,7 +367,6 @@ async def test_aggregate_raw_batches(db):
             called_with["op"] == "pymongo.collection.Collection.aggregate_raw_batches"
         )
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -406,7 +393,6 @@ async def test_watch(db):
         assert called_with["args"][0] == pipeline
         assert called_with["op"] == "pymongo.collection.Collection.watch"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -433,7 +419,6 @@ async def test_watch_key(db):
         assert called_with["args"][0] == pipeline
         assert called_with["op"] == "pymongo.collection.Collection.watch"
         assert called_with["kind"] == "nosql_injection"
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 @pytest.mark.asyncio

aikido_zen/sinks/tests/mysqlclient_test.py

@@ -27,13 +27,13 @@ def test_cursor_execute(database_conn):
         assert called_with_args[0] == query
         assert isinstance(called_with_args[1], MySQL)
 
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
         cursor.fetchall()
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
 
         cursor.close()
         database_conn.close()
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
 
 
 def test_cursor_execute_no_args(database_conn):
@@ -53,13 +53,11 @@ def test_cursor_execute_no_args(database_conn):
         )
         assert isinstance(called_with_args[1], MySQL)
 
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
         cursor.fetchall()
-        mock_run_vulnerability_scan.assert_called_once()
 
         cursor.close()
         database_conn.close()
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 def test_cursor_executemany(database_conn):
@@ -80,13 +78,11 @@ def test_cursor_executemany(database_conn):
             == "INSERT INTO dogs (dog_name, isAdmin) VALUES (%s, %s)"
         )
         assert isinstance(called_with_args[1], MySQL)
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
         database_conn.commit()
-        mock_run_vulnerability_scan.assert_called_once()
 
         cursor.close()
         database_conn.close()
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 def test_cursor_execute_with_fstring(database_conn):
@@ -106,10 +102,8 @@ def test_cursor_execute_with_fstring(database_conn):
             called_with_args[0] == "INSERT INTO dogs (dog_name, isAdmin) VALUES (%s, 1)"
         )
         assert isinstance(called_with_args[1], MySQL)
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
         database_conn.commit()
-        mock_run_vulnerability_scan.assert_called_once()
 
         cursor.close()
         database_conn.close()
-        mock_run_vulnerability_scan.assert_called_once()

aikido_zen/sinks/tests/psycopg2_test.py

@@ -29,11 +29,10 @@ def test_cursor_execute(database_conn):
         called_with_args = mock_run_vulnerability_scan.call_args[1]["args"]
         assert called_with_args[0] == query
         assert isinstance(called_with_args[1], Postgres)
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
 
         cursor.fetchall()
         database_conn.close()
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 def test_cursor_execute_parameterized(database_conn):
@@ -48,11 +47,10 @@ def test_cursor_execute_parameterized(database_conn):
         called_with_args = mock_run_vulnerability_scan.call_args[1]["args"]
         assert called_with_args[0] == query
         assert isinstance(called_with_args[1], Postgres)
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
 
         database_conn.commit()
         database_conn.close()
-        mock_run_vulnerability_scan.assert_called_once()
 
 
 def test_cursor_executemany(database_conn):
@@ -73,9 +71,8 @@ def test_cursor_executemany(database_conn):
             == "INSERT INTO dogs (dog_name, isadmin) VALUES (%s, %s)"
         )
         assert isinstance(called_with_args[1], Postgres)
-        mock_run_vulnerability_scan.assert_called_once()
+        mock_run_vulnerability_scan.assert_called()
 
         database_conn.commit()
         cursor.close()
         database_conn.close()
-        mock_run_vulnerability_scan.assert_called_once()