add Test.Lib

add test_pcap constant
update TODOs and FIXMEs
add clean_up function

code/Test/Lib.py

@@ -0,0 +1,17 @@
+import hashlib
+
+from definitions import ROOT_DIR
+
+# TODO: generate better test pcap (1000-2000 packets)
+test_pcap = ROOT_DIR + "/../resources/test/test.pcap"
+
+
+def get_sha256(file):
+    sha = hashlib.sha256()
+    with open(file, 'rb') as f:
+        while True:
+            data = f.read(0x100000)
+            if not data:
+                break
+            sha.update(data)
+    return sha.hexdigest()

code/Test/test_SMBLoris.py

@@ -1,33 +1,28 @@
+import os
 import random
 import unittest
 import unittest.mock as mock
-import hashlib
 
-from definitions import ROOT_DIR
 from ID2TLib.Controller import *
+from Test.Lib import test_pcap, get_sha256
 
+# FIXME: create new hashes
 sha_one_attacker = '887226f047456608c5c8746c91d387ffa35777650f083564e0104e381155c58e'
 
 
 class UnitTestSMBLoris(unittest.TestCase):
 
-    def get_sha256(self, file):
-        sha = hashlib.sha256()
-        with open(file, 'rb') as f:
-            while True:
-                data = f.read(0x100000)
-                if not data:
-                    break
-                sha.update(data)
-        return sha.hexdigest()
+    def clean_up(self, controller):
+        os.remove(controller.pcap_dest_path)
+        os.remove(controller.label_manager.label_file_path)
 
     def test_one_attacker(self):
-        # TODO: implement test
         random.seed(5)
-        controller = Controller(pcap_file_path=ROOT_DIR+"/../resources/test/test.pcap", do_extra_tests=False)
+        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(self.get_sha256(controller.pcap_dest_path), sha_one_attacker)
+        self.assertEqual(get_sha256(controller.pcap_dest_path), sha_one_attacker)
+        self.clean_up(controller)
 
     #def one_hundred_attacker(self):
         # TODO: implement test
@@ -38,6 +33,5 @@ class UnitTestSMBLoris(unittest.TestCase):
     #def five_minutes(self):
         # TODO: implement test
 
-
 if __name__ == '__main__':
     unittest.main()

code/definitions.py

@@ -1,3 +1,3 @@
 import os
 
-ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))