update get_rev_profile() unit tests

slips_files/core/database/database_manager.py

@@ -622,8 +622,11 @@ def add_out_notice(self, *args, **kwargs):
     def add_out_ssl(self, *args, **kwargs):
         return self.rdb.add_out_ssl(*args, **kwargs)
 
-    def getProfileIdFromIP(self, *args, **kwargs):
-        return self.rdb.getProfileIdFromIP(*args, **kwargs)
+    def get_profileid_from_ip(self, *args, **kwargs):
+        return self.rdb.get_profileid_from_ip(*args, **kwargs)
+
+    def get_first_flow_time(self, *args, **kwargs):
+        return self.rdb.get_first_flow_time(*args, **kwargs)
 
     def getProfiles(self, *args, **kwargs):
         return self.rdb.getProfiles(*args, **kwargs)

slips_files/core/database/redis_db/profile_handler.py

@@ -109,8 +109,8 @@ def get_timewindow(self, flowtime, profileid):
 
                 if starttime_of_first_tw:
                     starttime_of_first_tw = float(starttime_of_first_tw)
-                    tw_number: int = floor((flowtime - starttime_of_first_tw) /
-                                      self.width) + 1
+                    tw_number: int = floor((flowtime - starttime_of_first_tw)
+                                           / self.width) + 1
 
                     tw_start: float = starttime_of_first_tw + (
                             self.width * (tw_number-1) )
@@ -1071,10 +1071,13 @@ def add_out_ssl(
                         break
 
 
-    def getProfileIdFromIP(self, daddr_as_obj):
-        """Receive an IP and we want the profileid"""
+    def get_profileid_from_ip(self, ip: str) -> Optional[str]:
+        """
+        returns the profile of the given IP only if it was registered in
+        slips before
+        """
         try:
-            profileid = f'profile{self.separator}{str(daddr_as_obj)}'
+            profileid = f'profile_{ip}'
             if self.r.sismember('profiles', profileid):
                 return profileid
             return False
@@ -1554,6 +1557,9 @@ def mark_profile_as_dhcp(self, profileid):
         if not is_dhcp_set:
             self.r.hset(profileid, 'dhcp', 'true')
 
+    def get_first_flow_time(self) -> Optional[str]:
+        return self.r.hget('analysis', 'file_start')
+
     def addProfile(self, profileid, starttime, duration):
         """
         Add a new profile to the DB. Both the list of profiles and the
@@ -1564,12 +1570,12 @@ def addProfile(self, profileid, starttime, duration):
         Nothing operational
         """
         try:
-            if self.r.sismember('profiles', str(profileid)):
+            if self.r.sismember('profiles', profileid):
                 # we already have this profile
                 return False
 
             # Add the profile to the index. The index is called 'profiles'
-            self.r.sadd('profiles', str(profileid))
+            self.r.sadd('profiles', profileid)
             # Create the hashmap with the profileid. The hasmap of each
             # profile is named with the profileid
 

slips_files/core/profiler.py

@@ -114,15 +114,16 @@ def get_rev_profile(self):
             # some flows don't have a daddr like software.log flows
             return False, False
 
-        rev_profileid = self.db.getProfileIdFromIP(self.daddr_as_obj)
+        rev_profileid: str = self.db.get_profileid_from_ip(self.flow.daddr)
         if not rev_profileid:
             # the profileid is not present in the db, create it
             rev_profileid = f'profile_{self.flow.daddr}'
             self.db.addProfile(rev_profileid, self.flow.starttime, self.width)
 
         # in the database, Find and register the id of the tw where the flow
         # belongs.
-        rev_twid = self.db.get_timewindow(self.flow.starttime, rev_profileid)
+        rev_twid: str = self.db.get_timewindow(
+            self.flow.starttime, rev_profileid)
         return rev_profileid, rev_twid
 
     def add_flow_to_profile(self):

tests/test_database.py

@@ -46,7 +46,7 @@ def test_getProfileIdFromIP():
     # add a profile
     db.addProfile('profile_192.168.1.1', '00:00', '1')
     # try to retrieve it
-    assert db.getProfileIdFromIP(test_ip) is not False
+    assert db.get_profileid_from_ip(test_ip) is not False
 
 
 def test_timewindows():
@@ -73,21 +73,6 @@ def test_add_ips():
     db.addProfile(profileid, '00:00', '1')
     # add a tw to that profile
     db.add_new_tw(profileid, 'timewindow1', 0.0)
-    columns = {
-        'dport': 80,
-        'sport': 80,
-        'totbytes': 80,
-        'pkts': 20,
-        'sbytes': 30,
-        'bytes': 30,
-        'spkts': 70,
-        'state': 'Not Established',
-        'uid': '1234',
-        'proto': 'TCP',
-        'saddr': '8.8.8.8',
-        'daddr': test_ip,
-        'starttime': '20.0',
-    }
     # make sure ip is added
     assert (
         db.add_ips(profileid, twid, flow, 'Server') is True

tests/test_profiler.py

@@ -185,8 +185,6 @@ def test_process_line(file, flow_type):
         )
     assert added_flow is not None
 
-
-
 def test_get_rev_profile(mock_rdb):
     profiler = ModuleFactory().create_profiler_obj()
     profiler.flow = Conn(
@@ -204,8 +202,8 @@ def test_get_rev_profile(mock_rdb):
                 'Established',''
             )
     profiler.daddr_as_obj = ipaddress.ip_address(profiler.flow.daddr)
-    mock_rdb.getProfileIdFromIP.return_value = None
-    mock_rdb.r.hget('analysis', 'file_start').return_value = 0
+    mock_rdb.get_profileid_from_ip.return_value = None
+    mock_rdb.get_timewindow.return_value = 'timewindow1'
     assert profiler.get_rev_profile() == ('profile_8.8.8.8', 'timewindow1')
 
 def test_get_rev_profile_no_daddr(flow):