Patched: "/tmp/tmpqik896dk/modules/flowmldetection/flowmldetection.py"

Author: patched.codes[bot]

modules/flowmldetection/flowmldetection.py

@@ -1,4 +1,6 @@
 from slips_files.common.imports import *
+import hmac
+import hashlib
 from sklearn.linear_model import SGDClassifier
 from sklearn.preprocessing import StandardScaler
 import pickle
@@ -332,18 +334,35 @@ def store_model(self):
         with open('./modules/flowmldetection/scaler.bin', 'wb') as g:
             data = pickle.dumps(self.scaler)
             g.write(data)
-
+    
+    class SecureUnpickler(pickle.Unpickler):
+        @classmethod
+        def loads(cls, data, key):
+            # Check HMAC for data integrity
+            hmac_new, data = data[:64], data[64:]
+            if hmac.new(key, data, hashlib.sha256).hexdigest().encode() != hmac_new:
+                raise ValueError("Data integrity check failed")
+            return cls.loads(data)
+    
+        def load(self):
+            raise NotImplementedError("Use SecureUnpickler.loads() instead")
+    
+        def loads(self, data):
+            return super().loads(data)
+    
     def read_model(self):
         """
         Read the trained model from disk
         """
         try:
             self.print('Reading the trained model from disk.', 0, 2)
             with open('./modules/flowmldetection/model.bin', 'rb') as f:
-                self.clf = pickle.load(f)
+                data = f.read()
+                self.clf = SecureUnpickler.loads(data, b'my_secret_key')
             self.print('Reading the trained scaler from disk.', 0, 2)
             with open('./modules/flowmldetection/scaler.bin', 'rb') as g:
-                self.scaler = pickle.load(g)
+                data = g.read()
+                self.scaler = SecureUnpickler.loads(data, b'my_secret_key')
         except FileNotFoundError:
             # If there is no model, create one empty
             self.print('There was no model. Creating a new empty model.', 0, 2)