瀏覽代碼

added memory_test to ID2TAttackTest.py as well as a default test for every attack

Jonathan Speth 7 年之前
父節點
當前提交
7bda6e6fa0
共有 2 個文件被更改,包括 75 次插入1 次删除
  1. 34 1
      code/Test/ID2TAttackTest.py
  2. 41 0
      code/Test/memory_testing.py

+ 34 - 1
code/Test/ID2TAttackTest.py

@@ -2,6 +2,7 @@ import inspect
 import unittest
 
 import scapy.utils as pcr
+import memory_profiler as memprof
 
 import Core.Controller as Ctrl
 import ID2TLib.TestLibrary as Lib
@@ -54,7 +55,7 @@ class ID2TAttackTest(unittest.TestCase):
                                  flag_write_file=False, flag_recalculate_stats=False, flag_print_statistics=False,
                                  attack_sub_dir=True, test_sub_dir=True):
         """
-        Runs the attack with given aruments and monitors time efficiency.
+        Runs the attack with given arguments and monitors time efficiency.
 
         :param attack_args: A list of attacks with their attack parameters (as defined in Controller.process_attacks).
         :param time_limit: The given time limit in seconds.
@@ -152,3 +153,35 @@ class ID2TAttackTest(unittest.TestCase):
             Lib.clean_up(controller)
         else:
             Lib.rename_test_result_files(controller, caller_function, attack_sub_dir, test_sub_dir)
+
+    def memory_test(self, attack_args, seed=None, 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):
+        """
+        Runs the attack with given arguments and monitors memory usage.
+
+        :param attack_args: A list of attacks with their attack parameters (as defined in Controller.process_attacks).
+        :param seed: A random seed to keep random values static (care for count and order of random generation).
+        :param cleanup: Clean up attack output after testing.
+        :param pcap: The input pcap for the attack.
+        :param flag_write_file: Writes the statistics to a file.
+        :param flag_recalculate_stats: Forces the recalculation of statistics.
+        :param flag_print_statistics: Prints the statistics on the terminal.
+        :param attack_sub_dir: create sub-directory for each attack-class if True
+        :param test_sub_dir: create sub-directory for each test-function/case if True
+        """
+
+        controller = Ctrl.Controller(pcap_file_path=pcap, do_extra_tests=False)
+        controller.load_pcap_statistics(flag_write_file, flag_recalculate_stats, flag_print_statistics)
+        controller.process_attacks(attack_args, [[seed]])
+
+        mem_usage = memprof.memory_usage()
+        print(attack_args[0][0] + ' uses: ' + str(mem_usage[0]) + ' MB of memory')
+
+        caller_function = inspect.stack()[1].function
+
+
+        if cleanup:
+            Lib.clean_up(controller)
+        else:
+            Lib.rename_test_result_files(controller, caller_function, attack_sub_dir, test_sub_dir)

+ 41 - 0
code/Test/memory_testing.py

@@ -0,0 +1,41 @@
+import unittest.mock as mock
+
+import ID2TLib.TestLibrary as Lib
+import Test.ID2TAttackTest as Test
+
+
+class MemoryTests(Test.ID2TAttackTest):
+    def test_SMBLoris(self):
+        self.memory_test([['SMBLorisAttack']])
+
+    def test_SMBScanself(self):
+        self.memory_test([['SMBScanAttack']])
+
+    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=Lib.get_bytes)
+    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.get_x86_nop)
+    def test_FTPExploit(self, mock_get_rnd_x86_nop, mock_get_rnd_bytes):
+        self.memory_test([['FTPWinaXeExploit', 'ip.src=192.168.178.1', 'ip.dst=192.168.178.10']])
+
+    def test_PortscanAttack(self):
+        self.memory_test([['PortscanAttack', 'ip.src=192.168.178.1']])
+
+    def test_SQLi(self):
+        self.memory_test([['SQLiAttack', 'ip.dst=192.168.0.1']])
+
+    def test_Joomla(self):
+        self.memory_test([['JoomlaRegPrivExploit', 'ip.src=192.168.178.1']])
+
+    def test_SalityBotnet(self):
+        self.memory_test([['SalityBotnet']])
+
+    @mock.patch('Attack.BaseAttack.BaseAttack.write_attack_pcap', side_effect=Lib.write_attack_pcap)
+    def test_DDoS(self, mock_write_attack_pcap):
+        self.memory_test([['DDoSAttack']])
+
+    def test_EternalBlue(self):
+        self.memory_test([['EternalBlueExploit']])
+
+    def test_MS17Scan(self):
+        self.memory_test([['MS17ScanAttack']])
+
+