psycopg2 : Only wrap psycopg2._ext.cursor class

aikido_firewall/sinks/psycopg2.py

@@ -9,84 +9,54 @@
 from aikido_firewall.vulnerabilities import run_vulnerability_scan
 
 
-class MutableAikidoConnection:
-    """Aikido's mutable connection class"""
+def wrap_cursor_factory(cursor_factory):
+    former_cursor_factory = copy.deepcopy(cursor_factory)
+    import psycopg2.extensions as ext
 
-    def __init__(self, former_conn):
-        self._former_conn = former_conn
-        self._cursor_func_copy = copy.deepcopy(former_conn.cursor)
-
-    def __getattr__(self, name):
-        if name != "cursor":
-            return getattr(self._former_conn, name)
-
-        # Return a function dynamically
-        def cursor(*args, **kwargs):
-            former_cursor = self._cursor_func_copy(*args, **kwargs)
-            return MutableAikidoCursor(former_cursor)
-
-        return cursor
-
-
-class MutableAikidoCursor:
-    """Aikido's mutable cursor class"""
-
-    def __init__(self, former_cursor):
-        self._former_cursor = former_cursor
-        self._execute_func_copy = copy.deepcopy(former_cursor.execute)
-        self._executemany_func_copy = copy.deepcopy(former_cursor.executemany)
-
-    def __getattr__(self, name):
-        if not name in ["execute", "executemany"]:
-            return getattr(self._former_cursor, name)
-
-        # Return a function dynamically
-        def execute(*args, **kwargs):
+    class AikidoWrappedCursor(ext.cursor):
+        def execute(self, *args, **kwargs):
+            """Aikido's wrapped execute function"""
             run_vulnerability_scan(
                 kind="sql_injection",
                 op="psycopg2.Connection.Cursor.execute",
                 args=(args[0], Postgres()),  #  args[0] : sql
             )
-            return self._execute_func_copy(*args, **kwargs)
+            if former_cursor_factory and hasattr(former_cursor_factory, "execute"):
+                return former_cursor_factory.execute(self, *args, **kwargs)
+            return super().execute(*args, **kwargs)
 
-        def executemany(*args, **kwargs):
+        def executemany(self, *args, **kwargs):
+            """Aikido's wrapped executemany function"""
             for sql in args[0]:
                 run_vulnerability_scan(
                     kind="sql_injection",
                     op="psycopg2.Connection.Cursor.executemany",
                     args=(sql, Postgres()),
                 )
-            return self._executemany_func_copy(*args, **kwargs)
+            if former_cursor_factory and hasattr(former_cursor_factory, "executemany"):
+                return former_cursor_factory.executemany(self, *args, **kwargs)
+            return super().executemany(*args, **kwargs)
 
-        if name == "execute":
-            return execute
-        return executemany
+    return AikidoWrappedCursor
 
 
-@importhook.on_import("psycopg2._psycopg")
+@importhook.on_import("psycopg2")
 def on_psycopg2_import(psycopg2):
     """
-    Hook 'n wrap on `psycopg2._psycopg._connect` function
-    1. We first instantiate a MutableAikidoConnection, because the connection
-    class is immutable.
-    2. This class has an adapted __getattr__ so that everything redirects to
-    the created actual connection, except for "cursor()" function!
-    3. When the cursor() function is executed, we instantiate a MutableAikidoCursor
-    which is also because the cursor class is immutable
-    4. when .execute() is executed on this cursor we can intercept it, the rest
-    gets redirected back using __getattr__ to the original cursor
-    Returns : Modified psycopg2._psycopg._connect function
+    Hook 'n wrap on `psycopg2.connect` function, we modify the cursor_factory
+    of the result of this connect function.
     """
     modified_psycopg2 = importhook.copy_module(psycopg2)
-    prev__connect_create = copy.deepcopy(psycopg2._connect)
+    former_connect_function = copy.deepcopy(psycopg2.connect)
 
-    def aik__connect(*args, **kwargs):
-        conn = prev__connect_create(*args, **kwargs)
-        return MutableAikidoConnection(conn)
+    def aikido_connect(*args, **kwargs):
+        former_conn = former_connect_function(*args, **kwargs)
+        former_conn.cursor_factory = wrap_cursor_factory(former_conn.cursor_factory)
+        return former_conn
 
     # pylint: disable=no-member
-    setattr(psycopg2, "_connect", aik__connect)
-    setattr(modified_psycopg2, "_connect", aik__connect)
+    setattr(psycopg2, "connect", aikido_connect)
+    setattr(modified_psycopg2, "connect", aikido_connect)
     add_wrapped_package("psycopg2")
     add_wrapped_package("psycopg2-binary")
     return modified_psycopg2

sample-apps/django-postgres/.env.example

@@ -0,0 +1,6 @@
+SECRET_KEY="test_key"
+
+# Aikido keys
+AIKIDO_DEBUG=true
+AIKIDO_TOKEN="AIK_secret_token"
+AIKIDO_BLOCKING=true

sample-apps/django-postgres/Dockerfile

@@ -0,0 +1,23 @@
+# Use an official Python runtime as a parent image
+FROM python:3
+
+#Copy code base
+COPY ./ /tmp
+
+# Set the working directory
+WORKDIR /app
+
+# Install dependencies
+RUN mv /tmp/sample-apps/django-postgres/requirements.txt ./
+RUN pip install -r requirements.txt
+
+# Build and install aikido_firewall from source
+WORKDIR /tmp
+RUN pip install poetry
+RUN rm -rf ./dist
+RUN make build
+RUN mv ./dist/aikido_firewall-*.tar.gz ./dist/aikido_firewall.tar.gz
+RUN pip install ./dist/aikido_firewall.tar.gz
+RUN pip list
+
+WORKDIR /app

sample-apps/django-postgres/README.md

@@ -0,0 +1,16 @@
+# Sample Django/Postgres App
+It runs **multi-threaded**
+
+## Getting started
+With docker-compose installed run
+```bash
+docker-compose up --build
+```
+This will expose a Django web server at [localhost:8094](http://localhost:8094)
+
+## URLS : 
+- Homepage : `http://localhost:8094/app`
+- Create a dog : `http://localhost:8094/app/create/<dog_name>`
+- MySQL attack : `Malicious dog", "Injected wrong boss name"); -- `
+
+To verify your attack was successfully note that the boss_name usualy is 'N/A', if you open the dog page (you can do this from the home page). You should see a "malicious dog" with a boss name that is not permitted.

sample-apps/django-postgres/docker-compose.benchmark.yml

@@ -0,0 +1,26 @@
+version: "3"
+services:
+  backend_firewall_disabled:
+    image: sample_django_mysql
+    command: sh -c "python3 manage.py migrate --noinput && python manage.py runserver 0.0.0.0:8000"
+    restart: always
+    volumes:
+      - .:/app
+    ports:
+      - "8095:8000"
+    depends_on:
+      - db
+    extra_hosts:
+      - "app.local.aikido.io:host-gateway"
+    environment:
+      FIREWALL_DISABLED: 1
+      SECRET_KEY: 'Test key'
+      DB_HOST: 'db'
+      DB_NAME: 'db'
+      DB_USER: 'user'
+      DB_PASSWORD: 'password'
+  backend:
+    environment:
+      AIKIDO_TOKEN: "test_aikido_token"
+      AIKIDO_BLOCKING: 1
+

sample-apps/django-postgres/docker-compose.yml

@@ -0,0 +1,50 @@
+version: '3'
+services:
+  db:
+    image: postgres:14-alpine
+    container_name: django_postgres_db
+    restart: always
+    volumes:
+      - db_data:/var/lib/postgresql/data
+    environment:
+      POSTGRES_DB: 'db'
+      POSTGRES_USER: 'user'
+      POSTGRES_PASSWORD: 'password'
+    ports:
+      - "5432:5432"
+    networks:
+      - default_network
+
+  backend:
+    image: sample_django_postgres
+    build: 
+      context: ./../../
+      dockerfile: ./sample-apps/django-postgres/Dockerfile
+    container_name: django_postgres_backend
+    command: sh -c "python3 manage.py migrate --noinput && python manage.py runserver 0.0.0.0:8000"
+    restart: always
+    volumes:
+      - .:/app
+    ports:
+      - "8094:8000"
+    depends_on:
+      - db
+    networks:
+      - default_network
+    extra_hosts:
+      - "app.local.aikido.io:host-gateway"
+    environment:
+      SECRET_KEY: 'Test key'
+      DB_HOST: 'db'
+      DB_NAME: 'db'
+      DB_USER: 'user'
+      DB_PASSWORD: 'password'
+      FIREWALL_DISABLED: 0
+
+
+volumes:
+  db_data:
+
+networks:
+  default_network:
+    driver: bridge

sample-apps/django-postgres/manage.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+"""Django's command-line utility for administrative tasks."""
+from dotenv import load_dotenv
+import os
+load_dotenv()
+firewall_disabled = os.getenv("FIREWALL_DISABLED")
+if firewall_disabled is not None:
+    if firewall_disabled.lower() != "1":
+        import aikido_firewall # Aikido package import
+        aikido_firewall.protect()
+
+import os
+import sys
+
+
+def main():
+    """Run administrative tasks."""
+    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'sample-django-postgres-app.settings')
+    try:
+        from django.core.management import execute_from_command_line
+    except ImportError as exc:
+        raise ImportError(
+            "Couldn't import Django. Are you sure it's installed and "
+            "available on your PYTHONPATH environment variable? Did you "
+            "forget to activate a virtual environment?"
+        ) from exc
+    execute_from_command_line(sys.argv)
+
+
+if __name__ == '__main__':
+    main()

sample-apps/django-postgres/requirements.txt

@@ -0,0 +1,4 @@
+django
+python-decouple
+cryptography
+psycopg2-binary

sample-apps/django-postgres/sample-django-postgres-app/asgi.py

@@ -0,0 +1,16 @@
+"""
+ASGI config for sample-django-postgres-app project.
+
+It exposes the ASGI callable as a module-level variable named ``application``.
+
+For more information on this file, see
+https://docs.djangoproject.com/en/5.0/howto/deployment/asgi/
+"""
+
+import os
+
+from django.core.asgi import get_asgi_application
+
+os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'sample-django-postgres-app.settings')
+
+application = get_asgi_application()