From 8b166f7dfcaa8afc57b255418d1dfc0ee387184b Mon Sep 17 00:00:00 2001
From: Wout Feys <me@woutfeys.be>
Date: Tue, 30 Jul 2024 11:03:38 +0200
Subject: [PATCH] Move gunicorn code to a file so our gunicorn config file is
 smaller

---
 aikido_firewall/middleware/django_gunicorn.py | 67 +++++++++++++++++++
 .../django-mysql-gunicorn/gunicorn_config.py  | 46 ++-----------
 2 files changed, 74 insertions(+), 39 deletions(-)
 create mode 100644 aikido_firewall/middleware/django_gunicorn.py

diff --git a/aikido_firewall/middleware/django_gunicorn.py b/aikido_firewall/middleware/django_gunicorn.py
new file mode 100644
index 000000000..2a1d36111
--- /dev/null
+++ b/aikido_firewall/middleware/django_gunicorn.py
@@ -0,0 +1,67 @@
+"""
+Includes all the wrappers for gunicorn config file
+"""
+
+from gunicorn.http.body import Body
+from io import BytesIO
+import aikido_firewall
+from aikido_firewall.context import Context
+
+
+def when_ready(prev_func):
+    """
+    Aikido decorator for gunicorn config
+    Function: pre_request(worker, req)
+    """
+
+    def aik_when_ready(server):
+        aikido_firewall.protect("server-only")
+        prev_func(server)
+
+    return aik_when_ready
+
+
+def pre_request(prev_func):
+    """
+    Aikido decorator for gunicorn config
+    Function: pre_request(worker, req)
+    """
+
+    def aik_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()
+        prev_func(worker, req)
+
+    return aik_pre_request
+
+
+def post_fork(prev_func):
+    """
+    Aikido decorator for gunicorn config
+    Function: post_fork(server, worker)
+    """
+
+    def aik_post_fork(server, worker):
+        aikido_firewall.protect("django-gunicorn", False)
+        prev_func(server, worker)
+
+    return aik_post_fork
+
+
+def clone_body(body):
+    """
+    Clones the body by creating a new stream
+    """
+    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)
diff --git a/sample-apps/django-mysql-gunicorn/gunicorn_config.py b/sample-apps/django-mysql-gunicorn/gunicorn_config.py
index c722b0f04..75761e7e0 100644
--- a/sample-apps/django-mysql-gunicorn/gunicorn_config.py
+++ b/sample-apps/django-mysql-gunicorn/gunicorn_config.py
@@ -1,39 +1,7 @@
-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)
+import aikido_firewall.middleware.django_gunicorn as aik
+@aik.when_ready
+def when_ready(server): pass
+@aik.post_fork
+def post_fork(server, worker): pass
+@aik.pre_request
+def pre_request(worker, req): pass