Improve the package compatibility check

aikido_zen/background_process/packages.py

@@ -1,37 +1,77 @@
 """Helper functions for packages"""
 
-import importlib.metadata as metadata
+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"
 
 
-def pkg_compat_check(pkg_name, required_version):
-    """Reports a newly wrapped package to the bg process"""
+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)"""
     # 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:
+            # 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
+        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)
+
+
+# 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

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)

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)
 

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)

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)

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,11 @@ 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)

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:

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)
 

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)
 

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)

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

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)
 

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)
 

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

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