Merge pull request #966 from Sekhar-Kumar-Dash/patch-59

Created unittest for profile_handler.py

Author: AlyaGomaa

.github/workflows/unit-tests.yml

@@ -71,6 +71,7 @@ jobs:
           - test_timeline.py
           - test_database.py
           - test_symbols_handler.py
+          - test_profile_handler.py
 
     steps:
     - uses: actions/checkout@v4

slips_files/core/database/database_manager.py

@@ -748,8 +748,8 @@ def mark_profile_as_dhcp(self, *args, **kwargs):
     def add_profile(self, *args, **kwargs):
         return self.rdb.add_profile(*args, **kwargs)
 
-    def set_profile_module_label(self, *args, **kwargs):
-        return self.rdb.set_profile_module_label(*args, **kwargs)
+    def set_module_label_for_profile(self, *args, **kwargs):
+        return self.rdb.set_module_label_for_profile(*args, **kwargs)
 
     def check_tw_to_close(self, *args, **kwargs):
         return self.rdb.check_tw_to_close(*args, **kwargs)
@@ -773,7 +773,11 @@ def search_tws_for_flow(self, twid, uid, go_back=False):
         """
 
         # TODO test this
-        tws_to_search = self.rdb.get_tws_to_search(go_back)
+        # how many tws so search back in?
+        tws_to_search = float("inf")
+        if go_back:
+            hrs_to_search = float(go_back)
+            tws_to_search = self.rdb.get_equivalent_tws(hrs_to_search)
 
         twid_number: int = int(twid.split("timewindow")[-1])
         while twid_number > -1 and tws_to_search > 0:
@@ -790,8 +794,8 @@ def search_tws_for_flow(self, twid, uid, go_back=False):
         # uid isn't in this twid or any of the previous ones
         return {uid: None}
 
-    def get_profile_modules_labels(self, *args, **kwargs):
-        return self.rdb.get_profile_modules_labels(*args, **kwargs)
+    def get_modules_labels_of_a_profile(self, *args, **kwargs):
+        return self.rdb.get_modules_labels_of_a_profile(*args, **kwargs)
 
     def add_timeline_line(self, *args, **kwargs):
         return self.rdb.add_timeline_line(*args, **kwargs)

slips_files/core/database/redis_db/profile_handler.py

@@ -76,7 +76,7 @@ def get_timewindow(self, flowtime, profileid):
             belong to that tw
             if it is == the end of a tw, it will belong to the next one
             for example,
-            a flow with ts = 2 belongs to tw1
+            a flow with ts = 2 belongs to tw2
             a flow with ts = 4 belongs to tw3
 
                tw1   tw2   tw3   tw4
@@ -1178,6 +1178,7 @@ def get_modified_profiles_since(
         time_of_last_modified_tw: float = modified_tws[-1][-1]
 
         # this list will store modified profiles without tws
+        # this is a list of ips. not profileids
         profiles = []
         profiles.extend(
             modified_tw[0].split("_")[1] for modified_tw in modified_tws
@@ -1472,32 +1473,35 @@ def add_profile(self, profileid, starttime):
             self.print(type(inst), 0, 1)
             self.print(inst, 0, 1)
 
-    def set_profile_module_label(self, profileid, module, label):
+    def set_module_label_for_profile(self, profileid, module, label):
         """
         Set a module label for a profile.
         A module label is a label set by a module, and not
         a groundtruth label
         """
-        data = self.get_profile_modules_labels(profileid)
+        data = self.get_modules_labels_of_a_profile(profileid)
         data[module] = label
         data = json.dumps(data)
         self.r.hset(profileid, "modules_labels", data)
 
     def check_tw_to_close(self, close_all=False):
         """
-        Check if we should close some TW
-        Search in the modifed tw list and compare when they
+        Check if we should close a TW
+        Search in the modified tw list and compare when they
         were modified with the slips internal time
         """
 
         sit = self.get_slips_internal_time()
 
-        # for each modified profile
+        # sit is the ts of the last tw modification detected by slips
+        # so this line means if 1h(width) passed since the last
+        # modification detected, then it's time to close the tw
         modification_time = float(sit) - self.width
         if close_all:
             # close all tws no matter when they were last modified
             modification_time = float("inf")
 
+        # these are the tws that havent been modified in the last 1h
         profiles_tws_to_close = self.r.zrangebyscore(
             self.constants.MODIFIED_TIMEWINDOWS,
             0,
@@ -1675,15 +1679,7 @@ def add_tuple(
             )
             self.print(traceback.format_exc(), 0, 1)
 
-    def get_tws_to_search(self, go_back):
-        tws_to_search = float("inf")
-
-        if go_back:
-            hrs_to_search = float(go_back)
-            tws_to_search = self.get_equivalent_tws(hrs_to_search)
-        return tws_to_search
-
-    def get_profile_modules_labels(self, profileid):
+    def get_modules_labels_of_a_profile(self, profileid):
         """
         Get labels set by modules in the profile.
         """
@@ -1710,7 +1706,7 @@ def get_timeline_last_lines(
         key = str(
             profileid + self.separator + twid + self.separator + "timeline"
         )
-        # The the amount of lines in this list
+        # The amount of lines in this list
         last_index = self.r.zcard(key)
         # Get the data in the list from the index asked (first_index) until the last
         data = self.r.zrange(key, first_index, last_index - 1)

tests/module_factory.py

@@ -22,6 +22,7 @@
 from modules.flowalerts.dns import DNS
 from modules.flowalerts.downloaded_file import DownloadedFile
 from slips_files.core.helpers.symbols_handler import SymbolHandler
+from slips_files.core.database.redis_db.profile_handler import ProfileHandler
 from modules.flowalerts.notice import Notice
 from modules.flowalerts.smtp import SMTP
 from modules.flowalerts.software import Software
@@ -619,11 +620,6 @@ def create_riskiq_obj(self, mock_db):
         riskiq.db = mock_db
         return riskiq
 
-    def create_alert_handler_obj(self):
-        alert_handler = AlertHandler()
-        alert_handler.constants = Constants()
-        return alert_handler
-
     @patch(MODULE_DB_MANAGER, name="mock_db")
     def create_timeline_object(self, mock_db):
         logger = Mock()
@@ -633,3 +629,18 @@ def create_timeline_object(self, mock_db):
         tl = Timeline(logger, output_dir, redis_port, termination_event)
         tl.db = mock_db
         return tl
+
+    def create_alert_handler_obj(self):
+        alert_handler = AlertHandler()
+        alert_handler.constants = Constants()
+        return alert_handler
+
+    def create_profile_handler_obj(self):
+        handler = ProfileHandler()
+        handler.constants = Constants()
+        handler.r = Mock()
+        handler.rcache = Mock()
+        handler.separator = "_"
+        handler.width = 3600
+        handler.print = Mock()
+        return handler

tests/test_database.py

@@ -153,8 +153,8 @@ def test_profile_moddule_labels():
     db = ModuleFactory().create_db_manager_obj(6387, flush_db=True)
     module_label = "malicious"
     module_name = "test"
-    db.set_profile_module_label(profileid, module_name, module_label)
-    labels = db.get_profile_modules_labels(profileid)
+    db.set_module_label_for_profile(profileid, module_name, module_label)
+    labels = db.get_modules_labels_of_a_profile(profileid)
     assert "test" in labels
     assert labels["test"] == "malicious"