From ca454420b9bd8163dc5802ad1c13b1c8cb62e3d3 Mon Sep 17 00:00:00 2001 From: alya Date: Sat, 27 Jan 2024 22:28:53 +0200 Subject: [PATCH] update get_rev_profile() unit tests --- slips_files/core/database/database_manager.py | 7 +++++-- .../core/database/redis_db/profile_handler.py | 20 ++++++++++++------- slips_files/core/profiler.py | 5 +++-- tests/test_database.py | 17 +--------------- tests/test_profiler.py | 6 ++---- 5 files changed, 24 insertions(+), 31 deletions(-) diff --git a/slips_files/core/database/database_manager.py b/slips_files/core/database/database_manager.py index 954b76c9e..dbd895b2f 100644 --- a/slips_files/core/database/database_manager.py +++ b/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) diff --git a/slips_files/core/database/redis_db/profile_handler.py b/slips_files/core/database/redis_db/profile_handler.py index acec6bbf5..de6869d7a 100644 --- a/slips_files/core/database/redis_db/profile_handler.py +++ b/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 diff --git a/slips_files/core/profiler.py b/slips_files/core/profiler.py index 00870283b..b3d33c168 100644 --- a/slips_files/core/profiler.py +++ b/slips_files/core/profiler.py @@ -114,7 +114,7 @@ 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}' @@ -122,7 +122,8 @@ def get_rev_profile(self): # 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): diff --git a/tests/test_database.py b/tests/test_database.py index 44fc4cd6c..57b501986 100644 --- a/tests/test_database.py +++ b/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 diff --git a/tests/test_profiler.py b/tests/test_profiler.py index b41b15fc2..12e9de49f 100644 --- a/tests/test_profiler.py +++ b/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):