From 2308f59f55ce03638d8eee2e7ada6813a7c0182c Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Mon, 2 Sep 2024 12:14:09 +0200 Subject: [PATCH 1/7] Install psycopg as dev package --- poetry.lock | 65 +++++++++++++++++++++++++++++++++++++++++++++++++- pyproject.toml | 1 + 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/poetry.lock b/poetry.lock index 4260bf5e..e29b43c4 100644 --- a/poetry.lock +++ b/poetry.lock @@ -82,6 +82,34 @@ async-timeout = {version = ">=4.0.3", markers = "python_version < \"3.12.0\""} docs = ["Sphinx (>=5.3.0,<5.4.0)", "sphinx-rtd-theme (>=1.2.2)", "sphinxcontrib-asyncio (>=0.3.0,<0.4.0)"] test = ["flake8 (>=6.1,<7.0)", "uvloop (>=0.15.3)"] +[[package]] +name = "backports-zoneinfo" +version = "0.2.1" +description = "Backport of the standard library zoneinfo module" +optional = false +python-versions = ">=3.6" +files = [ + {file = "backports.zoneinfo-0.2.1-cp36-cp36m-macosx_10_14_x86_64.whl", hash = "sha256:da6013fd84a690242c310d77ddb8441a559e9cb3d3d59ebac9aca1a57b2e18bc"}, + {file = "backports.zoneinfo-0.2.1-cp36-cp36m-manylinux1_i686.whl", hash = "sha256:89a48c0d158a3cc3f654da4c2de1ceba85263fafb861b98b59040a5086259722"}, + {file = "backports.zoneinfo-0.2.1-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:1c5742112073a563c81f786e77514969acb58649bcdf6cdf0b4ed31a348d4546"}, + {file = "backports.zoneinfo-0.2.1-cp36-cp36m-win32.whl", hash = "sha256:e8236383a20872c0cdf5a62b554b27538db7fa1bbec52429d8d106effbaeca08"}, + {file = "backports.zoneinfo-0.2.1-cp36-cp36m-win_amd64.whl", hash = "sha256:8439c030a11780786a2002261569bdf362264f605dfa4d65090b64b05c9f79a7"}, + {file = "backports.zoneinfo-0.2.1-cp37-cp37m-macosx_10_14_x86_64.whl", hash = "sha256:f04e857b59d9d1ccc39ce2da1021d196e47234873820cbeaad210724b1ee28ac"}, + {file = "backports.zoneinfo-0.2.1-cp37-cp37m-manylinux1_i686.whl", hash = "sha256:17746bd546106fa389c51dbea67c8b7c8f0d14b5526a579ca6ccf5ed72c526cf"}, + {file = "backports.zoneinfo-0.2.1-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:5c144945a7752ca544b4b78c8c41544cdfaf9786f25fe5ffb10e838e19a27570"}, + {file = "backports.zoneinfo-0.2.1-cp37-cp37m-win32.whl", hash = "sha256:e55b384612d93be96506932a786bbcde5a2db7a9e6a4bb4bffe8b733f5b9036b"}, + {file = "backports.zoneinfo-0.2.1-cp37-cp37m-win_amd64.whl", hash = "sha256:a76b38c52400b762e48131494ba26be363491ac4f9a04c1b7e92483d169f6582"}, + {file = "backports.zoneinfo-0.2.1-cp38-cp38-macosx_10_14_x86_64.whl", hash = "sha256:8961c0f32cd0336fb8e8ead11a1f8cd99ec07145ec2931122faaac1c8f7fd987"}, + {file = "backports.zoneinfo-0.2.1-cp38-cp38-manylinux1_i686.whl", hash = "sha256:e81b76cace8eda1fca50e345242ba977f9be6ae3945af8d46326d776b4cf78d1"}, + {file = "backports.zoneinfo-0.2.1-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:7b0a64cda4145548fed9efc10322770f929b944ce5cee6c0dfe0c87bf4c0c8c9"}, + {file = "backports.zoneinfo-0.2.1-cp38-cp38-win32.whl", hash = "sha256:1b13e654a55cd45672cb54ed12148cd33628f672548f373963b0bff67b217328"}, + {file = "backports.zoneinfo-0.2.1-cp38-cp38-win_amd64.whl", hash = "sha256:4a0f800587060bf8880f954dbef70de6c11bbe59c673c3d818921f042f9954a6"}, + {file = "backports.zoneinfo-0.2.1.tar.gz", hash = "sha256:fadbfe37f74051d024037f223b8e001611eac868b5c5b06144ef4d8b799862f2"}, +] + +[package.extras] +tzdata = ["tzdata"] + [[package]] name = "black" version = "24.8.0" @@ -519,6 +547,30 @@ files = [ dev = ["pre-commit", "tox"] testing = ["pytest", "pytest-benchmark"] +[[package]] +name = "psycopg" +version = "3.2.1" +description = "PostgreSQL database adapter for Python" +optional = false +python-versions = ">=3.8" +files = [ + {file = "psycopg-3.2.1-py3-none-any.whl", hash = "sha256:ece385fb413a37db332f97c49208b36cf030ff02b199d7635ed2fbd378724175"}, + {file = "psycopg-3.2.1.tar.gz", hash = "sha256:dc8da6dc8729dacacda3cc2f17d2c9397a70a66cf0d2b69c91065d60d5f00cb7"}, +] + +[package.dependencies] +"backports.zoneinfo" = {version = ">=0.2.0", markers = "python_version < \"3.9\""} +typing-extensions = ">=4.4" +tzdata = {version = "*", markers = "sys_platform == \"win32\""} + +[package.extras] +binary = ["psycopg-binary (==3.2.1)"] +c = ["psycopg-c (==3.2.1)"] +dev = ["ast-comments (>=1.1.2)", "black (>=24.1.0)", "codespell (>=2.2)", "dnspython (>=2.1)", "flake8 (>=4.0)", "mypy (>=1.6)", "types-setuptools (>=57.4)", "wheel (>=0.37)"] +docs = ["Sphinx (>=5.0)", "furo (==2022.6.21)", "sphinx-autobuild (>=2021.3.14)", "sphinx-autodoc-typehints (>=1.12)"] +pool = ["psycopg-pool"] +test = ["anyio (>=4.0)", "mypy (>=1.6)", "pproxy (>=2.7)", "pytest (>=6.2.5)", "pytest-cov (>=3.0)", "pytest-randomly (>=3.5)"] + [[package]] name = "psycopg2-binary" version = "2.9.9" @@ -876,6 +928,17 @@ files = [ {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"}, ] +[[package]] +name = "tzdata" +version = "2024.1" +description = "Provider of IANA time zone data" +optional = false +python-versions = ">=2" +files = [ + {file = "tzdata-2024.1-py2.py3-none-any.whl", hash = "sha256:9068bc196136463f5245e51efda838afa15aaeca9903f49050dfa2679db4d252"}, + {file = "tzdata-2024.1.tar.gz", hash = "sha256:2674120f8d891909751c38abcdfd386ac0a5a1127954fbc332af6b5ceae07efd"}, +] + [[package]] name = "urllib3" version = "2.2.2" @@ -896,4 +959,4 @@ zstd = ["zstandard (>=0.18.0)"] [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "2f3af0dd01770c243706bb1ffb25aaff5c1fb3e702e1f1dcfba8ef87d1b8e021" +content-hash = "11a7e4834f9f496b5d7dba9b56a7c48d238db04647f8ac64abfb6bc397d0dc86" diff --git a/pyproject.toml b/pyproject.toml index 5cae60d0..75295567 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ pymysql = "^1.1.1" psycopg2-binary = "^2.9.9" pytest-asyncio = "^0.24.0" asyncpg = "^0.29.0" +psycopg = "^3.2.1" [build-system] requires = ["poetry-core>=1.0.0"] From 10f3edcc7ca3435ecb228eca8f61eead72847557 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Mon, 2 Sep 2024 12:19:18 +0200 Subject: [PATCH 2/7] Add psycopg sink --- aikido_firewall/__init__.py | 1 + aikido_firewall/sinks/psycopg.py | 49 ++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) create mode 100644 aikido_firewall/sinks/psycopg.py diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 49a4b205..688521ef 100644 --- a/aikido_firewall/__init__.py +++ b/aikido_firewall/__init__.py @@ -51,6 +51,7 @@ def protect(mode="daemon"): import aikido_firewall.sinks.mysqlclient import aikido_firewall.sinks.pymongo import aikido_firewall.sinks.psycopg2 + import aikido_firewall.sinks.psycopg import aikido_firewall.sinks.asyncpg import aikido_firewall.sinks.builtins import aikido_firewall.sinks.os diff --git a/aikido_firewall/sinks/psycopg.py b/aikido_firewall/sinks/psycopg.py new file mode 100644 index 00000000..d9ec408c --- /dev/null +++ b/aikido_firewall/sinks/psycopg.py @@ -0,0 +1,49 @@ +""" +Sink module for `psycopg` +""" + +import copy +import importhook +from aikido_firewall.vulnerabilities.sql_injection.dialects import Postgres +from aikido_firewall.background_process.packages import add_wrapped_package +import aikido_firewall.vulnerabilities as vulns + + +@importhook.on_import("psycopg.cursor") +def on_psycopg_import(psycopg): + """ + Hook 'n wrap on `psycopg.connect` function, we modify the cursor_factory + of the result of this connect function. + """ + modified_psycopg = importhook.copy_module(psycopg) + former_copy_funtcion = copy.deepcopy(psycopg.Cursor.copy) + former_execute_function = copy.deepcopy(psycopg.Cursor.execute) + former_executemany_function = copy.deepcopy(psycopg.Cursor.executemany) + + def aikido_copy(self, statement, params=None, *args, **kwargs): + sql = statement + vulns.run_vulnerability_scan( + kind="sql_injection", op="psycopg.Cursor.copy", args=(sql, Postgres()) + ) + return former_copy_funtcion(self, statement, params, *args, **kwargs) + + def aikido_execute(self, query, params=None, *args, **kwargs): + sql = query + vulns.run_vulnerability_scan( + kind="sql_injection", op="psycopg.Cursor.execute", args=(sql, Postgres()) + ) + return former_execute_function(self, query, params, *args, **kwargs) + + def aikido_executemany(self, query, params_seq, *args, **kwargs): + args = (query, Postgres()) + op = "psycopg.Cursor.executemany" + vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=args) + return former_executemany_function(self, query, params_seq, *args, **kwargs) + + setattr(psycopg.Cursor, "copy", aikido_copy) # pylint: disable=no-member + setattr(psycopg.Cursor, "execute", aikido_execute) # pylint: disable=no-member + # pylint: disable=no-member + setattr(psycopg.Cursor, "executemany", aikido_executemany) + + add_wrapped_package("psycopg") + return modified_psycopg From 57d8960b9459c31464df06a14bf78c5d80877582 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Mon, 2 Sep 2024 12:19:52 +0200 Subject: [PATCH 3/7] Add simple sink test for psycopg --- aikido_firewall/sinks/tests/psycopg_test.py | 36 +++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 aikido_firewall/sinks/tests/psycopg_test.py diff --git a/aikido_firewall/sinks/tests/psycopg_test.py b/aikido_firewall/sinks/tests/psycopg_test.py new file mode 100644 index 00000000..0bd51ee3 --- /dev/null +++ b/aikido_firewall/sinks/tests/psycopg_test.py @@ -0,0 +1,36 @@ +import pytest +from unittest.mock import patch +import aikido_firewall.sinks.psycopg +from aikido_firewall.background_process.comms import reset_comms +from aikido_firewall.vulnerabilities.sql_injection.dialects import Postgres + +kind = "sql_injection" +op = "pymysql.connections.query" + + +@pytest.fixture +def database_conn(): + import psycopg + + return psycopg.connect( + host="127.0.0.1", user="user", password="password", dbname="db" + ) + + +def test_cursor_execute(database_conn): + reset_comms() + with patch( + "aikido_firewall.vulnerabilities.run_vulnerability_scan" + ) as mock_run_vulnerability_scan: + cursor = database_conn.cursor() + query = "SELECT * FROM dogs" + cursor.execute(query) + + 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() + + cursor.fetchall() + database_conn.close() + mock_run_vulnerability_scan.assert_called_once() From da6e24075f494a095be6457e887f0de6f9b114f2 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Mon, 2 Sep 2024 12:22:38 +0200 Subject: [PATCH 4/7] Update psycopg tests to check for operation and kind --- aikido_firewall/sinks/tests/psycopg_test.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/aikido_firewall/sinks/tests/psycopg_test.py b/aikido_firewall/sinks/tests/psycopg_test.py index 0bd51ee3..2b4bf097 100644 --- a/aikido_firewall/sinks/tests/psycopg_test.py +++ b/aikido_firewall/sinks/tests/psycopg_test.py @@ -4,9 +4,6 @@ from aikido_firewall.background_process.comms import reset_comms from aikido_firewall.vulnerabilities.sql_injection.dialects import Postgres -kind = "sql_injection" -op = "pymysql.connections.query" - @pytest.fixture def database_conn(): @@ -26,9 +23,11 @@ def test_cursor_execute(database_conn): query = "SELECT * FROM dogs" cursor.execute(query) - called_with_args = mock_run_vulnerability_scan.call_args[1]["args"] - assert called_with_args[0] == query - assert isinstance(called_with_args[1], Postgres) + called_with = mock_run_vulnerability_scan.call_args[1] + assert called_with["args"][0] == query + assert isinstance(called_with["args"][1], Postgres) + assert called_with["op"] == "psycopg.Cursor.execute" + assert called_with["kind"] == "sql_injection" mock_run_vulnerability_scan.assert_called_once() cursor.fetchall() From a003311fa4f083c2cac1ec0bcd658e17fa15f572 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Mon, 2 Sep 2024 12:30:39 +0200 Subject: [PATCH 5/7] Add a lot more psycopg test cases --- aikido_firewall/sinks/tests/psycopg_test.py | 67 +++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/aikido_firewall/sinks/tests/psycopg_test.py b/aikido_firewall/sinks/tests/psycopg_test.py index 2b4bf097..358ca1f2 100644 --- a/aikido_firewall/sinks/tests/psycopg_test.py +++ b/aikido_firewall/sinks/tests/psycopg_test.py @@ -33,3 +33,70 @@ def test_cursor_execute(database_conn): cursor.fetchall() database_conn.close() mock_run_vulnerability_scan.assert_called_once() + + +def test_cursor_execute_parameterized(database_conn): + reset_comms() + cursor = database_conn.cursor() + query = "SELECT * FROM dogs WHERE dog_name = %s" + params = ("Fido",) + with patch( + "aikido_firewall.vulnerabilities.run_vulnerability_scan" + ) as mock_run_vulnerability_scan: + cursor.execute(query, params) + + called_with = mock_run_vulnerability_scan.call_args[1] + assert called_with["args"][0] == "SELECT * FROM dogs WHERE dog_name = %s" + assert isinstance(called_with["args"][1], Postgres) + assert called_with["op"] == "psycopg.Cursor.execute" + assert called_with["kind"] == "sql_injection" + mock_run_vulnerability_scan.assert_called_once() + + cursor.fetchall() + database_conn.close() + + +def test_cursor_executemany(database_conn): + reset_comms() + cursor = database_conn.cursor() + query = "INSERT INTO dogs (dog_name, isadmin) VALUES (%s, %s)" + params = [("doggo1", False), ("Rex", False), ("Buddy", True)] + + with patch( + "aikido_firewall.vulnerabilities.run_vulnerability_scan" + ) as mock_run_vulnerability_scan: + cursor.executemany(query, params) + + # Check the last call to run_vulnerability_scan + called_with = mock_run_vulnerability_scan.call_args[1] + assert ( + called_with["args"][0] + == "INSERT INTO dogs (dog_name, isadmin) VALUES (%s, %s)" + ) + assert isinstance(called_with["args"][1], Postgres) + assert called_with["op"] == "psycopg.Cursor.executemany" + assert called_with["kind"] == "sql_injection" + mock_run_vulnerability_scan.assert_called() + + database_conn.commit() + database_conn.close() + + +def test_cursor_copy(database_conn): + reset_comms() + cursor = database_conn.cursor() + query = "COPY dogs FROM STDIN" + + with patch( + "aikido_firewall.vulnerabilities.run_vulnerability_scan" + ) as mock_run_vulnerability_scan: + cursor.copy(query) + + called_with = mock_run_vulnerability_scan.call_args[1] + assert called_with["args"][0] == query + assert isinstance(called_with["args"][1], Postgres) + assert called_with["op"] == "psycopg.Cursor.copy" + assert called_with["kind"] == "sql_injection" + mock_run_vulnerability_scan.assert_called_once() + + database_conn.close() From faf25226f56d7f56dd4c0ec30c0cf71c523bb9b7 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Mon, 2 Sep 2024 12:30:52 +0200 Subject: [PATCH 6/7] Update psycopg code to remove args and kwargs --- aikido_firewall/sinks/psycopg.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aikido_firewall/sinks/psycopg.py b/aikido_firewall/sinks/psycopg.py index d9ec408c..d000733f 100644 --- a/aikido_firewall/sinks/psycopg.py +++ b/aikido_firewall/sinks/psycopg.py @@ -34,11 +34,11 @@ def aikido_execute(self, query, params=None, *args, **kwargs): ) return former_execute_function(self, query, params, *args, **kwargs) - def aikido_executemany(self, query, params_seq, *args, **kwargs): + def aikido_executemany(self, query, params_seq): args = (query, Postgres()) op = "psycopg.Cursor.executemany" vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=args) - return former_executemany_function(self, query, params_seq, *args, **kwargs) + return former_executemany_function(self, query, params_seq) setattr(psycopg.Cursor, "copy", aikido_copy) # pylint: disable=no-member setattr(psycopg.Cursor, "execute", aikido_execute) # pylint: disable=no-member From 198982b74697af8b1e5e0b229d7e6ab03b34a566 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Mon, 2 Sep 2024 12:37:09 +0200 Subject: [PATCH 7/7] Add psycopg as supported to docs --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index f495d876..336bc0dc 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,7 @@ Aikido Firewall for Python 3 is compatible with: * ✅ [`PyMySQL`](https://pypi.org/project/PyMySQL/) * ✅ [`pymongo`](https://pypi.org/project/pymongo/) * ✅ [`psycopg2`](https://pypi.org/project/psycopg2) +* ✅ [`psycopg`](https://pypi.org/project/psycopg) * ✅ [`asyncpg`](https://pypi.org/project/asyncpg) ## Reporting to your Aikido Security dashboard