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

@@ -3,124 +3,24 @@
 and listen for data sent by our sources and sinks
 """
 
-import time
-import os
 import secrets
-import multiprocessing.connection as con
-from multiprocessing import Process
-from threading import Thread
-from queue import Queue
 from aikido_firewall.helpers.logging import logger
+from aikido_firewall.background_process.comms import (
+    AikidoIPCCommunications,
+    get_comms,
+    reset_comms,
+)
 
-REPORT_SEC_INTERVAL = 600  # 10 minutes
 IPC_ADDRESS = ("localhost", 9898)  # Specify the IP address and port
 
 
-class AikidoBackgroundProcess:
-    """
-    Aikido's background process consists of 2 threads :
-    - (main) Listening thread which listens on an IPC socket for incoming data
-    - (spawned) reporting thread which will collect the IPC data and send it to a Reporter
-    """
-
-    def __init__(self, address, key):
-        logger.debug("Background process started")
-        listener = con.Listener(address, authkey=key)
-        self.queue = Queue()
-        # Start reporting thread :
-        Thread(target=self.reporting_thread).start()
-
-        while True:
-            conn = listener.accept()
-            logger.debug("connection accepted from %s", listener.last_accepted)
-            while True:
-                data = conn.recv()
-                logger.error(data)  # Temporary debugging
-                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
-
-    def reporting_thread(self):
-        """Reporting thread"""
-        logger.debug("Started reporting thread")
-        while True:
-            self.send_to_reporter()
-            time.sleep(REPORT_SEC_INTERVAL)
-
-    def send_to_reporter(self):
-        """
-        Reports the found data to an Aikido server
-        """
-        items_to_report = []
-        while not self.queue.empty():
-            items_to_report.append(self.queue.get())
-        logger.debug("Reporting to aikido server")
-        logger.critical("Items to report : %s", items_to_report)
-        # Currently not making API calls
-
-
-# pylint: disable=invalid-name # This variable does change
-ipc = None
-
-
-def get_comms():
-    """
-    Returns the globally stored IPC object, which you need
-    to communicate to our background process.
-    """
-    return ipc
-
-
 def start_background_process():
     """
     Starts a process to handle incoming/outgoing data
     """
-    # pylint: disable=global-statement # We need this to be global
-    global ipc
+
     # Generate a secret key :
     generated_key_bytes = secrets.token_bytes(32)
 
-    ipc = IPC(IPC_ADDRESS, generated_key_bytes)
-    ipc.start_aikido_listener()
-
-
-class IPC:
-    """
-    Facilitates Inter-Process communication
-    """
-
-    def __init__(self, address, key):
-        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()
-
-    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)
+    comms = AikidoIPCCommunications(IPC_ADDRESS, generated_key_bytes)
+    comms.start_aikido_listener()

aikido_firewall/background_process/aikido_background_process.py

@@ -0,0 +1,72 @@
+"""
+Simply exports the aikido background process
+"""
+
+import multiprocessing.connection as con
+import os
+import time
+import signal
+from threading import Thread
+from queue import Queue
+from aikido_firewall.helpers.logging import logger
+
+REPORT_SEC_INTERVAL = 600  # 10 minutes
+
+
+class AikidoBackgroundProcess:
+    """
+    Aikido's background process consists of 2 threads :
+    - (main) Listening thread which listens on an IPC socket for incoming data
+    - (spawned) reporting thread which will collect the IPC data and send it to a Reporter
+    """
+
+    def __init__(self, address, key):
+        logger.debug("Background process started")
+        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()
+
+        while True:
+            conn = listener.accept()
+            logger.debug("connection accepted from %s", listener.last_accepted)
+            while True:
+                data = conn.recv()
+                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"
+                ):  # when main process quits , or during testing etc
+                    logger.debug("Killing subprocess")
+                    conn.close()
+                    pid = os.getpid()
+                    os.kill(pid, signal.SIGTERM)  # Kill this subprocess
+
+    def reporting_thread(self):
+        """Reporting thread"""
+        logger.debug("Started reporting thread")
+        while True:
+            self.send_to_reporter()
+            time.sleep(REPORT_SEC_INTERVAL)
+
+    def send_to_reporter(self):
+        """
+        Reports the found data to an Aikido server
+        """
+        items_to_report = []
+        while not self.queue.empty():
+            items_to_report.append(self.queue.get())
+        logger.debug("Reporting to aikido server")
+        logger.critical("Items to report : %s", items_to_report)
+        # Currently not making API calls

aikido_firewall/background_process/comms.py

@@ -0,0 +1,83 @@
+"""
+Holds the globally stored comms object
+Exports the AikidoIPCCommunications class
+"""
+
+import os
+import multiprocessing.connection as con
+from threading import Thread
+from aikido_firewall.helpers.logging import logger
+from aikido_firewall.background_process.aikido_background_process import (
+    AikidoBackgroundProcess,
+)
+
+# pylint: disable=invalid-name # This variable does change
+comms = None
+
+
+def get_comms():
+    """
+    Returns the globally stored IPC object, which you need
+    to communicate to our background process.
+    """
+    return comms
+
+
+def reset_comms():
+    """This will reset communications"""
+    # pylint: disable=global-statement # This needs to be global
+    global comms
+    if comms:
+        comms.send_data_to_bg_process("KILL", {})
+        comms = None
+
+
+class AikidoIPCCommunications:
+    """
+    Facilitates Inter-Process communication
+    """
+
+    def __init__(self, address, key):
+        # The key needs to be in byte form
+        self.address = address
+        self.key = key
+
+        # Set as global ipc object :
+        reset_comms()
+        # pylint: disable=global-statement # This needs to be global
+        global comms
+        comms = self
+
+    def start_aikido_listener(self):
+        """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_to_bg_process(self, action, obj):
+        """
+        This creates a new client for comms to the background process
+        """
+
+        # 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()
+        # This joins the thread for 3 seconds, afterwards the thread is forced to close (daemon=True)
+        t.join(timeout=3)

aikido_firewall/background_process/comms_test.py

@@ -0,0 +1,31 @@
+import pytest
+from aikido_firewall.background_process.comms import AikidoIPCCommunications
+
+
+def test_comms_init():
+    address = ("localhost", 9898)
+    key = "secret_key"
+    comms = AikidoIPCCommunications(address, key)
+
+    assert comms.address == address
+    assert comms.key == key
+
+
+def test_send_data_to_bg_process_exception(monkeypatch, caplog):
+    def mock_client(address, authkey):
+        raise Exception("Connection Error")
+
+    monkeypatch.setitem(globals(), "Client", mock_client)
+    monkeypatch.setitem(globals(), "logger", caplog)
+
+    comms = AikidoIPCCommunications(("localhost", 9898), "mock_key")
+    comms.send_data_to_bg_process("ACTION", "Test Object")
+
+
+def test_send_data_to_bg_process_successful(monkeypatch, caplog, mocker):
+    comms = AikidoIPCCommunications(("localhost"), "mock_key")
+    mock_client = mocker.MagicMock()
+    monkeypatch.setattr("multiprocessing.connection.Client", mock_client)
+
+    # Call the send_data_to_bg_process function
+    comms.send_data_to_bg_process("ACTION", {"key": "value"})