Merge pull request #909 from stratosphereips/alya/fix-modules-stoppin…

…g-before-termination-event

Improve how Slips stops

modules/rnn_cc_detection/rnn_cc_detection.py

@@ -2,7 +2,6 @@
 import json
 from typing import Dict
 import numpy as np
-from typing import Optional
 from tensorflow.keras.models import load_model
 
 from slips_files.common.slips_utils import utils
@@ -186,23 +185,21 @@ def handle_new_letters(self, msg: Dict):
 
         if "tcp" not in tupleid.lower():
             return
-        
+
         if "established" not in state.lower():
             return
 
         # to reduce false positives
         threshold = 0.99
         # function to convert each letter of behavioral model to ascii
-        behavioral_model = self.convert_input_for_module(
-            pre_behavioral_model
-        )
+        behavioral_model = self.convert_input_for_module(pre_behavioral_model)
         # predict the score of behavioral model being c&c channel
         self.print(
             f"predicting the sequence: {pre_behavioral_model}",
             3,
             0,
         )
-        score = self.tcpmodel.predict(behavioral_model, verbose = 0)
+        score = self.tcpmodel.predict(behavioral_model, verbose=0)
         self.print(
             f" >> sequence: {pre_behavioral_model}. "
             f"final prediction score: {score[0][0]:.20f}",
@@ -216,9 +213,7 @@ def handle_new_letters(self, msg: Dict):
             if len(pre_behavioral_model) >= threshold_confidence:
                 confidence = 1
             else:
-                confidence = (
-                    len(pre_behavioral_model) / threshold_confidence
-                )
+                confidence = len(pre_behavioral_model) / threshold_confidence
             uid = msg["uid"]
             stime = flow["starttime"]
             self.set_evidence_cc_channel(
@@ -240,7 +235,6 @@ def handle_new_letters(self, msg: Dict):
             # detection
             self.db.publish("check_jarm_hash", json.dumps(to_send))
 
-
     def handle_tw_closed(self, msg: Dict):
         """handles msgs from the tw_closed channel"""
         profileid_tw = msg["data"].split("_")
@@ -263,5 +257,6 @@ def pre_main(self):
     def main(self):
         if msg := self.get_msg("new_letters"):
             self.handle_new_letters(msg)
-        elif msg := self.get_msg("tw_closed"):
+
+        if msg := self.get_msg("tw_closed"):
             self.handle_tw_closed(msg)

slips_files/common/abstracts/module.py

@@ -2,7 +2,10 @@
 import traceback
 from abc import ABC, abstractmethod
 from multiprocessing import Process, Event
-from typing import Dict
+from typing import (
+    Dict,
+    Optional,
+)
 
 from slips_files.core.output import Output
 from slips_files.common.slips_utils import utils
@@ -104,7 +107,6 @@ def should_stop(self) -> bool:
             # this module is still receiving msgs,
             # don't stop
             return False
-
         return True
 
     def print(self, text, verbose=1, debug=0, log_to_logfiles_only=False):
@@ -157,16 +159,20 @@ def pre_main(self):
         executed once before the main loop
         """
 
-    def get_msg(self, channel_name):
+    def get_msg(self, channel_name: str) -> Optional[dict]:
         message = self.db.get_message(self.channels[channel_name])
         if utils.is_msg_intended_for(message, channel_name):
             self.channel_tracker[channel_name] = True
             return message
-        else:
-            self.channel_tracker[channel_name] = False
-            return False
 
-    def run(self):
+        self.channel_tracker[channel_name] = False
+
+    def print_traceback(self):
+        exception_line = sys.exc_info()[2].tb_lineno
+        self.print(f"Problem in pre_main() line {exception_line}", 0, 1)
+        self.print(traceback.format_exc(), 0, 1)
+
+    def run(self) -> bool:
         """
         This is the loop function, it runs non-stop as long as
         the module is running
@@ -180,24 +186,32 @@ def run(self):
             self.shutdown_gracefully()
             return True
         except Exception:
-            exception_line = sys.exc_info()[2].tb_lineno
-            self.print(f"Problem in pre_main() line {exception_line}", 0, 1)
-            self.print(traceback.format_exc(), 0, 1)
+            self.print_traceback()
             return True
 
-        try:
-            while not self.should_stop():
-                # keep running main() in a loop as long as the module is
-                # online
+        keyboard_int_ctr = 0
+        while True:
+            try:
+                if self.should_stop():
+                    self.shutdown_gracefully()
+                    return True
+
                 # if a module's main() returns 1, it means there's an
                 # error and it needs to stop immediately
                 error: bool = self.main()
                 if error:
                     self.shutdown_gracefully()
 
-        except KeyboardInterrupt:
-            self.shutdown_gracefully()
-        except Exception:
-            self.print(f"Problem in {self.name}", 0, 1)
-            self.print(traceback.format_exc(), 0, 1)
-        return True
+            except KeyboardInterrupt:
+                keyboard_int_ctr += 1
+
+                if keyboard_int_ctr >= 2:
+                    # on the second ctrl+c Slips immediately stop
+                    return True
+
+                # on the first ctrl + C keep looping until the should_stop
+                # returns true
+                continue
+            except Exception:
+                self.print_traceback()
+                return False

slips_files/core/input.py

@@ -339,7 +339,6 @@ def get_earliest_line(self):
             # and there is still no files for us.
             # To cover this case, just refresh the list of files
             self.zeek_files = self.db.get_all_zeek_files()
-            # time.sleep(1)
             return False, False
 
         # to fix the problem of evidence being generated BEFORE their corresponding flows are added to our db
@@ -901,7 +900,6 @@ def handle_cyst(self):
                 self.give_profiler(line_info)
                 self.lines += 1
                 self.print("Done reading 1 CYST flow.\n ", 0, 3)
-                time.sleep(2)
 
         self.is_done_processing()