diff --git a/aikido_firewall/background_process/aikido_background_process.py b/aikido_firewall/background_process/aikido_background_process.py index 9c9c2fef3..4f4505180 100644 --- a/aikido_firewall/background_process/aikido_background_process.py +++ b/aikido_firewall/background_process/aikido_background_process.py @@ -44,33 +44,45 @@ def __init__(self, address, key): conn = listener.accept() logger.debug("connection accepted from %s", listener.last_accepted) while True: - data = conn.recv() # because of this no sleep needed in thread - logger.debug("Incoming data : %s", data) - 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 - elif ( - data[0] == "KILL" - ): # when main process quits , or during testing etc - logger.debug("Killing subprocess") - conn.close() - sys.exit(0) - elif data[0] == "READ_PROPERTY": # meant to get config props - if hasattr(self.reporter, data[1]): - conn.send(self.reporter.__dict__[data[1]]) - elif data[0] == "ROUTE": - # Called every time the user visits a route - self.reporter.routes.add_route(method=data[1][0], path=data[1][1]) - elif data[0] == "SHOULD_RATELIMIT": - # Called to check if the context passed along as data should be - # Rate limited - conn.send( - should_ratelimit_request( - context=data[1], reporter=self.reporter + try: + data = conn.recv() # because of this no sleep needed in thread + logger.debug("Incoming data : %s", data) + 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 + elif ( + data[0] == "KILL" + ): # when main process quits , or during testing etc + logger.debug("Killing subprocess") + conn.close() + sys.exit(0) + elif data[0] == "READ_PROPERTY": # meant to get config props + if hasattr(self.reporter, data[1]): + conn.send(self.reporter.__dict__[data[1]]) + else: + logger.debug( + "Reporter has no attribute %s, current reporter: %s", + data[1], + self.reporter, + ) + conn.send(None) + elif data[0] == "ROUTE": + # Called every time the user visits a route + self.reporter.routes.add_route( + method=data[1][0], path=data[1][1] ) - ) + elif data[0] == "SHOULD_RATELIMIT": + # Called to check if the context passed along as data should be + # Rate limited + conn.send( + should_ratelimit_request( + context=data[1], reporter=self.reporter + ) + ) + except Exception as e: + logger.error("Exception occured in server thread : %s", e) def reporting_thread(self): """Reporting thread"""