handle lists in SalityBotnet

DNS query in SalityBotnet can now return lists, is also handled
removed unnecessary list handling in SalityBotnet
removed now unnecessary testcase for this handling
handle_most_used_outputs in Utility can now handle empty lists

code/Attack/SalityBotnet.py

@@ -8,7 +8,7 @@ from definitions import ROOT_DIR
 from Attack import BaseAttack
 from Attack.AttackParameters import Parameter as Param
 from Attack.AttackParameters import ParameterTypes
-from ID2TLib.Utility import update_timestamp, get_interval_pps
+from ID2TLib.Utility import update_timestamp, get_interval_pps, handle_most_used_outputs
 
 logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
 # noinspection PyPep8
@@ -46,8 +46,7 @@ class SalityBotnet(BaseAttack.BaseAttack):
         # PARAMETERS: initialize with default utilsvalues
         # (values are overwritten if user specifies them)
         most_used_ip_address = self.statistics.get_most_used_ip_address()
-        if isinstance(most_used_ip_address, list):
-            most_used_ip_address = most_used_ip_address[0]
+
         self.add_param_value(Param.IP_SOURCE, most_used_ip_address)
         self.add_param_value(Param.MAC_SOURCE, self.statistics.get_mac_address(most_used_ip_address))
 
@@ -73,7 +72,8 @@ class SalityBotnet(BaseAttack.BaseAttack):
         ip_source = self.get_param_value(Param.IP_SOURCE)
 
         # Pick a DNS server from the background traffic
-        ip_dns_server = self.statistics.process_db_query("SELECT ipAddress FROM ip_protocols WHERE protocolName='DNS' ORDER BY protocolCount DESC LIMIT 1;")
+        ip_dns_server = self.statistics.process_db_query("SELECT ipAddress FROM ip_protocols WHERE protocolName='DNS' AND protocolCount=(SELECT MAX(protocolCount) FROM ip_protocols WHERE protocolName='DNS');")
+        ip_dns_server = handle_most_used_outputs(ip_dns_server)
         if not ip_dns_server or ip_source == ip_dns_server:
             ip_dns_server = self.statistics.get_random_ip_address()
         mac_dns_server = self.statistics.get_mac_address(ip_dns_server)

code/ID2TLib/Utility.py

@@ -286,9 +286,12 @@ def get_bytes_from_file(filepath):
 def handle_most_used_outputs(most_used_x):
     """
     :param most_used_x: Element or list (e.g. from SQL-query output) which should only be one element
-    :return: most_used_x if it's not a list. The first element of most_used_x after being sorted if it's a list
+    :return: most_used_x if it's not a list. The first element of most_used_x after being sorted if it's a list.
+    None if that list is empty.
     """
     if isinstance(most_used_x, list):
+        if len(most_used_x) == 0:
+            return None
         most_used_x.sort()
         return most_used_x[0]
     else:

code/Test/test_SalityBotnet.py

@@ -5,7 +5,6 @@ import ID2TLib.TestLibrary as Lib
 import Test.GenericTest as GenericTest
 
 sha_botnet_basic = '8ff1e400dcf01d2d2cb97312cecdb71473ea140f6406ea935f74970aecdd7305'
-sha_botnet_most_used_ip_in_list = '8ff1e400dcf01d2d2cb97312cecdb71473ea140f6406ea935f74970aecdd7305'
 
 """
 CURRENT COVERAGE
@@ -21,11 +20,6 @@ class UnitTestSalityBotnet(GenericTest.GenericTest):
     def test_botnet_basic(self):
         self.generic_test([['SalityBotnet']], sha_botnet_basic)
 
-    @mock.patch('ID2TLib.Statistics.Statistics.get_most_used_ip_address')
-    def test_botnet_most_used_ips(self, mock_most_used_ip_address):
-        mock_most_used_ip_address.return_value = Lib.test_pcap_ips
-        self.generic_test([['SalityBotnet']], sha_botnet_most_used_ip_in_list)
-
 
 if __name__ == '__main__':
     unittest.main()