From 4fc5f0a96dcfe81a9c7c1c965d434f9fdd047c21 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Thu, 1 Aug 2024 09:20:53 +0200 Subject: [PATCH 1/3] Merge IPC into one function --- aikido_firewall/background_process/comms.py | 48 ++++++++++----------- aikido_firewall/sinks/mysqlclient.py | 4 +- aikido_firewall/sinks/pymongo.py | 4 +- aikido_firewall/sinks/pymysql.py | 4 +- 4 files changed, 31 insertions(+), 29 deletions(-) diff --git a/aikido_firewall/background_process/comms.py b/aikido_firewall/background_process/comms.py index 2c08d4be4..fd146a2d8 100644 --- a/aikido_firewall/background_process/comms.py +++ b/aikido_firewall/background_process/comms.py @@ -56,40 +56,36 @@ def start_aikido_listener(self): else: # Parent process logger.debug("Started background process, PID: %d", pid) - def send_data_to_bg_process(self, action, obj): + def send_data_to_bg_process(self, action, obj, receive=False): """ 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) + # If something goes wrong, it will also be encapsulated in the thread i.e. no crashes + def target(address, key, receive, data, result_obj): + # Create a connection, this can get stuck : + conn = con.Client(address, authkey=key) + # Send/Receive data : + conn.send(data) + if receive: + result_obj = conn.recv() + + # Close the connection : + conn.send(("CLOSE", {})) + conn.close() + + # Create a shared result object between the thread and this process : + result_obj = None t = Thread( - target=target, args=(self.address, self.key, [(action, obj)]), daemon=True + target=target, + args=(self.address, self.key, receive, (action, obj), result_obj), + daemon=True, ) + + # Start and join the thread for 3 seconds, afterwards the thread is forced to close (daemon=True) t.start() - # This joins the thread for 3 seconds, afterwards the thread is forced to close (daemon=True) t.join(timeout=3) - - def poll_config(self, prop): - """ - This will poll the config from the Background Process - """ - conn = con.Client(self.address, authkey=self.key) - conn.send(("READ_PROPERTY", prop)) - prop_value = conn.recv() - conn.send(("CLOSE", {})) - conn.close() - logger.debug("Received property %s as %s", prop, prop_value) - return prop_value + return result_obj diff --git a/aikido_firewall/sinks/mysqlclient.py b/aikido_firewall/sinks/mysqlclient.py index a8d999c30..e41b05e12 100644 --- a/aikido_firewall/sinks/mysqlclient.py +++ b/aikido_firewall/sinks/mysqlclient.py @@ -38,7 +38,9 @@ def aikido_new_query(_self, sql): logger.debug("sql_injection results : %s", json.dumps(contains_injection)) if contains_injection: get_comms().send_data_to_bg_process("ATTACK", (contains_injection, context)) - should_block = get_comms().poll_config("block") + should_block = get_comms().send_data_to_bg_process( + "READ_PROPERTY", "block", True + ) if should_block: raise Exception("SQL Injection [aikido_firewall]") diff --git a/aikido_firewall/sinks/pymongo.py b/aikido_firewall/sinks/pymongo.py index 9b41f1a83..760454839 100644 --- a/aikido_firewall/sinks/pymongo.py +++ b/aikido_firewall/sinks/pymongo.py @@ -51,7 +51,9 @@ def wrapped_operation_function( context = get_current_context() injection_results = detect_nosql_injection(context, _filter) if injection_results["injection"]: - get_comms().send_data("ATTACK", injection_results) + get_comms().send_data_to_bg_process( + "ATTACK", (injection_results, context) + ) raise Exception("NOSQL Injection [aikido_firewall]") return prev_func(_self, _filter, *args, **kwargs) diff --git a/aikido_firewall/sinks/pymysql.py b/aikido_firewall/sinks/pymysql.py index 88f5e1353..80c3ca051 100644 --- a/aikido_firewall/sinks/pymysql.py +++ b/aikido_firewall/sinks/pymysql.py @@ -40,7 +40,9 @@ def aikido_new_query(_self, sql, unbuffered=False): logger.info("sql_injection results : %s", json.dumps(contains_injection)) if contains_injection: get_comms().send_data_to_bg_process("ATTACK", (contains_injection, context)) - should_block = get_comms().poll_config("block") + should_block = get_comms().send_data_to_bg_process( + "READ_PROPERTY", "block", True + ) if should_block: raise Exception("SQL Injection [aikido_firewall]") From 00009d5fd0bb5acd5967591a11f43a9962eccb95 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 2 Aug 2024 14:42:39 +0200 Subject: [PATCH 2/3] Parametrized receive=True --- aikido_firewall/sinks/mysqlclient.py | 2 +- aikido_firewall/sinks/pymysql.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/aikido_firewall/sinks/mysqlclient.py b/aikido_firewall/sinks/mysqlclient.py index 9663a093d..246a8b175 100644 --- a/aikido_firewall/sinks/mysqlclient.py +++ b/aikido_firewall/sinks/mysqlclient.py @@ -40,7 +40,7 @@ def aikido_new_query(_self, sql): if contains_injection: get_comms().send_data_to_bg_process("ATTACK", (contains_injection, context)) should_block = get_comms().send_data_to_bg_process( - "READ_PROPERTY", "block", True + action="READ_PROPERTY", obj="block", receive=True ) if should_block: raise AikidoSQLInjection("SQL Injection [aikido_firewall]") diff --git a/aikido_firewall/sinks/pymysql.py b/aikido_firewall/sinks/pymysql.py index 084828069..2422cd655 100644 --- a/aikido_firewall/sinks/pymysql.py +++ b/aikido_firewall/sinks/pymysql.py @@ -42,7 +42,7 @@ def aikido_new_query(_self, sql, unbuffered=False): if contains_injection: get_comms().send_data_to_bg_process("ATTACK", (contains_injection, context)) should_block = get_comms().send_data_to_bg_process( - "READ_PROPERTY", "block", True + action="READ_PROPERTY", obj="block", receive=True ) if should_block: raise AikidoSQLInjection("SQL Injection [aikido_firewall]") From d0ee370fb149b0d1e4c87099b7194ba30766ee51 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Fri, 2 Aug 2024 14:42:51 +0200 Subject: [PATCH 3/3] Timeout to 100ms and add extra comments for Daemon=true --- aikido_firewall/background_process/comms.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/aikido_firewall/background_process/comms.py b/aikido_firewall/background_process/comms.py index fd146a2d8..aa709fe81 100644 --- a/aikido_firewall/background_process/comms.py +++ b/aikido_firewall/background_process/comms.py @@ -62,7 +62,7 @@ def send_data_to_bg_process(self, action, obj, receive=False): """ # 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 + # So we run it inside a seperate thread with a timeout of 100ms # If something goes wrong, it will also be encapsulated in the thread i.e. no crashes def target(address, key, receive, data, result_obj): # Create a connection, this can get stuck : @@ -82,10 +82,10 @@ def target(address, key, receive, data, result_obj): t = Thread( target=target, args=(self.address, self.key, receive, (action, obj), result_obj), - daemon=True, + daemon=True, # This allows us to join and set a timeout after which the thread closes ) - # Start and join the thread for 3 seconds, afterwards the thread is forced to close (daemon=True) + # Start and join the thread for 100ms, afterwards the thread is forced to close (daemon=True) t.start() - t.join(timeout=3) + t.join(timeout=0.1) return result_obj