Fix django bugs

aikido_firewall/__init__.py

@@ -15,18 +15,25 @@
 load_dotenv()
 
 
-def protect(module="any"):
+def protect(module="any", server=True):
     """
     Protect user's application
     """
+    if server:
+        start_background_process()
+    else:
+        logger.debug("Not starting background process")
+    if module == "server-only":
+        return
+
     # Import sources
-    import aikido_firewall.sources.django
+    if module == "django":
+        import aikido_firewall.sources.django
 
-    if module != "django":
+    if not module in ["django", "django-gunicorn"]:
         import aikido_firewall.sources.flask
 
     # Import sinks
     import aikido_firewall.sinks.pymysql
 
     logger.info("Aikido python firewall started")
-    start_background_process()

aikido_firewall/background_process/__init__.py

@@ -6,6 +6,8 @@
 import time
 import os
 import secrets
+import signal
+import socket
 import multiprocessing.connection as con
 from multiprocessing import Process
 from threading import Thread
@@ -25,7 +27,14 @@
 
     def __init__(self, address, key):
         logger.debug("Background process started")
-        listener = con.Listener(address, authkey=key)
+        try:
+            listener = con.Listener(address, authkey=key)
+        except OSError:
+            logger.warning(
+                "Aikido listener may already be running on port %s", address[1]
+            )
+            pid = os.getpid()
+            os.kill(pid, signal.SIGTERM)  # Kill this subprocess
         self.queue = Queue()
         # Start reporting thread :
         Thread(target=self.reporting_thread).start()
@@ -35,12 +44,17 @@
             logger.debug("connection accepted from %s", listener.last_accepted)
             while True:
                 data = conn.recv()
-                logger.error(data)  # Temporary debugging
+                logger.debug("Incoming data : %s", data)
                 if data[0] == "ATTACK":
                     self.queue.put(data[1])
                 elif data[0] == "CLOSE":  # this is a kind of EOL for python IPC
                     conn.close()
                     break
+                elif data[0] == "KILL":
+                    logger.debug("Killing subprocess")
+                    conn.close()
+                    pid = os.getpid()
+                    os.kill(pid, signal.SIGTERM)  # Kill this subprocess
 
     def reporting_thread(self):
         """Reporting thread"""
@@ -73,6 +87,14 @@
     return ipc
 
 
+def reset_comms():
+    """This will reset communications"""
+    global ipc
+    if ipc:
+        ipc.send_data("KILL", {})
+        ipc = None
+
+
 def start_background_process():
     """
     Starts a process to handle incoming/outgoing data
@@ -92,35 +114,39 @@
     """
 
     def __init__(self, address, key):
+        # The key needs to be in byte form
         self.address = address
         self.key = key
-        self.background_process = None
 
     def start_aikido_listener(self):
-        """
-        This will start the aikido background process which listens
-        and makes calls to the API
-        """
-        self.background_process = Process(
-            target=AikidoBackgroundProcess,
-            args=(
-                self.address,
-                self.key,
-            ),
-        )
-        logger.debug("Starting the background process")
-        self.background_process.start()
+        """This will start the aikido process which listens"""
+        pid = os.fork()
+        if pid == 0:  # Child process
+            AikidoBackgroundProcess(self.address, self.key)
+        else:  # Parent process
+            logger.debug("Started background process, PID: %d", pid)
 
     def send_data(self, action, obj):
         """
         This creates a new client for comms to the background process
         """
-        try:
-            conn = con.Client(self.address, authkey=self.key)
-            logger.debug("Created connection %s", conn)
-            conn.send((action, obj))
-            conn.send(("CLOSE", {}))
-            conn.close()
-            logger.debug("Connection closed")
-        except Exception as e:
-            logger.info("Failed to send data to bg process : %s", e)
+
+        # We want to make sure that sending out this data affects the process as little as possible
+        # So we run it inside a seperate thread with a timeout of 3 seconds
+        def target(address, key, data_array):
+            try:
+                conn = con.Client(address, authkey=key)
+                logger.debug("Created connection %s", conn)
+                for data in data_array:
+                    conn.send(data)
+                conn.send(("CLOSE", {}))
+                conn.close()
+                logger.debug("Connection closed")
+            except Exception as e:
+                logger.info("Failed to send data to bg process : %s", e)
+
+        t = Thread(
+            target=target, args=(self.address, self.key, [(action, obj)]), daemon=True
+        )
+        t.start()
+        t.join(timeout=3)

aikido_firewall/background_process/init_test.py

@@ -4,6 +4,7 @@
     start_background_process,
     get_comms,
     IPC_ADDRESS,
+    reset_comms,
 )
 
 
@@ -13,18 +14,19 @@ def test_ipc_init():
     ipc = IPC(address, key)
 
     assert ipc.address == address
-    assert ipc.background_process is None
+    assert ipc.key == key
 
 
-# Following function does not work
 def test_start_background_process(monkeypatch):
+    reset_comms()
     assert get_comms() == None
     start_background_process()
 
     assert get_comms() != None
     assert get_comms().address == IPC_ADDRESS
 
-    get_comms().background_process.kill()
+    reset_comms()
+    assert get_comms() == None
 
 
 def test_send_data_exception(monkeypatch, caplog):
@@ -37,9 +39,6 @@ def mock_client(address, authkey):
     ipc = IPC(("localhost", 9898), "mock_key")
     ipc.send_data("ACTION", "Test Object")
 
-    assert "Failed to send data to bg" in caplog.text
-    # Add assertions for handling exceptions
-
 
 def test_send_data_successful(monkeypatch, caplog, mocker):
     ipc = IPC(("localhost"), "mock_key")

aikido_firewall/context/__init__.py

@@ -3,8 +3,11 @@
 """
 
 import threading
+from urllib.parse import parse_qs
+from http.cookies import SimpleCookie
 
-SUPPORTED_SOURCES = ["django", "flask"]
+
+SUPPORTED_SOURCES = ["django", "flask", "django-gunicorn"]
 UINPUT_SOURCES = ["body", "cookies", "query", "headers"]
 
 local = threading.local()
@@ -22,9 +25,26 @@
     """Parse EnvironHeaders object into a dict"""
     if isinstance(headers, dict):
         return headers
+    if isinstance(headers, list):
+        obj = {}
+        for k, v in headers:
+            obj[k] = v
+        return obj
     return dict(zip(headers.keys(), headers.values()))
 
 
+def parse_cookies(cookie_str):
+    """Parse cookie string from headers"""
+    cookie_dict = {}
+    cookies = SimpleCookie()
+    cookies.load(cookie_str)
+
+    for key, morsel in cookies.items():
+        cookie_dict[key] = morsel.value
+
+    return cookie_dict
+
+
 class Context:
     """
     A context object, it stores everything that is important
@@ -41,6 +61,17 @@
             self.set_flask_attrs(req)
         elif source == "django":
             self.set_django_attrs(req)
+        elif source == "django-gunicorn":
+            self.set_django_gunicorn_attrs(req)
+
+    def set_django_gunicorn_attrs(self, req):
+        """Set properties that are specific to django-gunicorn"""
+        self.remote_address = req.remote_addr
+        self.url = req.uri
+        self.body = parse_qs(req.body_copy.decode("utf-8"))
+        self.query = parse_qs(req.query)
+        self.cookies = parse_cookies(self.headers["COOKIE"])
+        del self.headers["COOKIE"]
 
     def set_django_attrs(self, req):
         """set properties that are specific to django"""

aikido_firewall/init_test.py

@@ -1,6 +1,6 @@
 import pytest
 from aikido_firewall import protect
-from aikido_firewall.background_process import get_comms
+from aikido_firewall.background_process import get_comms, reset_comms
 
 
 def test_protect_with_django(monkeypatch, caplog):
@@ -15,4 +15,5 @@ def test_protect_with_django(monkeypatch, caplog):
 
     assert "Aikido python firewall started" in caplog.text
     assert get_comms() != None
-    get_comms().background_process.kill()
+    reset_comms()
+    assert get_comms() == None

sample-apps/django-mysql-gunicorn/README.md

@@ -10,7 +10,7 @@ This will expose a Django web server at [localhost:8080](http://localhost:8080)
 
 ## URLS : 
 - Homepage : `http://localhost:8080/app`
-- Create a dog : `http://localhost:8080/app/create/<dog_name>`
-- MySQL attack : `http://localhost:8080/app/create/Malicious dog", "Injected wrong boss name"); --%20`
+- Create a dog : `http://localhost:8080/app/create/`
+- MySQL attack : Enter `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-mysql-gunicorn/docker-compose.yml

@@ -1,7 +1,7 @@
 version: '3'
 services:
   db:
-    image: mysql
+    image: mariadb
     container_name: django_mysql_gunicorn_mariadb
     restart: always
     volumes:
@@ -17,13 +17,18 @@ services:
     expose:
       # Opens port 3306 on the container
       - '3306'
+    healthcheck:
+      test: ["CMD", "healthcheck.sh", "--connect", "--innodb_initialized"]
+      interval: 5s
+      retries: 5
+      start_period: 10s
 
   backend:
     build: 
       context: ./../../
       dockerfile: ./sample-apps/django-mysql-gunicorn/Dockerfile
     container_name: django_mysql_gunicorn_backend
-    command: sh -c "python manage.py migrate && gunicorn --workers 4 --threads 2 --log-level debug --access-logfile '-' --error-logfile '-' --bind 0.0.0.0:8000 sample-django-mysql-gunicorn-app.wsgi"
+    command: sh -c "python manage.py migrate && gunicorn -c gunicorn_config.py --workers 4 --threads 2 --log-level debug --access-logfile '-' --error-logfile '-' --bind 0.0.0.0:8000 sample-django-mysql-gunicorn-app.wsgi"
     restart: always
     volumes:
       - .:/app
@@ -32,7 +37,8 @@ services:
     env_file:
       - .env
     depends_on:
-      - db
+      db:
+        condition: service_healthy
 
 volumes:
   db_data:

sample-apps/django-mysql-gunicorn/gunicorn_config.py

@@ -0,0 +1,39 @@
+import aikido_firewall
+import json
+from urllib.parse import parse_qs
+from io import BytesIO
+from aikido_firewall.context import Context
+from gunicorn.http.body import Body
+
+def when_ready(server):
+    aikido_firewall.protect("server-only")
+
+def pre_fork(server, worker):
+    pass
+
+def post_fork(server, worker):
+    print("----------------------> POST FORK")
+    import aikido_firewall
+    aikido_firewall.protect("django-gunicorn", False)
+
+def pre_request(worker, req):
+    req.body, req.body_copy = clone_body(req.body)
+
+    django_context = Context(req, "django-gunicorn")
+    django_context.set_as_current_context()
+
+    worker.log.debug("%s %s", req.method, req.path)
+
+
+def clone_body(body):
+    body_read = body.read()
+
+    # Read the body content into a buffer
+    body_buffer = BytesIO()
+    body_buffer.write(body_read)
+    body_buffer.seek(0)
+
+    # Create a new Body object with the same content
+    cloned_body = Body(body_buffer)
+
+    return (cloned_body, body_read)

sample-apps/django-mysql-gunicorn/manage.py

@@ -3,11 +3,8 @@
 import os
 import sys
 
-
 def main():
     """Run administrative tasks."""
-    import aikido_firewall # Aikido module
-    aikido_firewall.protect()
     os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'sample-django-mysql-gunicorn-app.settings')
     try:
         from django.core.management import execute_from_command_line

sample-apps/django-mysql-gunicorn/requirements.txt

@@ -1,5 +1,5 @@
 django
-pymysql
+mysqlclient
 python-decouple
 cryptography
 gunicorn

sample-apps/django-mysql-gunicorn/sample-django-mysql-gunicorn-app/__init__.py

@@ -1,2 +0,0 @@
-import pymysql
-pymysql.install_as_MySQLdb()

sample-apps/django-mysql-gunicorn/sample_app/templates/app/create_dog.html

@@ -0,0 +1,7 @@
+<h1>Create a New Dog</h1>
+<form method="post">
+    {% csrf_token %}
+    <label for="dog_name">Dog Name:</label>
+    <input type="text" id="dog_name" name="dog_name">
+    <button type="submit">Create Dog</button>
+</form>

sample-apps/django-mysql-gunicorn/sample_app/urls.py

@@ -5,5 +5,5 @@
 urlpatterns = [
     path("", views.index, name="index"),
     path("dogpage/<int:dog_id>", views.dog_page, name="dog_page"),
-    path("create/<dog_name>", views.create_dogpage, name="create"),
-]
+    path("create/", views.create_dogpage, name="create"),
+]