Merge branch 'main' into AIK-3531

README.md

@@ -28,6 +28,7 @@ Aikido Firewall for Python 3 is compatible with:
 
 * ✅ [Django](docs/django.md)
 * ✅ [Flask](docs/flask.md)
+* ✅ [Quart](docs/quart.md)
 
 ### WSGI servers
 * ✅ [Gunicorn](docs/gunicorn.md)
@@ -38,6 +39,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

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

aikido_firewall/sinks/mysqlclient.py

@@ -6,30 +6,37 @@
 import importhook
 from aikido_firewall.vulnerabilities.sql_injection.dialects import MySQL
 from aikido_firewall.background_process.packages import add_wrapped_package
+from aikido_firewall.helpers.logging import logger
 import aikido_firewall.vulnerabilities as vulns
 
 
-@importhook.on_import("MySQLdb.connections")
+@importhook.on_import("MySQLdb.cursors")
 def on_mysqlclient_import(mysql):
     """
-    Hook 'n wrap on `MySQLdb.connections`
+    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
     """
     modified_mysql = importhook.copy_module(mysql)
-    prev_query_function = copy.deepcopy(mysql.Connection.query)
+    prev_execute_func = copy.deepcopy(mysql.Cursor.execute)
+    prev_executemany_func = copy.deepcopy(mysql.Cursor.executemany)
 
-    def aikido_new_query(_self, sql):
+    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.connections.query",
-            args=(sql.decode("utf-8"), MySQL()),
+            kind="sql_injection", op="MySQLdb.Cursor.execute", args=(query, MySQL())
         )
+        return prev_execute_func(self, query, args)
 
-        return prev_query_function(_self, sql)
+    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)
 
-    # pylint: disable=no-member
-    setattr(mysql.Connection, "query", aikido_new_query)
+    setattr(mysql.Cursor, "execute", aikido_new_execute)
+    setattr(mysql.Cursor, "executemany", aikido_new_executemany)
     add_wrapped_package("mysqlclient")
     return modified_mysql

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

aikido_firewall/sinks/pymongo.py

@@ -6,19 +6,26 @@
 import importhook
 from aikido_firewall.helpers.logging import logger
 import aikido_firewall.background_process.packages as pkgs
-from aikido_firewall.vulnerabilities import run_vulnerability_scan
+import aikido_firewall.vulnerabilities as vulns
+
+# find_one not present in list since find_one calls find function.
 
 OPERATIONS_WITH_FILTER = [
-    "replace_one",
-    "update_one",
-    "update_many",
-    "delete_one",
-    "delete_many",
-    "find_one",
-    "count_documents",
-    "find_one_and_delete",
-    "find_one_and_replace",
-    "find_one_and_update",
+    ("replace_one", [0, "filter"]),
+    ("update_one", [0, "filter"]),
+    ("update_many", [0, "filter"]),
+    ("delete_one", [0, "filter"]),
+    ("delete_many", [0, "filter"]),
+    ("count_documents", [0, "filter"]),
+    ("find_one_and_delete", [0, "filter"]),
+    ("find_one_and_replace", [0, "filter"]),
+    ("find_one_and_update", [0, "filter"]),
+    ("find", [0, "filter"]),
+    ("find_raw_batches", [0, "filter"]),
+    ("distinct", [1, "filter"]),
+    ("watch", [0, "pipeline"]),
+    ("aggregate", [0, "pipeline"]),
+    ("aggregate_raw_batches", [0, "pipeline"]),
 ]
 
 
@@ -32,23 +39,53 @@ def on_pymongo_import(pymongo):
     Returns : Modified pymongo.collection.Collection object
     """
     modified_pymongo = importhook.copy_module(pymongo)
-    for operation in OPERATIONS_WITH_FILTER:
-        if not hasattr(pymongo.Collection, operation):
-            logger.warning("Operation `%s` not found on Collection object.", operation)
+    for op_data in OPERATIONS_WITH_FILTER:
+        op = op_data[0]
+        if not hasattr(pymongo.Collection, op):
+            logger.warning("Operation `%s` not found on Collection object.", op)
 
-        prev_func = deepcopy(getattr(pymongo.Collection, operation))
+        prev_func = deepcopy(getattr(pymongo.Collection, op))
 
-        def wrapped_operation_function(
-            _self, _filter, *args, prev_func=prev_func, op=operation, **kwargs
+        def wrapped_op_func(
+            self,
+            *args,
+            prev_func=prev_func,
+            op_data=op_data,
+            **kwargs,
         ):
-            run_vulnerability_scan(
-                kind="nosql_injection",
-                op=f"pymongo.collection.Collection.{op}",
-                args=(_filter,),
-            )
-            return prev_func(_self, _filter, *args, **kwargs)
+            op, spot, key = op_data[0], op_data[1][0], op_data[1][1]
+            data = None
+            print(kwargs, key)
+            if kwargs.get(key, None):
+                # Keyword found, setting data
+                data = kwargs.get(key)
+            elif len(args) > spot and args[spot]:
+                data = args[spot]
+            if data:
+                vulns.run_vulnerability_scan(
+                    kind="nosql_injection",
+                    op=f"pymongo.collection.Collection.{op}",
+                    args=(data,),
+                )
+
+            return prev_func(self, *args, **kwargs)
+
+        setattr(modified_pymongo.Collection, op, wrapped_op_func)
+
+    # Add bulk_write support :
+    former_bulk_write = deepcopy(pymongo.Collection.bulk_write)
 
-        setattr(modified_pymongo.Collection, operation, wrapped_operation_function)
+    def aikido_bulk_write(self, requests, *args, **kwargs):
+        for request in requests:
+            if hasattr(request, "_filter"):
+                # Requested operation has a filter
+                vulns.run_vulnerability_scan(
+                    kind="nosql_injection",
+                    op="pymongo.collection.Collection.bulk_write",
+                    args=(request._filter,),
+                )
+        return former_bulk_write(self, requests, *args, **kwargs)
 
+    setattr(modified_pymongo.Collection, "bulk_write", aikido_bulk_write)
     pkgs.add_wrapped_package("pymongo")
     return modified_pymongo

aikido_firewall/sinks/pymysql.py

@@ -12,29 +12,35 @@
 logger = logging.getLogger("aikido_firewall")
 
 
-@importhook.on_import("pymysql.connections")
+@importhook.on_import("pymysql.cursors")
 def on_pymysql_import(mysql):
     """
-    Hook 'n wrap on `pymysql.connections`
-    Our goal is to wrap the query() function of the Connection class :
-    https://github.com/PyMySQL/PyMySQL/blob/95635f587ba9076e71a223b113efb08ac34a361d/pymysql/connections.py#L557
-    Returns : Modified pymysql.connections object
+    Hook 'n wrap on `pymysql.cursors`
+    Our goal is to wrap execute() and executemany() on Cursor class
+    https://github.com/PyMySQL/PyMySQL/blob/95635f587ba9076e71a223b113efb08ac34a361d/pymysql/cursors.py#L133
+    Returns : Modified pymysql.cursors object
     """
     modified_mysql = importhook.copy_module(mysql)
 
-    prev_query_function = copy.deepcopy(mysql.Connection.query)
+    prev_execute_func = copy.deepcopy(mysql.Cursor.execute)
+    prev_executemany_func = copy.deepcopy(mysql.Cursor.executemany)
 
-    def aikido_new_query(_self, sql, unbuffered=False):
-        if isinstance(sql, bytearray):
-            # executemany() gives a bytearray, convert to string.
-            sql = sql.decode("utf-8")
+    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="pymysql.connections.query", args=(sql, MySQL())
+            kind="sql_injection", op="pymysql.Cursor.execute", args=(query, MySQL())
         )
+        return prev_execute_func(self, query, args)
 
-        return prev_query_function(_self, sql, unbuffered=False)
+    def aikido_new_executemany(self, query, args):
+        op = "pymysql.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)
 
-    # pylint: disable=no-member
-    setattr(mysql.Connection, "query", aikido_new_query)
     add_wrapped_package("pymysql")
     return modified_mysql