flowmldetection.py: better error handling

stratosphereips · Jan 8, 2025 · ba1c23e · ba1c23e
1 parent 9a95de8
commit ba1c23e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/modules/flowmldetection/flowmldetection.py b/modules/flowmldetection/flowmldetection.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 import numpy
 from sklearn.linear_model import SGDClassifier
 from sklearn.preprocessing import StandardScaler
@@ -294,12 +296,13 @@ def process_flow(self, flow_to_process: dict):
             self.print("Error in process_flow()")
             self.print(traceback.format_exc(), 0, 1)
 
-    def detect(self, x_flow) -> numpy.ndarray:
+    def detect(self, x_flow) -> Optional[numpy.ndarray]:
         """
         Detects the given flow with the current model stored
         and returns the predection array
         """
         try:
+            given_x_flow = x_flow
             # clean the flow
             fields_to_drop = [
                 "label",
@@ -323,7 +326,9 @@ def detect(self, x_flow) -> numpy.ndarray:
             pred: numpy.ndarray = self.clf.predict(x_flow)
             return pred
         except Exception as e:
-            self.print(f"Error in detect(): {e}")
+            self.print(
+                f"Error in detect() while processing " f"\n{given_x_flow}\n{e}"
+            )
             self.print(traceback.format_exc(), 0, 1)
 
     def store_model(self):
@@ -468,6 +473,10 @@ def main(self):
                 if processed_flow is not None and not processed_flow.empty:
                     # Predict
                     pred: numpy.ndarray = self.detect(processed_flow)
+                    if not pred:
+                        # an error occurred
+                        return
+
                     label = self.flow["label"]
                     if label and label != "unknown" and label != pred[0]:
                         # If the user specified a label in test mode,

diff --git a/tests/integration_tests/test_pcap_dataset.py b/tests/integration_tests/test_pcap_dataset.py
@@ -38,7 +38,10 @@ def test_pcap(
 ):
     output_dir = create_output_dir(output_dir)
     output_file = os.path.join(output_dir, "slips_output.txt")
-    command = f"./slips.py  -e 1 -t -f {pcap_path} -o {output_dir}  -P {redis_port} > {output_file} 2>&1"
+    command = (
+        f"./slips.py -e 1 -t -f {pcap_path} -o {output_dir} "
+        f" -P {redis_port} > {output_file} 2>&1"
+    )
     # this function returns when slips is done
     run_slips(command)
     assert_no_errors(output_dir)