diff --git a/aikido_zen/sinks/asyncpg.py b/aikido_zen/sinks/asyncpg.py index b5a7a0347..b956131ae 100644 --- a/aikido_zen/sinks/asyncpg.py +++ b/aikido_zen/sinks/asyncpg.py @@ -18,6 +18,18 @@ def _execute(func, instance, args, kwargs): vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=(query, "postgres")) +@before +def _executemany(func, instance, args, kwargs): + command = get_argument(args, kwargs, 0, "command") + + op = f"asyncpg.connection.Connection.{func.__name__}" + register_call(op, "sql_op") + + vulns.run_vulnerability_scan( + kind="sql_injection", op=op, args=(command, "postgres") + ) + + @on_import("asyncpg.connection", "asyncpg", version_requirement="0.27.0") def patch(m): """ @@ -28,5 +40,5 @@ def patch(m): src: https://github.com/MagicStack/asyncpg/blob/85d7eed40637e7cad73a44ed2439ffeb2a8dc1c2/asyncpg/connection.py#L43 """ patch_function(m, "Connection.execute", _execute) - patch_function(m, "Connection.executemany", _execute) + patch_function(m, "Connection.executemany", _executemany) patch_function(m, "Connection._execute", _execute) diff --git a/aikido_zen/sinks/tests/asyncpg_test.py b/aikido_zen/sinks/tests/asyncpg_test.py index 706a665ac..e0d4f3e08 100644 --- a/aikido_zen/sinks/tests/asyncpg_test.py +++ b/aikido_zen/sinks/tests/asyncpg_test.py @@ -13,7 +13,7 @@ async def database_conn(): import asyncpg conn = await asyncpg.connect( - host="127.0.0.1", user="user", password="password", database="db" + host="127.0.0.1", user="user", ****="****", database="db" ) return conn @@ -154,3 +154,47 @@ async def test_conn_cursor(database_conn): assert called_with_args[1] == "postgres" await conn.close() + + +@pytest.mark.asyncio +async def test_conn_executemany_positional(database_conn): + reset_comms() + with patch( + "aikido_zen.vulnerabilities.run_vulnerability_scan" + ) as mock_run_vulnerability_scan: + conn = await database_conn + command = "INSERT INTO dogs (dog_name, isadmin) VALUES ($1, $2)" + data = [("Dog1", False), ("Dog2", True)] + await conn.executemany(command, data) + + called_with_args = mock_run_vulnerability_scan.call_args[1]["args"] + called_with_op = mock_run_vulnerability_scan.call_args[1]["op"] + called_with_kind = mock_run_vulnerability_scan.call_args[1]["kind"] + assert called_with_args[0] == command + assert called_with_args[1] == "postgres" + assert called_with_op == "asyncpg.connection.Connection.executemany" + assert called_with_kind == "sql_injection" + + await conn.close() + + +@pytest.mark.asyncio +async def test_conn_executemany_keyword(database_conn): + reset_comms() + with patch( + "aikido_zen.vulnerabilities.run_vulnerability_scan" + ) as mock_run_vulnerability_scan: + conn = await database_conn + command = "INSERT INTO dogs (dog_name, isadmin) VALUES ($1, $2)" + data = [("Dog1", False), ("Dog2", True)] + await conn.executemany(command=command, args=data) + + called_with_args = mock_run_vulnerability_scan.call_args[1]["args"] + called_with_op = mock_run_vulnerability_scan.call_args[1]["op"] + called_with_kind = mock_run_vulnerability_scan.call_args[1]["kind"] + assert called_with_args[0] == command + assert called_with_args[1] == "postgres" + assert called_with_op == "asyncpg.connection.Connection.executemany" + assert called_with_kind == "sql_injection" + + await conn.close()