Python 3.13 support: wrapping update

.github/workflows/benchmark.yml

@@ -1,7 +1,8 @@
-name: Benchmark
+name: 📊 Benchmarks
 on:
   push: {}
   workflow_call: {}
+
 jobs:
   benchmark_sql_algorithm:
     runs-on: ubuntu-latest
@@ -13,7 +14,7 @@ jobs:
         working-directory:  ./sample-apps/databases
         run: docker compose up --build -d
       - name: Set up Python 3.9
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: 3.9
 
@@ -30,15 +31,24 @@ jobs:
         run: |
           poetry run python ./benchmarks/sql_benchmark/sql_benchmark_fw.py
           poetry run python ./benchmarks/sql_benchmark/sql_benchmark_no_fw.py
+
   benchmark_with_flask_mysql:
     runs-on: ubuntu-latest
     timeout-minutes: 10
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
       - name: Start databases
         working-directory:  ./sample-apps/databases
         run: docker compose up --build -d
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
       - name: Install dependencies & build
         run: |
           python -m pip install --upgrade pip
@@ -51,15 +61,24 @@ jobs:
       - name: Run flask-mysql k6 Benchmark
         run: |
           k6 run -q ./benchmarks/flask-mysql-benchmarks.js
+
   benchmark_with_starlette_app:
     runs-on: ubuntu-latest
     timeout-minutes: 10
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
       - name: Start databases
         working-directory:  ./sample-apps/databases
         run: docker compose up --build -d
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
       - name: Install dependencies & build
         run: |
           python -m pip install --upgrade pip
@@ -71,9 +90,5 @@ jobs:
         run: |
           sudo apt-get update
           sudo apt-get install -y wrk
-      - name: Set up Python 3.9
-        uses: actions/setup-python@v2
-        with:
-          python-version: 3.9
       - name: Run benchmark
         run: python ./benchmarks/starlette_benchmark.py

.github/workflows/end2end.yml

@@ -35,7 +35,7 @@ jobs:
           - { name: flask-postgres-xml, testfile: end2end/flask_postgres_xml_lxml_test.py }
           - { name: quart-postgres-uvicorn, testfile: end2end/quart_postgres_uvicorn_test.py }
           - { name: starlette-postgres-uvicorn, testfile: end2end/starlette_postgres_uvicorn_test.py }
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Install packages
         run: sudo apt update && sudo apt install python3-dev libmysqlclient-dev
@@ -53,7 +53,7 @@ jobs:
             docker run --name mock_core -d -p 5000:5000 mock_core
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies

.github/workflows/lint.yml

@@ -9,7 +9,7 @@ jobs:
       uses: actions/checkout@v4
 
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: '3.x'
 

.github/workflows/smoke-test-ffi.yml

@@ -20,7 +20,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: 3.12
 

.github/workflows/unit-test.yml

@@ -10,7 +10,7 @@ jobs:
       # Don't cancel jobs if one fails
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
     steps:
     - name: Checkout code
       uses: actions/checkout@v4
@@ -22,7 +22,7 @@ jobs:
       run: |
           sudo echo "127.0.0.1 local.aikido.io" | sudo tee -a /etc/hosts
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
 

aikido_zen/sinks/__init__.py

@@ -1,0 +1,70 @@
+from wrapt import wrap_object, FunctionWrapper, when_imported
+from aikido_zen.background_process.packages import ANY_VERSION, is_package_compatible
+from aikido_zen.errors import AikidoException
+from aikido_zen.helpers.logging import logger
+
+
+def on_import(name, package="", version_requirement=ANY_VERSION):
+    """
+    Decorator to register a function to be called when a package is imported.
+    It checks if the package is compatible with the specified version requirement.
+    """
+
+    def decorator(func):
+        if package and not is_package_compatible(package, version_requirement):
+            return
+
+        when_imported(name)(func)  # Register the function to be called on import
+
+    return decorator
+
+
+def patch_function(module, name, wrapper):
+    """
+    Patches a function in the specified module with a wrapper function.
+    """
+    try:
+        wrap_object(module, name, FunctionWrapper, (wrapper,))
+    except Exception as e:
+        logger.info("Failed to wrap %s:%s, due to: %s", module, name, e)
+
+
+def before(wrapper):
+    """
+    Surrounds a patch with try-except and calls the original function at the end
+    """
+
+    def decorator(func, instance, args, kwargs):
+        try:
+            wrapper(func, instance, args, kwargs)  # Call the patch
+        except AikidoException as e:
+            raise e  # Re-raise AikidoException
+        except Exception as e:
+            logger.debug(
+                "%s:%s wrapping-before error: %s", func.__module__, func.__name__, e
+            )
+
+        return func(*args, **kwargs)  # Call the original function
+
+    return decorator
+
+
+def after(wrapper):
+    """
+    Surrounds a patch with try-except, calls the original function and gives the return value to the patch
+    """
+
+    def decorator(func, instance, args, kwargs):
+        return_value = func(*args, **kwargs)  # Call the original function
+        try:
+            wrapper(func, instance, args, kwargs, return_value)  # Call the patch
+        except AikidoException as e:
+            raise e  # Re-raise AikidoException
+        except Exception as e:
+            logger.debug(
+                "%s:%s wrapping-after error: %s", func.__module__, func.__name__, e
+            )
+
+        return return_value
+
+    return decorator

aikido_zen/sinks/asyncpg.py

@@ -2,64 +2,28 @@
 Sink module for `asyncpg`
 """
 
-import copy
-import aikido_zen.importhook as importhook
-from aikido_zen.background_process.packages import is_package_compatible
 import aikido_zen.vulnerabilities as vulns
-from aikido_zen.helpers.logging import logger
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import patch_function, before, on_import
 
-REQUIRED_ASYNCPG_VERSION = "0.27.0"
 
-
-@importhook.on_import("asyncpg.connection")
-def on_asyncpg_import(asyncpg):
+@on_import("asyncpg.connection", "asyncpg", version_requirement="0.27.0")
+def patch(m):
     """
-    Hook 'n wrap on `asyncpg.connection`
-    * the Cursor classes in asyncpg.cursor are only used to fetch data. (Currently not supported)
-    * Pool class uses Connection class (Wrapping supported for Connection class)
-    * _execute(...) get's called by all except execute and executemany
-    Our goal is to wrap the _execute(), execute(), executemany() functions in Connection class :
-    https://github.com/MagicStack/asyncpg/blob/85d7eed40637e7cad73a44ed2439ffeb2a8dc1c2/asyncpg/connection.py#L43
-    Returns : Modified asyncpg.connection object
+    patching module asyncpg.connection
+    - patches Connection.execute, Connection.executemany, Connection._execute
+    - doesn't patch Cursor class -> are only used to fetch data.
+    - doesn't patch Pool class -> uses Connection class
+    src: https://github.com/MagicStack/asyncpg/blob/85d7eed40637e7cad73a44ed2439ffeb2a8dc1c2/asyncpg/connection.py#L43
     """
-    if not is_package_compatible("asyncpg", REQUIRED_ASYNCPG_VERSION):
-        return asyncpg
-    modified_asyncpg = importhook.copy_module(asyncpg)
-
-    # pylint: disable=protected-access # We need to wrap this function
-    former__execute = copy.deepcopy(asyncpg.Connection._execute)
-    former_executemany = copy.deepcopy(asyncpg.Connection.executemany)
-    former_execute = copy.deepcopy(asyncpg.Connection.execute)
-
-    def aikido_new__execute(_self, query, *args, **kwargs):
-        vulns.run_vulnerability_scan(
-            kind="sql_injection",
-            op="asyncpg.connection.Connection._execute",
-            args=(query, "postgres"),
-        )
-
-        return former__execute(_self, query, *args, **kwargs)
-
-    def aikido_new_executemany(_self, query, *args, **kwargs):
-        # This query is just a string, not a list, see docs.
-        vulns.run_vulnerability_scan(
-            kind="sql_injection",
-            op="asyncpg.connection.Connection.executemany",
-            args=(query, "postgres"),
-        )
-        return former_executemany(_self, query, *args, **kwargs)
+    patch_function(m, "Connection.execute", _execute)
+    patch_function(m, "Connection.executemany", _execute)
+    patch_function(m, "Connection._execute", _execute)
 
-    def aikido_new_execute(_self, query, *args, **kwargs):
-        vulns.run_vulnerability_scan(
-            kind="sql_injection",
-            op="asyncpg.connection.Connection.execute",
-            args=(query, "postgres"),
-        )
-        return former_execute(_self, query, *args, **kwargs)
 
-    # pylint: disable=no-member
-    setattr(asyncpg.Connection, "_execute", aikido_new__execute)
-    setattr(asyncpg.Connection, "executemany", aikido_new_executemany)
-    setattr(asyncpg.Connection, "execute", aikido_new_execute)
+@before
+def _execute(func, instance, args, kwargs):
+    query = get_argument(args, kwargs, 0, "query")
 
-    return modified_asyncpg
+    op = f"asyncpg.connection.Connection.{func.__name__}"
+    vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=(query, "postgres"))

aikido_zen/sinks/builtins.py

@@ -3,34 +3,26 @@
 """
 
 from pathlib import PurePath
-import aikido_zen.importhook as importhook
 import aikido_zen.vulnerabilities as vulns
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import patch_function, on_import, before
 
 
-def aikido_open_decorator(func):
-    """Decorator for open(...)"""
+@before
+def _open(func, instance, args, kwargs):
+    filename = get_argument(args, kwargs, 0, "filename")
+    if not isinstance(filename, (str, bytes, PurePath)):
+        return
 
-    def wrapper(*args, **kwargs):
-        #  args[0] is thefunc_name filename
-        if len(args) > 0 and isinstance(args[0], (str, bytes, PurePath)):
-            vulns.run_vulnerability_scan(
-                kind="path_traversal", op="builtins.open", args=(args[0],)
-            )
-        return func(*args, **kwargs)
+    vulns.run_vulnerability_scan(
+        kind="path_traversal", op="builtins.open", args=(filename,)
+    )
 
-    return wrapper
 
-
-@importhook.on_import("builtins")
-def on_builtins_import(builtins):
+@on_import("builtins")
+def patch(m):
     """
-    Hook 'n wrap on `builtins`, python's built-in functions
-    Our goal is to wrap the open() function, which you use when opening files
-    Returns : Modified builtins object
+    patching module builtins
+    - patches builtins.open
     """
-    modified_builtins = importhook.copy_module(builtins)
-
-    # pylint: disable=no-member
-    setattr(builtins, "open", aikido_open_decorator(builtins.open))
-    setattr(modified_builtins, "open", aikido_open_decorator(builtins.open))
-    return modified_builtins
+    patch_function(m, "open", _open)