From 0d2104256d3b104d544d36c1cc1ff13deb066d70 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 21 Mar 2025 12:26:46 +0100 Subject: [PATCH 1/5] Use new `is_package_compatible` function --- aikido_zen/background_process/packages_test.py | 8 ++++---- aikido_zen/sinks/asyncpg.py | 4 ++-- aikido_zen/sinks/mysqlclient.py | 4 ++-- aikido_zen/sinks/psycopg.py | 4 ++-- aikido_zen/sinks/psycopg2.py | 8 +++----- aikido_zen/sinks/pymongo.py | 2 +- aikido_zen/sinks/pymysql.py | 4 ++-- aikido_zen/sources/django/__init__.py | 4 ++-- aikido_zen/sources/flask.py | 4 ++-- aikido_zen/sources/gunicorn.py | 4 ++-- aikido_zen/sources/lxml.py | 4 ++-- aikido_zen/sources/quart.py | 4 ++-- aikido_zen/sources/starlette/__init__.py | 4 ++-- aikido_zen/sources/uwsgi.py | 4 ++-- 14 files changed, 30 insertions(+), 32 deletions(-) diff --git a/aikido_zen/background_process/packages_test.py b/aikido_zen/background_process/packages_test.py index bc45de2b7..66e3419a4 100644 --- a/aikido_zen/background_process/packages_test.py +++ b/aikido_zen/background_process/packages_test.py @@ -1,6 +1,6 @@ import pytest from unittest.mock import MagicMock -from .packages import pkg_compat_check, ANY_VERSION +from .packages import is_package_compatible, ANY_VERSION @pytest.fixture @@ -34,7 +34,7 @@ def test_pkg_compat_check_success(mock_get_comms, mocker): ) # Call the function under test - pkg_compat_check(pkg_name, ANY_VERSION) + is_package_compatible(pkg_name, ANY_VERSION) def test_pkg_compat_check_retry(mock_get_comms, mocker): @@ -59,7 +59,7 @@ def test_pkg_compat_check_retry(mock_get_comms, mocker): ) # Call the function under test - pkg_compat_check(pkg_name, ANY_VERSION) + is_package_compatible(pkg_name, ANY_VERSION) def test_pkg_compat_check_partial_success(mock_get_comms, mocker): @@ -84,4 +84,4 @@ def test_pkg_compat_check_partial_success(mock_get_comms, mocker): ) # Call the function under test - pkg_compat_check(pkg_name, ANY_VERSION) + is_package_compatible(pkg_name, ANY_VERSION) diff --git a/aikido_zen/sinks/asyncpg.py b/aikido_zen/sinks/asyncpg.py index f6d669426..bfca59fc0 100644 --- a/aikido_zen/sinks/asyncpg.py +++ b/aikido_zen/sinks/asyncpg.py @@ -4,7 +4,7 @@ import copy import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check +from aikido_zen.background_process.packages import is_package_compatible import aikido_zen.vulnerabilities as vulns from aikido_zen.helpers.logging import logger @@ -22,7 +22,7 @@ def on_asyncpg_import(asyncpg): https://github.com/MagicStack/asyncpg/blob/85d7eed40637e7cad73a44ed2439ffeb2a8dc1c2/asyncpg/connection.py#L43 Returns : Modified asyncpg.connection object """ - if not pkg_compat_check("asyncpg", REQUIRED_ASYNCPG_VERSION): + if not is_package_compatible("asyncpg", REQUIRED_ASYNCPG_VERSION): return asyncpg modified_asyncpg = importhook.copy_module(asyncpg) diff --git a/aikido_zen/sinks/mysqlclient.py b/aikido_zen/sinks/mysqlclient.py index 371b7abd6..ca712aca3 100644 --- a/aikido_zen/sinks/mysqlclient.py +++ b/aikido_zen/sinks/mysqlclient.py @@ -4,7 +4,7 @@ import copy import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check +from aikido_zen.background_process.packages import is_package_compatible from aikido_zen.helpers.logging import logger import aikido_zen.vulnerabilities as vulns @@ -19,7 +19,7 @@ def on_mysqlclient_import(mysql): https://github.com/PyMySQL/mysqlclient/blob/9fd238b9e3105dcbed2b009a916828a38d1f0904/src/MySQLdb/connections.py#L257 Returns : Modified MySQLdb.connections object """ - if not pkg_compat_check("mysqlclient", REQUIRED_MYSQLCLIENT_VERSION): + if not is_package_compatible("mysqlclient", REQUIRED_MYSQLCLIENT_VERSION): return mysql modified_mysql = importhook.copy_module(mysql) prev_execute_func = copy.deepcopy(mysql.Cursor.execute) diff --git a/aikido_zen/sinks/psycopg.py b/aikido_zen/sinks/psycopg.py index ba0787da1..18672c869 100644 --- a/aikido_zen/sinks/psycopg.py +++ b/aikido_zen/sinks/psycopg.py @@ -4,7 +4,7 @@ import copy import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check +from aikido_zen.background_process.packages import is_package_compatible import aikido_zen.vulnerabilities as vulns REQUIRED_PSYCOPG_VERSION = "3.1.0" @@ -16,7 +16,7 @@ def on_psycopg_import(psycopg): Hook 'n wrap on `psycopg.connect` function, we modify the cursor_factory of the result of this connect function. """ - if not pkg_compat_check("psycopg", REQUIRED_PSYCOPG_VERSION): + if not is_package_compatible("psycopg", REQUIRED_PSYCOPG_VERSION): return psycopg modified_psycopg = importhook.copy_module(psycopg) former_copy_funtcion = copy.deepcopy(psycopg.Cursor.copy) diff --git a/aikido_zen/sinks/psycopg2.py b/aikido_zen/sinks/psycopg2.py index bd99a5904..523a06988 100644 --- a/aikido_zen/sinks/psycopg2.py +++ b/aikido_zen/sinks/psycopg2.py @@ -4,7 +4,7 @@ import copy import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check +from aikido_zen.background_process.packages import is_package_compatible import aikido_zen.vulnerabilities as vulns PSYCOPG2_REQUIRED_VERSION = "2.9.2" @@ -49,10 +49,8 @@ def on_psycopg2_import(psycopg2): """ # Users can install either psycopg2 or psycopg2-binary, we need to check if at least # one is installed and if they meet version requirements : - if not pkg_compat_check( - "psycopg2", PSYCOPG2_REQUIRED_VERSION - ) and not pkg_compat_check("psycopg2-binary", PSYCOPG2_REQUIRED_VERSION): - # Both pyscopg2 and psycopg2-binary are not supported, abort wrapping + if not is_package_compatible(required_version=PSYCOPG2_REQUIRED_VERSION, packages=["psycopg2", "psycopg2-binary"]): + # Both psycopg2 and psycopg2-binary are not supported, don't wrapping return psycopg2 modified_psycopg2 = importhook.copy_module(psycopg2) former_connect_function = copy.deepcopy(psycopg2.connect) diff --git a/aikido_zen/sinks/pymongo.py b/aikido_zen/sinks/pymongo.py index 860232105..0db9bebe4 100644 --- a/aikido_zen/sinks/pymongo.py +++ b/aikido_zen/sinks/pymongo.py @@ -40,7 +40,7 @@ def on_pymongo_import(pymongo): https://github.com/mongodb/mongo-python-driver/blob/98658cfd1fea42680a178373333bf27f41153759/pymongo/synchronous/collection.py#L136 Returns : Modified pymongo.collection.Collection object """ - if not pkgs.pkg_compat_check("pymongo", REQUIRED_PYMONGO_VERSION): + if not pkgs.is_package_compatible("pymongo", REQUIRED_PYMONGO_VERSION): return pymongo modified_pymongo = importhook.copy_module(pymongo) for op_data in OPERATIONS_WITH_FILTER: diff --git a/aikido_zen/sinks/pymysql.py b/aikido_zen/sinks/pymysql.py index f85935408..ae1e5fdb1 100644 --- a/aikido_zen/sinks/pymysql.py +++ b/aikido_zen/sinks/pymysql.py @@ -5,7 +5,7 @@ import copy import logging import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check +from aikido_zen.background_process.packages import is_package_compatible import aikido_zen.vulnerabilities as vulns logger = logging.getLogger("aikido_zen") @@ -21,7 +21,7 @@ def on_pymysql_import(mysql): https://github.com/PyMySQL/PyMySQL/blob/95635f587ba9076e71a223b113efb08ac34a361d/pymysql/cursors.py#L133 Returns : Modified pymysql.cursors object """ - if not pkg_compat_check("pymysql", REQUIRED_PYMYSQL_VERSION): + if not is_package_compatible("pymysql", REQUIRED_PYMYSQL_VERSION): return mysql modified_mysql = importhook.copy_module(mysql) diff --git a/aikido_zen/sources/django/__init__.py b/aikido_zen/sources/django/__init__.py index 05f0416f2..9c16ed579 100644 --- a/aikido_zen/sources/django/__init__.py +++ b/aikido_zen/sources/django/__init__.py @@ -3,7 +3,7 @@ import copy import aikido_zen.importhook as importhook from aikido_zen.helpers.logging import logger -from aikido_zen.background_process.packages import pkg_compat_check, ANY_VERSION +from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION from ..functions.request_handler import request_handler from .run_init_stage import run_init_stage from .pre_response_middleware import pre_response_middleware @@ -17,7 +17,7 @@ def on_django_gunicorn_import(django): # https://github.com/django/django/blob/5865ff5adcf64da03d306dc32b36e87ae6927c85/django/core/handlers/base.py#L174 Returns : Modified django.core.handlers.base object """ - if not pkg_compat_check("django", required_version=ANY_VERSION): + if not is_package_compatible("django", required_version=ANY_VERSION): return django modified_django = importhook.copy_module(django) diff --git a/aikido_zen/sources/flask.py b/aikido_zen/sources/flask.py index e0f119c77..6fc862749 100644 --- a/aikido_zen/sources/flask.py +++ b/aikido_zen/sources/flask.py @@ -6,7 +6,7 @@ import aikido_zen.importhook as importhook from aikido_zen.helpers.logging import logger from aikido_zen.context import Context -from aikido_zen.background_process.packages import pkg_compat_check, ANY_VERSION +from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION from aikido_zen.context import get_current_context import aikido_zen.sources.functions.request_handler as funcs @@ -103,7 +103,7 @@ def on_flask_import(flask): @app.route |-> `add_url_rule` |-> self.view_functions. these get called via full_dispatch_request, which we wrap. We also wrap __call__ to run our middleware. """ - if not pkg_compat_check("flask", required_version=FLASK_REQUIRED_VERSION): + if not is_package_compatible("flask", required_version=FLASK_REQUIRED_VERSION): return flask modified_flask = importhook.copy_module(flask) former_fdr = copy.deepcopy(flask.Flask.full_dispatch_request) diff --git a/aikido_zen/sources/gunicorn.py b/aikido_zen/sources/gunicorn.py index dad88d269..caa704e80 100644 --- a/aikido_zen/sources/gunicorn.py +++ b/aikido_zen/sources/gunicorn.py @@ -1,11 +1,11 @@ """Gunicorn Module, report if module was found""" import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check, ANY_VERSION +from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION @importhook.on_import("gunicorn") def on_gunicorn_import(gunicorn): """Report to the core when gunicorn gets imported""" - pkg_compat_check("gunicorn", required_version=ANY_VERSION) + is_package_compatible("gunicorn", required_version=ANY_VERSION) return gunicorn diff --git a/aikido_zen/sources/lxml.py b/aikido_zen/sources/lxml.py index b046d62fd..0519ce7da 100644 --- a/aikido_zen/sources/lxml.py +++ b/aikido_zen/sources/lxml.py @@ -7,7 +7,7 @@ from aikido_zen.helpers.extract_data_from_xml_body import ( extract_data_from_xml_body, ) -from aikido_zen.background_process.packages import pkg_compat_check, ANY_VERSION +from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION @importhook.on_import("lxml.etree") @@ -18,7 +18,7 @@ def on_lxml_import(eltree): - Wrap on Returns : Modified `lxml.etree` object """ - if not pkg_compat_check("lxml", required_version=ANY_VERSION): + if not is_package_compatible("lxml", required_version=ANY_VERSION): return eltree modified_eltree = importhook.copy_module(eltree) diff --git a/aikido_zen/sources/quart.py b/aikido_zen/sources/quart.py index d133b394a..a2057362f 100644 --- a/aikido_zen/sources/quart.py +++ b/aikido_zen/sources/quart.py @@ -6,7 +6,7 @@ import aikido_zen.importhook as importhook from aikido_zen.helpers.logging import logger from aikido_zen.context import Context, get_current_context -from aikido_zen.background_process.packages import pkg_compat_check, ANY_VERSION +from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION from .functions.request_handler import request_handler @@ -87,7 +87,7 @@ def on_quart_import(quart): Hook 'n wrap on `quart.app` Our goal is to wrap the __call__, handle_request, asgi_app functios of the "Quart" class """ - if not pkg_compat_check("quart", required_version=ANY_VERSION): + if not is_package_compatible("quart", required_version=ANY_VERSION): return quart modified_quart = importhook.copy_module(quart) diff --git a/aikido_zen/sources/starlette/__init__.py b/aikido_zen/sources/starlette/__init__.py index 48deabede..e22cb1891 100644 --- a/aikido_zen/sources/starlette/__init__.py +++ b/aikido_zen/sources/starlette/__init__.py @@ -12,7 +12,7 @@ """ import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check, ANY_VERSION +from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION @importhook.on_import("starlette") @@ -21,7 +21,7 @@ def on_starlette_import(starlette): This checks for the package version of starlette so you don't have to do it twice, once in starlette_applications and once in starlette_applications. """ - if not pkg_compat_check("starlette", required_version=ANY_VERSION): + if not is_package_compatible("starlette", required_version=ANY_VERSION): return starlette # Package is compatible, start wrapping : import aikido_zen.sources.starlette.starlette_applications diff --git a/aikido_zen/sources/uwsgi.py b/aikido_zen/sources/uwsgi.py index f374084c6..f2a61f4ee 100644 --- a/aikido_zen/sources/uwsgi.py +++ b/aikido_zen/sources/uwsgi.py @@ -1,11 +1,11 @@ """UWSGI Module, report if module was found""" import aikido_zen.importhook as importhook -from aikido_zen.background_process.packages import pkg_compat_check, ANY_VERSION +from aikido_zen.background_process.packages import is_package_compatible, ANY_VERSION @importhook.on_import("uwsgi") def on_uwsgi_import(uwsgi): """Report to the core when uwsgi gets imported""" - pkg_compat_check("uwsgi", required_version=ANY_VERSION) + is_package_compatible("uwsgi", required_version=ANY_VERSION) return uwsgi From 12c4a34093bc189a6c068b2b062220bf3e92a221 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 21 Mar 2025 12:27:04 +0100 Subject: [PATCH 2/5] Update is_package_compatible function --- aikido_zen/background_process/packages.py | 45 ++++++++++++----------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index 23987e752..9f34ab7d6 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -1,6 +1,7 @@ """Helper functions for packages""" -import importlib.metadata as metadata +import importlib.metadata as importilb_metadata +import importlib.util as importlib_util from packaging.version import Version from aikido_zen.helpers.logging import logger import aikido_zen.background_process.comms as comms @@ -11,27 +12,29 @@ ANY_VERSION = "0.0.0" -def pkg_compat_check(pkg_name, required_version): +def is_package_compatible(package=None, required_version=ANY_VERSION, packages=None): """Reports a newly wrapped package to the bg process""" # Fetch package version : + if package is not None: + packages = [package] + if packages is None: + return False # no package names provided, return false. try: - pkg_version = metadata.version(pkg_name) - except metadata.PackageNotFoundError: - logger.info( - "Version for %s is undetermined. Zen is unable to protect this module.", - pkg_name, - ) - return False # We don't support it since we are not sure what it is. - - # Check if the package version is supported : - version_supported = is_version_supported(pkg_version, required_version) - if version_supported: - logger.debug("Instrumentation for %s=%s supported", pkg_name, pkg_version) - else: - logger.info("Zen does not support %s=%s", pkg_name, pkg_version) - return version_supported - - -def is_version_supported(pkg_verion, required_version): + for package in packages: + if importlib_util.find_spec(package) is None: + continue # package name is not installed + package_version = importilb_metadata.version(package) + if is_version_supported(package_version, required_version): + logger.debug("Instrumentation for %s=%s supported", package, package_version) + return True + + # No match found + logger.info("Zen does not support %s", packages) + return False + except Exception as e: + logger.debug("Exception occurred in is_package_compatible: %s", e) + return False # We don't support it, since something unexpected happened in checking compatibility. + +def is_version_supported(version, required_version): """Checks if the package version is supported""" - return Version(pkg_verion) >= Version(required_version) + return Version(version) >= Version(required_version) From ac1d3989672b55a134967595870ef9e4b1b92412 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 21 Mar 2025 12:55:34 +0100 Subject: [PATCH 3/5] Add packagesStore and improve psycopg2 writing --- aikido_zen/background_process/packages.py | 59 ++++++++++++++++++----- aikido_zen/sinks/psycopg2.py | 5 +- 2 files changed, 52 insertions(+), 12 deletions(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index 9f34ab7d6..f573aa3aa 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -1,12 +1,9 @@ """Helper functions for packages""" -import importlib.metadata as importilb_metadata -import importlib.util as importlib_util +import importlib.metadata as importlib_metadata + from packaging.version import Version from aikido_zen.helpers.logging import logger -import aikido_zen.background_process.comms as comms - -MAX_REPORT_TRIES = 5 # If any version is supported, this constant can be used ANY_VERSION = "0.0.0" @@ -18,14 +15,27 @@ def is_package_compatible(package=None, required_version=ANY_VERSION, packages=N if package is not None: packages = [package] if packages is None: - return False # no package names provided, return false. + return False # no package names provided, return false. try: for package in packages: - if importlib_util.find_spec(package) is None: - continue # package name is not installed - package_version = importilb_metadata.version(package) - if is_version_supported(package_version, required_version): - logger.debug("Instrumentation for %s=%s supported", package, package_version) + # Checks if we already looked up the package : + if PackagesStore.get_package(package) is not None: + return PackagesStore.get_package(package)["supported"] + + # Safely get the package version, with an exception for when the package was not found + try: + package_version = importlib_metadata.version(package) + except importlib_metadata.PackageNotFoundError: + continue + + # Check support and store package for later + supported = is_version_supported(package_version, required_version) + PackagesStore.add_package(package, package_version, supported) + + if supported: + logger.debug( + "Instrumentation for %s=%s supported", package, package_version + ) return True # No match found @@ -35,6 +45,33 @@ def is_package_compatible(package=None, required_version=ANY_VERSION, packages=N logger.debug("Exception occurred in is_package_compatible: %s", e) return False # We don't support it, since something unexpected happened in checking compatibility. + def is_version_supported(version, required_version): """Checks if the package version is supported""" return Version(version) >= Version(required_version) + + +# packages store, uses python's built in GlobalInterpreterLock (GIL) +packages = dict() + + +class PackagesStore: + @staticmethod + def get_packages(): + global packages + return packages + + @staticmethod + def add_package(package, version, supported): + global packages + packages[package] = { + "version": version, + "supported": bool(supported), + } + + @staticmethod + def get_package(package_name): + global packages + if package_name in packages: + return packages[package_name] + return None diff --git a/aikido_zen/sinks/psycopg2.py b/aikido_zen/sinks/psycopg2.py index 523a06988..afd5a0fe5 100644 --- a/aikido_zen/sinks/psycopg2.py +++ b/aikido_zen/sinks/psycopg2.py @@ -49,7 +49,10 @@ def on_psycopg2_import(psycopg2): """ # Users can install either psycopg2 or psycopg2-binary, we need to check if at least # one is installed and if they meet version requirements : - if not is_package_compatible(required_version=PSYCOPG2_REQUIRED_VERSION, packages=["psycopg2", "psycopg2-binary"]): + if not is_package_compatible( + required_version=PSYCOPG2_REQUIRED_VERSION, + packages=["psycopg2", "psycopg2-binary"], + ): # Both psycopg2 and psycopg2-binary are not supported, don't wrapping return psycopg2 modified_psycopg2 = importhook.copy_module(psycopg2) From 31e0855a87c5a938542efd4f8a7d8105b3c7b403 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Tue, 1 Apr 2025 08:34:40 +0000 Subject: [PATCH 4/5] Update aikido_zen/background_process/packages.py --- aikido_zen/background_process/packages.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index f573aa3aa..cff35bc22 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -10,7 +10,7 @@ def is_package_compatible(package=None, required_version=ANY_VERSION, packages=None): - """Reports a newly wrapped package to the bg process""" + """Checks for compatibility of one or multiple package names (in the case of psycopg, you need to check multiple names.""" # Fetch package version : if package is not None: packages = [package] From 67030b441ef7881e04a8f798fd3bf9c9ce3cebe8 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Tue, 1 Apr 2025 08:35:04 +0000 Subject: [PATCH 5/5] Update aikido_zen/background_process/packages.py --- aikido_zen/background_process/packages.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index cff35bc22..19425e37a 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -10,7 +10,7 @@ def is_package_compatible(package=None, required_version=ANY_VERSION, packages=None): - """Checks for compatibility of one or multiple package names (in the case of psycopg, you need to check multiple names.""" + """Checks for compatibility of one or multiple package names (in the case of psycopg, you need to check multiple names)""" # Fetch package version : if package is not None: packages = [package]