add unittests for SMBLoris

add test_one_hundred_attackers
add test_ips_in_pcap
move clean_up to Test.Lib

code/Test/Lib.py

@@ -1,9 +1,11 @@
+import os
 import hashlib
 
 from definitions import ROOT_DIR
 
 # TODO: generate better test pcap (1000-2000 packets)
 test_pcap = ROOT_DIR + "/../resources/test/test.pcap"
+test_pcap_ips = ["192.168.189.143", "192.168.189.1"]
 
 
 def get_sha256(file):
@@ -15,3 +17,8 @@ def get_sha256(file):
                 break
             sha.update(data)
     return sha.hexdigest()
+
+
+def clean_up(controller):
+    os.remove(controller.pcap_dest_path)
+    os.remove(controller.label_manager.label_file_path)

code/Test/test_SMBLoris.py

@@ -1,37 +1,48 @@
-import os
 import random
 import unittest
 import unittest.mock as mock
 
 from ID2TLib.Controller import *
-from Test.Lib import test_pcap, get_sha256
+from Test.Lib import test_pcap, test_pcap_ips, get_sha256, clean_up
 
 # FIXME: create new hashes
 sha_one_attacker = '887226f047456608c5c8746c91d387ffa35777650f083564e0104e381155c58e'
+sha_one_hundred_attackers = '3cab29e73bc048ea7cbffde51f5637fe00256d896a6788db27fbec3804f19cc9'
+sha_ips_in_pcap = '4a072c57bf97c8543305964c4df4de5f010df7da22fc8f414ccbbf88b962ae86'
 
 
 class UnitTestSMBLoris(unittest.TestCase):
 
-    def clean_up(self, controller):
-        os.remove(controller.pcap_dest_path)
-        os.remove(controller.label_manager.label_file_path)
-
     def test_one_attacker(self):
         random.seed(5)
         controller = Controller(pcap_file_path=test_pcap, do_extra_tests=False)
         controller.load_pcap_statistics(False, False, False)
         controller.process_attacks([['SMBLorisAttack', 'ip.src=192.168.1.240', 'ip.dst=192.168.1.210']])
         self.assertEqual(get_sha256(controller.pcap_dest_path), sha_one_attacker)
-        self.clean_up(controller)
-
-    #def one_hundred_attacker(self):
-        # TODO: implement test
+        clean_up(controller)
 
-    #def target_ip_not_in_pcap(self):
+    def test_ips_in_pcap(self):
         # TODO: implement test
+        ip_src = 'ip.src='+test_pcap_ips[0]
+        ip_dst = 'ip.dst='+test_pcap_ips[1]
+        random.seed(5)
+        controller = Controller(pcap_file_path=test_pcap, do_extra_tests=False)
+        controller.load_pcap_statistics(False, False, False)
+        controller.process_attacks([['SMBLorisAttack', ip_src, ip_dst]])
+        self.assertEqual(get_sha256(controller.pcap_dest_path), sha_ips_in_pcap)
+        clean_up(controller)
 
     #def five_minutes(self):
         # TODO: implement test
 
+    def test_one_hundred_attackers(self):
+        random.seed(5)
+        controller = Controller(pcap_file_path=test_pcap, do_extra_tests=False)
+        controller.load_pcap_statistics(False, False, False)
+        controller.process_attacks([['SMBLorisAttack', 'ip.dst=192.168.1.210', 'attackers.count=100']])
+        self.assertEqual(get_sha256(controller.pcap_dest_path), sha_one_hundred_attackers)
+        clean_up(controller)
+
+
 if __name__ == '__main__':
     unittest.main()