diff --git a/aikido_firewall/background_process/comms.py b/aikido_firewall/background_process/comms.py index 2c08d4be4..aa709fe81 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) + # 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 : + 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, # This allows us to join and set a timeout after which the thread closes ) - 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 + # Start and join the thread for 100ms, afterwards the thread is forced to close (daemon=True) + t.start() + t.join(timeout=0.1) + return result_obj diff --git a/aikido_firewall/sinks/mysqlclient.py b/aikido_firewall/sinks/mysqlclient.py index 33fcbd3c5..246a8b175 100644 --- a/aikido_firewall/sinks/mysqlclient.py +++ b/aikido_firewall/sinks/mysqlclient.py @@ -39,7 +39,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( + action="READ_PROPERTY", obj="block", receive=True + ) if should_block: raise AikidoSQLInjection("SQL Injection [aikido_firewall]") diff --git a/aikido_firewall/sinks/pymongo.py b/aikido_firewall/sinks/pymongo.py index c37bb30bb..66579b319 100644 --- a/aikido_firewall/sinks/pymongo.py +++ b/aikido_firewall/sinks/pymongo.py @@ -52,7 +52,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 AikidoNoSQLInjection("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 aded8e574..2422cd655 100644 --- a/aikido_firewall/sinks/pymysql.py +++ b/aikido_firewall/sinks/pymysql.py @@ -41,7 +41,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( + action="READ_PROPERTY", obj="block", receive=True + ) if should_block: raise AikidoSQLInjection("SQL Injection [aikido_firewall]")