add SalityBotnet tests

code/Attack/SalityBotnet.py

@@ -7,14 +7,14 @@ from scapy.layers.inet import Ether
 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, handle_most_used_outputs
+import ID2TLib.Utility as Util
 
 logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
 # noinspection PyPep8
 
 
 class SalityBotnet(BaseAttack.BaseAttack):
-    template_attack_pcap_path = "resources/sality_botnet.pcap"
+    template_attack_pcap_path = Util.RESOURCE_DIR + "/../resources/sality_botnet.pcap"
 
     def __init__(self):
         """
@@ -72,7 +72,7 @@ class SalityBotnet(BaseAttack.BaseAttack):
 
         # Pick a DNS server from the background traffic
         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)
+        ip_dns_server = Util.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)
@@ -122,8 +122,8 @@ class SalityBotnet(BaseAttack.BaseAttack):
             new_pkt = (eth_frame / ip_pkt)
             new_pkt.time = timestamp_next_pkt
 
-            pps = max(get_interval_pps(complement_interval_pps, timestamp_next_pkt), 10)
-            timestamp_next_pkt = update_timestamp(timestamp_next_pkt, pps)
+            pps = max(Util.get_interval_pps(complement_interval_pps, timestamp_next_pkt), 10)
+            timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps)
 
             packets.append(new_pkt)
 

code/Test/test_SalityBotnet.py

@@ -0,0 +1,18 @@
+import unittest.mock as mock
+
+import Test.ID2TAttackTest as Test
+import ID2TLib.TestLibrary as Lib
+
+sha_botnet_basic = '8ff1e400dcf01d2d2cb97312cecdb71473ea140f6406ea935f74970aecdd7305'
+sha_botnet_most_used_ip_in_list = '8ff1e400dcf01d2d2cb97312cecdb71473ea140f6406ea935f74970aecdd7305'
+
+
+class UnitTestSalityBotnet(Test.ID2TAttackTest):
+
+    def test_botnet_basic(self):
+        self.checksum_test([['SalityBotnet']], sha_botnet_basic)
+
+    @mock.patch('ID2TLib.Statistics.Statistics.get_most_used_ip_address')
+    def test_botnet_most_used_ip(self, mock_most_used_ip_address):
+        mock_most_used_ip_address.return_value = Lib.test_pcap_ips[0]
+        self.checksum_test([['SalityBotnet']], sha_botnet_most_used_ip_in_list)