Browse Source

get_most_used_mss and get_most_used_ttl can now handle lists
get_random_ip_address does not generate duplicates anymore if multiple IPs are requested
changed hash for default smb_scan testcase because of that

Stefano Acquaviti 6 years ago
parent
commit
c1c4f09f2f
3 changed files with 24 additions and 9 deletions
  1. 21 6
      code/ID2TLib/Statistics.py
  2. 2 2
      code/Test/GenericTest.py
  3. 1 1
      code/Test/test_SMBScan.py

+ 21 - 6
code/ID2TLib/Statistics.py

@@ -3,6 +3,7 @@ from math import sqrt, ceil, log
 
 import os
 import time
+import random
 import ID2TLib.libpcapreader as pr
 import matplotlib
 
@@ -514,10 +515,13 @@ class Statistics:
         if count == 1:
             return self.process_db_query("random(all(ipAddress))")
         else:
-            ip_address_list = []
+            ip_address_list = self.process_db_query("all(ipAddress)")
+            result_list = []
             for i in range(0, count):
-                ip_address_list.append(self.process_db_query("random(all(ipAddress))"))
-            return ip_address_list
+                random_ip = random.choice(ip_address_list)
+                result_list.append(random_ip)
+                ip_address_list.remove(random_ip)
+            return result_list
 
     def get_ip_address_from_mac(self, macAddress: str):
         """
@@ -538,9 +542,15 @@ class Statistics:
         :return: The TCP MSS value used by the IP address, or if the IP addresses never specified a MSS,
         then None is returned
         """
-        mss_value = self.process_db_query('SELECT mssValue from tcp_mss WHERE ipAddress="' + ipAddress + '" ORDER BY mssCount DESC LIMIT 1')
+        mss_value = self.process_db_query('SELECT mssValue from tcp_mss WHERE ipAddress="' + ipAddress + '" AND mssCount == (SELECT MAX(mssCount) from tcp_mss WHERE ipAddress="' + ipAddress + '")')
         if isinstance(mss_value, int):
             return mss_value
+        elif isinstance(mss_value, list):
+            if len(mss_value) == 0:
+                return None
+            else:
+                mss_value.sort()
+                return mss_value[0]
         else:
             return None
 
@@ -551,13 +561,18 @@ class Statistics:
         then None is returned
         """
         ttl_value = self.process_db_query(
-            'SELECT ttlValue from ip_ttl WHERE ipAddress="' + ipAddress + '" ORDER BY ttlCount DESC LIMIT 1')
+            'SELECT ttlValue from ip_ttl WHERE ipAddress="' + ipAddress + '" AND ttlCount == (SELECT MAX(ttlCount) from ip_ttl WHERE ipAddress="' + ipAddress + '")')
         if isinstance(ttl_value, int):
             return ttl_value
+        elif isinstance(ttl_value, list):
+            if len(ttl_value) == 0:
+                return None
+            else:
+                ttl_value.sort()
+                return ttl_value[0]
         else:
             return None
 
-
     def get_statistics_database(self):
         """
         :return: A reference to the statistics database object

+ 2 - 2
code/Test/GenericTest.py

@@ -8,7 +8,7 @@ import ID2TLib.TestLibrary as Lib
 
 class GenericTest(unittest.TestCase):
 
-    def generic_test(self, attack_args, sha_checksum, seed=5, cleanup=False, pcap=Lib.test_pcap, flag_write_file=False,
+    def generic_test(self, attack_args, sha_checksum, seed=5, cleanup=True, pcap=Lib.test_pcap, flag_write_file=False,
                      flag_recalculate_stats=False, flag_print_statistics=False, attack_sub_dir=True, test_sub_dir=True):
         # TODO: move seed to attacks
         random.seed(seed)
@@ -19,7 +19,7 @@ class GenericTest(unittest.TestCase):
         caller_function = inspect.stack()[1].function
 
         try:
-            self.assertEqual(Lib.get_sha256(controller.pcap_dest_path), sha_checksum)
+            self.assertEqual(sha_checksum, Lib.get_sha256(controller.pcap_dest_path))
         except self.failureException:
             Lib.rename_test_result_files(controller, caller_function, attack_sub_dir, test_sub_dir)
             raise

+ 1 - 1
code/Test/test_SMBScan.py

@@ -3,7 +3,7 @@ import unittest.mock as mock
 
 import Test.GenericTest as GenericTest
 
-sha_default = '264b243c9b67978f3c892327352f4b293c9a79f6023b06b53d0af7628d171c0b'
+sha_default = '213e194da7bc952cc093868c7450901b0fb93c7255d694eb37ea0b9b48bca65d'
 sha_one_victim_linux = '4928d421caaec8f2c4e5c5bb835b5521b705478779cbc8f343b77143a5a66995'
 sha_victim_range_winxp_hosting = '4c6cb5cb4f838e75b41af4feb2fd9a6fe7e1b226a38b3e8759ce3d31e5a2535e'
 sha_multiple_victims_macos = '0be79b9ad7346562f392e07a5156de978e02f4f25ae8d409b81cc6e0d726012c'