New Wrapping PR 4: Update all sinks to new wrapping system

aikido_zen/background_process/packages.py

@@ -39,7 +39,7 @@ def is_package_compatible(package=None, required_version=ANY_VERSION, packages=N
                 return True
 
         # No match found
-        logger.info("Zen does not support %s", packages)
+        logger.info("Zen does not support current version of %s", "/".join(packages))
         return False
     except Exception as e:
         logger.debug("Exception occurred in is_package_compatible: %s", e)

aikido_zen/sinks/__init__.py

@@ -1,5 +1,4 @@
 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
@@ -12,10 +11,17 @@ def on_import(name, package="", version_requirement=ANY_VERSION):
     """
 
     def decorator(func):
-        if package and not is_package_compatible(package, version_requirement):
-            return
+        def check_pkg_wrapper(f):
+            def wrapper(*args, **kwargs):
+                # This code runs only on import
+                if package and not is_package_compatible(package, version_requirement):
+                    return
+                return f(*args, **kwargs)
+
+            return wrapper
 
-        when_imported(name)(func)  # Register the function to be called on import
+        # Register the function to be called on import
+        when_imported(name)(check_pkg_wrapper(func))
 
     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)

aikido_zen/sinks/http_client.py

@@ -2,49 +2,35 @@
 Sink module for `http`
 """
 
-import copy
-import aikido_zen.importhook as importhook
-from aikido_zen.helpers.logging import logger
-from aikido_zen.vulnerabilities import run_vulnerability_scan
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import before, after, patch_function, on_import
 from aikido_zen.vulnerabilities.ssrf.handle_http_response import (
     handle_http_response,
 )
 from aikido_zen.helpers.try_parse_url import try_parse_url
-from aikido_zen.errors import AikidoException
 
 
-@importhook.on_import("http.client")
-def on_http_import(http):
-    """
-    Hook 'n wrap on `http.client.HTTPConnection.putrequest`
-    Our goal is to wrap the putrequest() function of the HTTPConnection class :
-    https://github.com/python/cpython/blob/372df1950817dfcf8b9bac099448934bf8657cf5/Lib/http/client.py#L1136
-    Returns : Modified http.client object
-    """
-    modified_http = importhook.copy_module(http)
-    former_putrequest = copy.deepcopy(http.HTTPConnection.putrequest)
-    former_getresponse = copy.deepcopy(http.HTTPConnection.getresponse)
+@before
+def _putrequest(func, instance, args, kwargs):
+    # putrequest(...) is called with path argument, store this on the HTTPConnection
+    # instance for later use in the getresponse(...) function.
+    path = get_argument(args, kwargs, 1, "path")
+    setattr(instance, "_aikido_var_path", path)
+
 
-    def aik_new_putrequest(_self, method, path, *args, **kwargs):
-        #  Aikido putrequest, gets called before the request goes through
-        # Set path for aik_new_getresponse :
-        _self.aikido_attr_path = path
-        return former_putrequest(_self, method, path, *args, **kwargs)
+@after
+def _getresponse(func, instance, args, kwargs, return_value):
+    path = getattr(instance, "_aikido_var_path")
+    source_url = try_parse_url(f"http://{instance.host}:{instance.port}{path}")
+    handle_http_response(http_response=return_value, source=source_url)
 
-    def aik_new_getresponse(_self):
-        #  Aikido getresponse, gets called after the request is complete
-        # And fetches the response
-        response = former_getresponse(_self)
-        try:
-            assembled_url = f"http://{_self.host}:{_self.port}{_self.aikido_attr_path}"
-            source_url = try_parse_url(assembled_url)
-            handle_http_response(http_response=response, source=source_url)
-        except Exception as e:
-            logger.debug("Exception occurred in custom getresponse function : %s", e)
-        return response
 
-    # pylint: disable=no-member
-    setattr(http.HTTPConnection, "putrequest", aik_new_putrequest)
-    # pylint: disable=no-member
-    setattr(http.HTTPConnection, "getresponse", aik_new_getresponse)
-    return modified_http
+@on_import("http.client")
+def patch(m):
+    """
+    patching module http.client
+    - patches HTTPConnection.putrequest -> stores path
+    - patches HTTPConnection.getresponse -> handles response object
+    """
+    patch_function(m, "HTTPConnection.putrequest", _putrequest)
+    patch_function(m, "HTTPConnection.getresponse", _getresponse)

aikido_zen/sinks/io.py

@@ -2,34 +2,35 @@
 Sink module for python's `io`
 """
 
-import copy
-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, before, on_import
 
-KIND = "path_traversal"
 
+@before
+def _open(func, instance, args, kwargs):
+    file = get_argument(args, kwargs, 0, "file")
+    if not file:
+        return
+
+    vulns.run_vulnerability_scan(kind="path_traversal", op="io.open", args=(file,))
 
-@importhook.on_import("io")
-def on_io_import(io):
-    """
-    Hook 'n wrap on `io`, wrapping io.open(...) and io.open_code(...)
-    Returns : Modified io object
-    """
-    modified_io = importhook.copy_module(io)
-    former_open_func = copy.deepcopy(io.open)
-    former_open_code_func = copy.deepcopy(io.open_code)
 
-    def aikido_open_func(file, *args, **kwargs):
-        if file:
-            vulns.run_vulnerability_scan(kind=KIND, op="io.open", args=(file,))
-        return former_open_func(file, *args, **kwargs)
+@before
+def _open_code(func, instance, args, kwargs):
+    path = get_argument(args, kwargs, 0, "path")
+    if not path:
+        return
 
-    def aikido_open_code_func(path):
-        if path:
-            vulns.run_vulnerability_scan(kind=KIND, op="io.open_code", args=(path,))
-        return former_open_code_func(path)
+    vulns.run_vulnerability_scan(kind="path_traversal", op="io.open_code", args=(path,))
 
-    setattr(modified_io, "open", aikido_open_func)
-    setattr(modified_io, "open_code", aikido_open_code_func)
 
-    return modified_io
+@on_import("io")
+def patch(m):
+    """
+    patching module io
+    - patches io.open(file, ...)
+    - patches io.open_code(path)
+    """
+    patch_function(m, "open", _open)
+    patch_function(m, "open_code", _open_code)

aikido_zen/sinks/mysqlclient.py

@@ -2,43 +2,38 @@
 Sink module for `mysqlclient`
 """
 
-import copy
-import aikido_zen.importhook as importhook
-from aikido_zen.background_process.packages import is_package_compatible
-from aikido_zen.helpers.logging import logger
+from aikido_zen.helpers.get_argument import get_argument
 import aikido_zen.vulnerabilities as vulns
+from aikido_zen.sinks import patch_function, on_import, before
 
-REQUIRED_MYSQLCLIENT_VERSION = "1.5.0"
 
-
-@importhook.on_import("MySQLdb.cursors")
-def on_mysqlclient_import(mysql):
+@on_import("MySQLdb.cursors", "mysqlclient", version_requirement="1.5.0")
+def patch(m):
     """
-    Hook 'n wrap on `MySQLdb.cursors`
-    Our goal is to wrap the query() function of the Connection class :
-    https://github.com/PyMySQL/mysqlclient/blob/9fd238b9e3105dcbed2b009a916828a38d1f0904/src/MySQLdb/connections.py#L257
-    Returns : Modified MySQLdb.connections object
+    patching MySQLdb.cursors (mysqlclient)
+    - patches Cursor.execute(query, ...)
+    - patches Cursor.executemany(query, ...)
     """
-    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)
-    prev_executemany_func = copy.deepcopy(mysql.Cursor.executemany)
-
-    def aikido_new_execute(self, query, args=None):
-        if isinstance(query, bytearray):
-            logger.debug("Query is bytearray, normally comes from executemany.")
-            return prev_execute_func(self, query, args)
-        vulns.run_vulnerability_scan(
-            kind="sql_injection", op="MySQLdb.Cursor.execute", args=(query, "mysql")
-        )
-        return prev_execute_func(self, query, args)
-
-    def aikido_new_executemany(self, query, args):
-        op = "MySQLdb.Cursor.executemany"
-        vulns.run_vulnerability_scan(kind="sql_injection", op=op, args=(query, "mysql"))
-        return prev_executemany_func(self, query, args)
-
-    setattr(mysql.Cursor, "execute", aikido_new_execute)
-    setattr(mysql.Cursor, "executemany", aikido_new_executemany)
-    return modified_mysql
+    patch_function(m, "Cursor.execute", _execute)
+    patch_function(m, "Cursor.executemany", _executemany)
+
+
+@before
+def _execute(func, instance, args, kwargs):
+    query = get_argument(args, kwargs, 0, "query")
+    if isinstance(query, bytearray):
+        # If query is type bytearray, it will be picked up by our wrapping of executemany
+        return
+
+    vulns.run_vulnerability_scan(
+        kind="sql_injection", op="MySQLdb.Cursor.execute", args=(query, "mysql")
+    )
+
+
+@before
+def _executemany(func, instance, args, kwargs):
+    query = get_argument(args, kwargs, 0, "query")
+
+    vulns.run_vulnerability_scan(
+        kind="sql_injection", op="MySQLdb.Cursor.executemany", args=(query, "mysql")
+    )