AIK-3167 : Switch the python firewall agent to an IPC socket

README.md

@@ -1 +1,5 @@
-# firewall-python
+# firewall-python
+
+## Environment variables
+- `AIKIDO_SECRET_KEY` : Secret to encrypt IPC comms
+- `AIKIDO_DEBUG` : Boolean value to enable debug logs

aikido_firewall/__init__.py

@@ -9,7 +9,7 @@
 from aikido_firewall.helpers.logging import logger
 
 # Import agent
-from aikido_firewall.agent import start_agent
+from aikido_firewall.agent import start_ipc
 
 # Load environment variables
 load_dotenv()
@@ -27,4 +27,4 @@ def protect(module="any"):
     import aikido_firewall.sinks.pymysql
 
     logger.info("Aikido python firewall started")
-    start_agent()
+    start_ipc()

aikido_firewall/agent/__init__.py

@@ -3,72 +3,110 @@
 """
 
 import time
-import queue
+import os
+import multiprocessing.connection as con
+from multiprocessing import Process
 from threading import Thread
+from queue import Queue
 from aikido_firewall.helpers.logging import logger
 
-AGENT_SEC_INTERVAL = 60
+AGENT_SEC_INTERVAL = 5
+IPC_ADDRESS = ("localhost", 9898)  # Specify the IP address and port
 
 
-class AikidoThread:
+class AikidoProc:
     """
     Our agent thread
     """
 
-    def __init__(self, q):
+    def __init__(self, address, key):
         logger.debug("Agent thread 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] == "SQL_INJECTION":
+                    self.queue.put(data[1])
+                elif data[0] == "CLOSE":
+                    conn.close()
+                    break
+
+    def reporting_thread(self):
+        """Reporting thread"""
+        logger.debug("Started reporting thread")
         while True:
-            while not q.empty():
-                self.process_data(q.get())
-        time.sleep(AGENT_SEC_INTERVAL)
-        self.q = q
-        self.current_context = None
-
-    def process_data(self, item):
-        """Will process the data added to the queue"""
-        action, data = item
-        logger.debug("Action %s, Data %s", action, data)
-        if action == "REPORT":
-            logger.debug("Report")
-            self.current_context = data
-        else:
-            logger.error("Action `%s` is not defined. (Aikido Agent)", action)
+            self.report_to_agent()
+            time.sleep(AGENT_SEC_INTERVAL)
+
+    def report_to_agent(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
-agent = None
+ipc = None
 
 
-def get_agent():
+def get_ipc():
     """Returns the globally stored agent"""
-    return agent
+    return ipc
 
 
-def start_agent():
+def start_ipc():
     """
     Starts a thread to handle incoming/outgoing data
     """
     # pylint: disable=global-statement # We need this to be global
-    global agent
+    global ipc
+    aikido_secret_key_env = os.getenv("AIKIDO_SECRET_KEY")
+    if not aikido_secret_key_env:
+        raise EnvironmentError("AIKIDO_SECRET_KEY is not set.")
+    ipc = IPC(IPC_ADDRESS, aikido_secret_key_env)
+    ipc.start_aikido_listener()
 
-    # This creates a queue for Inter-Process Communication
-    logger.debug("Creating IPC Queue")
-    q = queue.Queue()
 
-    logger.debug("Starting a new agent thread")
-    agent_thread = Thread(target=AikidoThread, args=(q,))
-    agent_thread.start()
-    agent = Agent(q)
-
-
-class Agent:
+class IPC:
     """Agent class"""
 
-    def __init__(self, q):
-        self.q = q
-
-    def report(self, obj, action):
-        """
-        Report something to the agent
-        """
-        self.q.put((action, obj))
+    def __init__(self, address, key):
+        self.address = address
+        self.key = str.encode(key)
+        self.agent_proc = None
+
+    def start_aikido_listener(self):
+        """This will start the aikido thread which listens"""
+        self.agent_proc = Process(
+            target=AikidoProc,
+            args=(
+                self.address,
+                self.key,
+            ),
+        )
+        logger.debug("Starting a new agent thread")
+        self.agent_proc.start()
+
+    def send_data(self, action, obj):
+        """This creates a new client for comms to the thread"""
+        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 agent %s", e)

aikido_firewall/agent/init_test.py

@@ -0,0 +1,58 @@
+import pytest
+from aikido_firewall.agent import IPC, start_ipc, get_ipc, IPC_ADDRESS
+
+
+def test_ipc_init():
+    address = ("localhost", 9898)
+    key = "secret_key"
+    ipc = IPC(address, key)
+
+    assert ipc.address == address
+    assert ipc.key == b"secret_key"  # Ensure key is encoded as bytes
+    assert ipc.agent_proc is None
+
+
+def test_start_ipc_missing_secret_key(mocker):
+    mocker.patch("os.environ", {})
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        start_ipc()
+
+    assert "AIKIDO_SECRET_KEY is not set." in str(exc_info.value)
+
+
+# Following function does not work
+def test_start_ipc(monkeypatch):
+    assert get_ipc() == None
+    monkeypatch.setenv("AIKIDO_SECRET_KEY", "mock_key")
+
+    start_ipc()
+
+    assert get_ipc() != None
+    assert get_ipc().address == IPC_ADDRESS
+    assert get_ipc().key == b"mock_key"
+
+    get_ipc().agent_proc.kill()
+
+
+def test_send_data_exception(monkeypatch, caplog):
+    def mock_client(address, authkey):
+        raise Exception("Connection Error")
+
+    monkeypatch.setitem(globals(), "Client", mock_client)
+    monkeypatch.setitem(globals(), "logger", caplog)
+
+    ipc = IPC(("localhost", 9898), "mock_key")
+    ipc.send_data("ACTION", "Test Object")
+
+    assert "Failed to send data to agent" in caplog.text
+    # Add assertions for handling exceptions
+
+
+def test_send_data_successful(monkeypatch, caplog, mocker):
+    ipc = IPC(("localhost"), "mock_key")
+    mock_client = mocker.MagicMock()
+    monkeypatch.setattr("multiprocessing.connection.Client", mock_client)
+
+    # Call the send_data function
+    ipc.send_data("ACTION", {"key": "value"})

aikido_firewall/init_test.py

@@ -0,0 +1,19 @@
+import pytest
+from aikido_firewall import protect
+from aikido_firewall.agent import get_ipc
+
+
+def test_protect_with_django(monkeypatch, caplog):
+    monkeypatch.setitem(
+        globals(), "aikido_firewall.sources.django", "dummy_django_module"
+    )
+    monkeypatch.setitem(
+        globals(), "aikido_firewall.sinks.pymysql", "dummy_pymysql_module"
+    )
+    monkeypatch.setenv("AIKIDO_SECRET_KEY", "mock_key")
+
+    protect(module="django")
+
+    assert "Aikido python firewall started" in caplog.text
+    assert get_ipc() != None
+    get_ipc().agent_proc.kill()

aikido_firewall/sinks/pymysql.py

@@ -12,6 +12,7 @@
     check_context_for_sql_injection,
 )
 from aikido_firewall.vulnerabilities.sql_injection.dialects import MySQL
+from aikido_firewall.agent import get_ipc
 
 logger = logging.getLogger("aikido_firewall")
 
@@ -36,6 +37,7 @@
 
         logger.info("sql_injection results : %s", json.dumps(result))
         if result:
+            get_ipc().send_data("SQL_INJECTION", result)
             raise Exception("SQL Injection [aikido_firewall]")
         return prev_query_function(_self, sql, unbuffered=False)
 

sample-apps/django-mysql-gunicorn/.env.example

@@ -3,4 +3,7 @@ MYSQL_DATABASE="db"
 MYSQL_USER="user"
 MYSQL_PASSWORD="password"
 MYSQL_ROOT_PASSWORD="password"
+
+# Aikido keys
 AIKIDO_DEBUG=true
+AIKIDO_SECRET_KEY="your_secret_key"

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

@@ -3,4 +3,7 @@ MYSQL_DATABASE="db"
 MYSQL_USER="user"
 MYSQL_PASSWORD="password"
 MYSQL_ROOT_PASSWORD="password"
+
+# Aikido keys
 AIKIDO_DEBUG=true
+AIKIDO_SECRET_KEY="your_secret_key"

sample-apps/flask-mysql-uwsgi/.env.example

@@ -1 +1,2 @@
 AIKIDO_DEBUG=true
+AIKIDO_SECRET_KEY="your_secret_key"

sample-apps/flask-mysql/.env.example

@@ -1 +1,2 @@
 AIKIDO_DEBUG=true
+AIKIDO_SECRET_KEY="your_secret_key"