Fly io bugfixes and improvements

aikido_firewall/__init__.py

@@ -13,7 +13,10 @@
 # Import background process
 from aikido_firewall.background_process import start_background_process
 
-# Load environment variables
+# Load environment variables and constants
+# Load environment variables and constants
+from aikido_firewall.config import PKG_VERSION
+
 load_dotenv()
 
 
@@ -51,4 +54,4 @@ def protect(module="any", server=True):
     import aikido_firewall.sinks.os_system
     import aikido_firewall.sinks.subprocess
 
-    logger.info("Aikido python firewall started")
+    logger.info("Aikido python firewall v%s starting.", PKG_VERSION)

aikido_firewall/background_process/__init__.py

@@ -26,7 +26,9 @@
 
     # Remove the socket file if it already exists
     if os.path.exists(IPC_ADDRESS):
+        logger.debug("Unix Domain Socket file already exists, deleting.")
         os.remove(IPC_ADDRESS)
 
+    logger.debug("Communication starting on UDS File : %s", IPC_ADDRESS)
     comms = AikidoIPCCommunications(IPC_ADDRESS, secret_key_bytes)
     comms.start_aikido_listener()

aikido_firewall/background_process/comms.py

@@ -29,7 +29,7 @@
     # pylint: disable=global-statement # This needs to be global
     global comms
     if comms:
-        comms.send_data_to_bg_process("KILL", {})
+        logger.debug("Resetting communications. (comms = None)")
         comms = None
 
 
@@ -68,6 +68,14 @@
         self.background_process.start()
 
     def send_data_to_bg_process(self, action, obj, receive=False):
+        """Try-catched send_data_to_bg_process"""
+        try:
+            return self._send_data_to_bg_process(action, obj, receive=False)
+        except Exception as e:
+            logger.debug("Exception happened in send_data_to_bg_process : %s", e)
+            return {"success": False, "error": "unknown"}
+
+    def _send_data_to_bg_process(self, action, obj, receive=False):
         """
         This creates a new client for comms to the background process
         """

aikido_firewall/init_test.py

@@ -13,7 +13,7 @@ def test_protect_with_django(monkeypatch, caplog):
 
     protect(module="django")
 
-    assert "Aikido python firewall started" in caplog.text
+    assert "starting" in caplog.text
     assert get_comms() != None
     reset_comms()
     assert get_comms() == None

aikido_firewall/sinks/psycopg2.py

@@ -44,7 +44,7 @@ def __getattr__(self, name):
         def execute(*args, **kwargs):
             run_vulnerability_scan(
                 kind="sql_injection",
-                op="pymysql.connection.cursor.execute",
+                op="psycopg2.Connection.Cursor.execute",
                 args=(args[0], Postgres()),  #  args[0] : sql
             )
             return self._execute_func_copy(*args, **kwargs)
@@ -53,7 +53,7 @@ def executemany(*args, **kwargs):
             for sql in args[0]:
                 run_vulnerability_scan(
                     kind="sql_injection",
-                    op="pymysql.connection.cursor.executemany",
+                    op="psycopg2.Connection.Cursor.executemany",
                     args=(sql, Postgres()),
                 )
             return self._executemany_func_copy(*args, **kwargs)