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 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..d000733f --- /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 = (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) + + 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 diff --git a/aikido_firewall/sinks/tests/psycopg_test.py b/aikido_firewall/sinks/tests/psycopg_test.py new file mode 100644 index 00000000..358ca1f2 --- /dev/null +++ b/aikido_firewall/sinks/tests/psycopg_test.py @@ -0,0 +1,102 @@ +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 + + +@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 = 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() + 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() 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"]