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

aikido_firewall/__init__.py

@@ -8,15 +8,17 @@
 # Import logger
 from aikido_firewall.helpers.logging import logger
 
-# Import agent
-from aikido_firewall.agent import start_agent
+# Import background process
+from aikido_firewall.background_process import start_background_process
 
 # Load environment variables
 load_dotenv()
 
 
 def protect(module="any"):
-    """Start Aikido agent"""
+    """
+    Protect user's application
+    """
     # Import sources
     import aikido_firewall.sources.django
 
@@ -27,4 +29,4 @@ def protect(module="any"):
     import aikido_firewall.sinks.pymysql
 
     logger.info("Aikido python firewall started")
-    start_agent()
+    start_background_process()

aikido_firewall/background_process/__init__.py

@@ -0,0 +1,126 @@
+"""
+Aikido background process, this will create a new process
+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
+
+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)

aikido_firewall/background_process/init_test.py

@@ -0,0 +1,50 @@
+import pytest
+from aikido_firewall.background_process import (
+    IPC,
+    start_background_process,
+    get_comms,
+    IPC_ADDRESS,
+)
+
+
+def test_ipc_init():
+    address = ("localhost", 9898)
+    key = "secret_key"
+    ipc = IPC(address, key)
+
+    assert ipc.address == address
+    assert ipc.background_process is None
+
+
+# Following function does not work
+def test_start_background_process(monkeypatch):
+    assert get_comms() == None
+    start_background_process()
+
+    assert get_comms() != None
+    assert get_comms().address == IPC_ADDRESS
+
+    get_comms().background_process.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 bg" 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,18 @@
+import pytest
+from aikido_firewall import protect
+from aikido_firewall.background_process import get_comms
+
+
+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"
+    )
+
+    protect(module="django")
+
+    assert "Aikido python firewall started" in caplog.text
+    assert get_comms() != None
+    get_comms().background_process.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.background_process import get_comms
 
 logger = logging.getLogger("aikido_firewall")
 
@@ -36,6 +37,7 @@
 
         logger.info("sql_injection results : %s", json.dumps(result))
         if result:
+            get_comms().send_data("ATTACK", 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,6 @@ MYSQL_DATABASE="db"
 MYSQL_USER="user"
 MYSQL_PASSWORD="password"
 MYSQL_ROOT_PASSWORD="password"
+
+# Aikido keys
 AIKIDO_DEBUG=true

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

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

sample-apps/flask-mongo/app.py

@@ -8,7 +8,7 @@
 
 app = Flask(__name__)
 if __name__ == '__main__':
-    app.run(threaded=True) # Run threaded so we can test our agent's capabilities
+    app.run(threaded=True) # Run threaded so we can test how our bg process works
 app.config["MONGO_URI"] = "mongodb://admin:password@db:27017/my_database?authSource=admin"
 mongo = PyMongo(app)
 

sample-apps/flask-postgres/app.py

@@ -6,7 +6,7 @@
 
 app = Flask(__name__)
 if __name__ == '__main__':
-    app.run(threaded=True) # Run threaded so we can test our agent's capabilities
+    app.run(threaded=True) # Run threaded so we can test how our bg process works
 
 def get_db_connection():
     return psycopg2.connect(