diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 56303f56..2afe5ab4 100644 --- a/aikido_firewall/__init__.py +++ b/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) diff --git a/aikido_firewall/background_process/__init__.py b/aikido_firewall/background_process/__init__.py index 48f26e3d..6a9b4aa0 100644 --- a/aikido_firewall/background_process/__init__.py +++ b/aikido_firewall/background_process/__init__.py @@ -26,7 +26,9 @@ def start_background_process(): # 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() diff --git a/aikido_firewall/background_process/comms.py b/aikido_firewall/background_process/comms.py index ce86cf5c..a32f8fea 100644 --- a/aikido_firewall/background_process/comms.py +++ b/aikido_firewall/background_process/comms.py @@ -29,7 +29,7 @@ def reset_comms(): # 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 @@ def start_aikido_listener(self): 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 """ diff --git a/aikido_firewall/init_test.py b/aikido_firewall/init_test.py index 54072447..6fb9b24c 100644 --- a/aikido_firewall/init_test.py +++ b/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 diff --git a/aikido_firewall/sinks/psycopg2.py b/aikido_firewall/sinks/psycopg2.py index 83b49fb0..2fa6c1f7 100644 --- a/aikido_firewall/sinks/psycopg2.py +++ b/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)