profiler: shutdown queue gracefully

Author: AlyaGomaa

slips_files/common/abstracts/core.py

@@ -39,7 +39,11 @@ def run(self):
             self.shutdown_gracefully()
 
         except KeyboardInterrupt:
+            self.keyboard_int_ctr += 1
+            if self.keyboard_int_ctr >= 2:
+                return
             self.shutdown_gracefully()
+
         except Exception:
             self.print(f"Problem in {self.name}", 0, 1)
             self.print(traceback.format_exc(), 0, 1)

slips_files/core/profiler.py

@@ -81,7 +81,7 @@ def init(
         # the input process signals the rest of the modules to stop
         self.done_processing: multiprocessing.Semaphore = is_profiler_done
         # every line put in this queue should be profiled
-        self.profiler_queue = profiler_queue
+        self.profiler_queue: multiprocessing.Queue = profiler_queue
         self.timeformat = None
         self.input_type = False
         self.rec_lines = 0
@@ -230,10 +230,10 @@ def get_gateway_info(self, flow):
             ):
                 self.gw_mac: str = flow.dmac
                 self.db.set_default_gateway("MAC", self.gw_mac)
-                self.print(
-                    f"MAC address of the gateway detected: "
-                    f"{green(self.gw_mac)}"
-                )
+                # self.print(
+                #     f"MAC address of the gateway detected: "
+                #     f"{green(self.gw_mac)}"
+                # )
                 gw_mac_found = True
 
         # we need the mac to be set to be able to find the ip using it
@@ -664,6 +664,12 @@ def should_stop(self):
         return False
 
     def shutdown_gracefully(self):
+        # signal the threads to stop
+        self.stop_profiler_threads.set()
+        self.flows_to_process_q.close()
+        self.profiler_queue.close()
+        # wait for them to finish
+        self.join_profiler_threads()
         self.db.set_new_incoming_flows(False)
         self.print(
             f"Stopping. Total lines read: {self.rec_lines}",
@@ -701,13 +707,8 @@ def main(self):
             # without it, there's no way this module will know it's
             # time to stop and no new flows are coming
             if self.is_stop_msg(msg):
-                # DO NOT return/exit this module before all profilers are
-                # done. if you do, the profiler threads will shutdown
-                # before reading all flows as soon as we receive the stop msg.
-                # signal the threads to stop
-                self.stop_profiler_threads.set()
-                # wait for them to finish
-                self.join_profiler_threads()
+                # shutdown gracefully will be called by icore once this
+                # function returns
                 return 1
 
             self.pending_flows_queue_lock.acquire()