Refactored and further improved performance test code structure

envs/monkey_zoo/blackbox/test_blackbox.py

@@ -4,12 +4,12 @@ import logging
 import pytest
 from time import sleep
 
-from envs.monkey_zoo.blackbox.analyzers.performance_analyzer import PerformanceAnalyzer
 from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
 from envs.monkey_zoo.blackbox.analyzers.communication_analyzer import CommunicationAnalyzer
 from envs.monkey_zoo.blackbox.island_client.island_config_parser import IslandConfigParser
+from envs.monkey_zoo.blackbox.tests.performance.report_generation import ReportGenerationTest
 from envs.monkey_zoo.blackbox.utils import gcp_machine_handlers
-from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
 from envs.monkey_zoo.blackbox.log_handlers.test_logs_handler import TestLogsHandler
 
 DEFAULT_TIMEOUT_SECONDS = 5*60
@@ -24,7 +24,7 @@ LOGGER = logging.getLogger(__name__)
 @pytest.fixture(autouse=True, scope='session')
 def GCPHandler(request):
     GCPHandler = gcp_machine_handlers.GCPHandler()
-    GCPHandler.start_machines(" ".join(GCP_TEST_MACHINE_LIST))
+    #GCPHandler.start_machines(" ".join(GCP_TEST_MACHINE_LIST))
     wait_machine_bootup()
 
     def fin():
@@ -55,34 +55,32 @@ def island_client(island):
 class TestMonkeyBlackbox(object):
 
     @staticmethod
-    def run_basic_test(island_client, conf_filename, test_name, timeout_in_seconds=DEFAULT_TIMEOUT_SECONDS):
+    def run_exploitation_test(island_client, conf_filename, test_name, timeout_in_seconds=DEFAULT_TIMEOUT_SECONDS):
         config_parser = IslandConfigParser(conf_filename)
         analyzer = CommunicationAnalyzer(island_client, config_parser.get_ips_of_targets())
         log_handler = TestLogsHandler(test_name, island_client, TestMonkeyBlackbox.get_log_dir_path())
-        BasicTest(
+        ExploitationTest(
             name=test_name,
             island_client=island_client,
             config_parser=config_parser,
             analyzers=[analyzer],
             timeout=timeout_in_seconds,
-            post_exec_analyzers=[],
             log_handler=log_handler).run()
 
     @staticmethod
-    def run_performance_test(island_client, conf_filename, test_name, timeout_in_seconds):
+    def run_performance_test(performance_test_class, island_client,
+                             conf_filename, timeout_in_seconds, break_on_timeout=False):
         config_parser = IslandConfigParser(conf_filename)
-        log_handler = TestLogsHandler(test_name, island_client, TestMonkeyBlackbox.get_log_dir_path())
+        log_handler = TestLogsHandler(performance_test_class.TEST_NAME,
-        BasicTest(
-            name=test_name,
-            island_client=island_client,
-            config_parser=config_parser,
-            analyzers=[CommunicationAnalyzer(island_client, config_parser.get_ips_of_targets())],
-            timeout=timeout_in_seconds,
-            post_exec_analyzers=[PerformanceAnalyzer(
                                       island_client,
-                break_if_took_too_long=False
+                                      TestMonkeyBlackbox.get_log_dir_path())
-            )],
+        analyzer = CommunicationAnalyzer(island_client, config_parser.get_ips_of_targets())
-            log_handler=log_handler).run()
+        performance_test_class.__init__(island_client=island_client,
+                                        config_parser=config_parser,
+                                        analyzers=analyzer,
+                                        timeout=timeout_in_seconds,
+                                        log_handler=log_handler,
+                                        break_on_timeout=break_on_timeout).run()
 
     @staticmethod
     def get_log_dir_path():
@@ -91,44 +89,44 @@ class TestMonkeyBlackbox(object):
     def test_server_online(self, island_client):
         assert island_client.get_api_status() is not None
 
-    def test_ssh_exploiter(self, island_client):
+    #def test_ssh_exploiter(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "SSH.conf", "SSH_exploiter_and_keys")
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "SSH.conf", "SSH_exploiter_and_keys")
+#
+    #def test_hadoop_exploiter(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "HADOOP.conf", "Hadoop_exploiter", 6 * 60)
+#
+    #def test_mssql_exploiter(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "MSSQL.conf", "MSSQL_exploiter")
+#
+    #def test_smb_and_mimikatz_exploiters(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "SMB_MIMIKATZ.conf", "SMB_exploiter_mimikatz")
+#
+    #def test_smb_pth(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "SMB_PTH.conf", "SMB_PTH")
+#
+    #def test_elastic_exploiter(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "ELASTIC.conf", "Elastic_exploiter")
+#
+    #def test_struts_exploiter(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "STRUTS2.conf", "Strtuts2_exploiter")
+#
+    #def test_weblogic_exploiter(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "WEBLOGIC.conf", "Weblogic_exploiter")
+#
+    #def test_shellshock_exploiter(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "SHELLSHOCK.conf", "Shellschock_exploiter")
+#
+    #@pytest.mark.xfail(reason="Test fails randomly - still investigating.")
+    #def test_tunneling(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "TUNNELING.conf", "Tunneling_exploiter", 10 * 60)
+#
+    #def test_wmi_and_mimikatz_exploiters(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "WMI_MIMIKATZ.conf", "WMI_exploiter,_mimikatz")
+#
+    #def test_wmi_pth(self, island_client):
+    #    TestMonkeyBlackbox.run_exploitation_test(island_client, "WMI_PTH.conf", "WMI_PTH")
 
-    def test_hadoop_exploiter(self, island_client):
+    def test_report_generation_performance(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "HADOOP.conf", "Hadoop_exploiter", 6*60)
-
-    def test_mssql_exploiter(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "MSSQL.conf", "MSSQL_exploiter")
-
-    def test_smb_and_mimikatz_exploiters(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "SMB_MIMIKATZ.conf", "SMB_exploiter_mimikatz")
-
-    def test_smb_pth(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "SMB_PTH.conf", "SMB_PTH")
-
-    def test_elastic_exploiter(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "ELASTIC.conf", "Elastic_exploiter")
-
-    def test_struts_exploiter(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "STRUTS2.conf", "Strtuts2_exploiter")
-
-    def test_weblogic_exploiter(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "WEBLOGIC.conf", "Weblogic_exploiter")
-
-    def test_shellshock_exploiter(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "SHELLSHOCK.conf", "Shellschock_exploiter")
-
-    @pytest.mark.xfail(reason="Test fails randomly - still investigating.")
-    def test_tunneling(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "TUNNELING.conf", "Tunneling_exploiter", 10*60)
-
-    def test_wmi_and_mimikatz_exploiters(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "WMI_MIMIKATZ.conf", "WMI_exploiter,_mimikatz")
-
-    def test_wmi_pth(self, island_client):
-        TestMonkeyBlackbox.run_basic_test(island_client, "WMI_PTH.conf", "WMI_PTH")
-
-    def test_performance(self, island_client):
         """
         This test includes the SSH + Elastic + Hadoop + MSSQL machines all in one test
         for a total of 8 machines including the Monkey Island.
@@ -136,8 +134,9 @@ class TestMonkeyBlackbox(object):
         Is has 2 analyzers - the regular one which checks all the Monkeys
         and the Timing one which checks how long the report took to execute
         """
-        TestMonkeyBlackbox.run_performance_test(
+        TestMonkeyBlackbox.run_performance_test(ReportGenerationTest,
                                                 island_client,
                                                 "PERFORMANCE.conf",
-            "test_report_performance",
                                                 timeout_in_seconds=10*60)
+
+

envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py

@@ -0,0 +1,42 @@
+import logging
+from datetime import timedelta
+
+from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.analyzers.performance_analyzer import PerformanceAnalyzer
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class EndpointPerformanceTest(BasicTest):
+
+    def __init__(self, name, test_config: PerformanceTestConfig, island_client: MonkeyIslandClient):
+        self.name = name
+        self.test_config = test_config
+        self.island_client = island_client
+
+    def run(self) -> bool:
+        if not self.island_client.is_all_monkeys_dead():
+            raise RuntimeError("Can't test report times since not all Monkeys have died.")
+
+        # Collect timings for all pages
+        self.island_client.clear_caches()
+        endpoint_timings = {}
+        for endpoint in self.test_config.endpoints_to_test:
+            endpoint_timings[endpoint] = self.get_elapsed_for_get_request(endpoint)
+
+        analyzer = PerformanceAnalyzer(self.test_config, endpoint_timings)
+
+        return analyzer.analyze_test_results()
+
+    def get_elapsed_for_get_request(self, url):
+        response = self.island_client.requests.get(url)
+        if response.ok:
+            LOGGER.debug(f"Got ok for {url} content peek:\n{response.content[:120].strip()}")
+            return response.elapsed
+        else:
+            LOGGER.error(f"Trying to get {url} but got unexpected {str(response)}")
+            # instead of raising for status, mark failed responses as maxtime
+            return timedelta.max()

envs/monkey_zoo/blackbox/tests/performance/map_generation.py

@@ -1,47 +1,35 @@
 from datetime import timedelta
 
-from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
 from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
 from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_workflow import PerformanceTestWorkflow
 
 MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
 MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
 
-REPORT_RESOURCES = [
+MAP_RESOURCES = [
-    "api/report/security",
+    "api/map",
-    "api/attack/report",
-    "api/report/zero_trust/findings",
-    "api/report/zero_trust/principles",
-    "api/report/zero_trust/pillars"
 ]
 
 
-class MapGenerationTest(BasicTest):
+class MapGenerationTest(PerformanceTest):
 
-    def __init__(self, name, island_client, config_parser, analyzers,
+    TEST_NAME = "Map generation performance test"
-                 timeout, log_handler, break_on_timeout=False):
+
-        self.name = name
+    def __init__(self, island_client, config_parser, analyzers,
+                 timeout, log_handler, break_on_timeout):
         self.island_client = island_client
         self.config_parser = config_parser
-        self.exploitation_test = ExploitationTest(name, island_client, config_parser, analyzers, timeout, log_handler)
+        exploitation_test = ExploitationTest(MapGenerationTest.TEST_NAME, island_client,
-        self.break_on_timeout = break_on_timeout
+                                             config_parser, analyzers, timeout, log_handler)
-
-    def run(self):
-        self.island_client.import_config(self.config_parser.config_raw)
-        self.exploitation_test.print_test_starting_info()
-        try:
-            self.island_client.run_monkey_local()
-            self.exploitation_test.test_until_timeout()
-        finally:
-            self.island_client.kill_all_monkeys()
-            self.exploitation_test.wait_until_monkeys_die()
-            self.exploitation_test.wait_for_monkey_process_to_finish()
         performance_config = PerformanceTestConfig(max_allowed_single_page_time=MAX_ALLOWED_SINGLE_PAGE_TIME,
                                                    max_allowed_total_time=MAX_ALLOWED_TOTAL_TIME,
-                                                       endpoints_to_test=REPORT_RESOURCES,
+                                                   endpoints_to_test=MAP_RESOURCES,
-                                                       break_on_timeout=self.break_on_timeout)
+                                                   break_on_timeout=break_on_timeout)
-            performance_test = PerformanceTest("Report generation test", performance_config, self.exploitation_test)
+        self.performance_test_workflow = PerformanceTestWorkflow(MapGenerationTest.TEST_NAME,
-            performance_test.run()
+                                                                 exploitation_test,
-            self.exploitation_test.parse_logs()
+                                                                 performance_config)
-            self.island_client.reset_env()
+
+    def run(self):
+        self.performance_test_workflow.run()

envs/monkey_zoo/blackbox/tests/performance/performance_test.py

@@ -1,42 +1,16 @@
-import logging
+from abc import ABCMeta, abstractmethod
-from datetime import timedelta
 
 from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
-from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
-from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
-from envs.monkey_zoo.blackbox.analyzers.performance_analyzer import PerformanceAnalyzer
 
 
-LOGGER = logging.getLogger(__name__)
+class PerformanceTest(BasicTest, metaclass=ABCMeta):
 
+    @abstractmethod
+    def __init__(self, island_client, config_parser, analyzers,
+                 timeout, log_handler, break_on_timeout):
+        pass
 
-class PerformanceTest(BasicTest):
+    @property
-
+    @abstractmethod
-    def __init__(self, name, test_config: PerformanceTestConfig, exploitation_test: ExploitationTest):
+    def TEST_NAME(self):
-        self.name = name
+        pass
-        self.test_config = test_config
-        self.exploitation_test = exploitation_test
-
-    def run(self) -> bool:
-        if not self.exploitation_test.island_client.is_all_monkeys_dead():
-            raise RuntimeError("Can't test report times since not all Monkeys have died.")
-
-        # Collect timings for all pages
-        self.exploitation_test.island_client.clear_caches()
-        endpoint_timings = {}
-        for endpoint in self.test_config.endpoints_to_test:
-            endpoint_timings[endpoint] = self.exploitation_test.island_client.get_elapsed_for_get_request(endpoint)
-
-        analyzer = PerformanceAnalyzer(self.test_config, endpoint_timings)
-
-        return analyzer.analyze_test_results()
-
-    def get_elapsed_for_get_request(self, url):
-        response = self.exploitation_test.island_client.requests.get(url)
-        if response.ok:
-            LOGGER.debug(f"Got ok for {url} content peek:\n{response.content[:120].strip()}")
-            return response.elapsed
-        else:
-            LOGGER.error(f"Trying to get {url} but got unexpected {str(response)}")
-            # instead of raising for status, mark failed responses as maxtime
-            return timedelta.max()

envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py

@@ -0,0 +1,29 @@
+from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.endpoint_performance_test import EndpointPerformanceTest
+
+
+class PerformanceTestWorkflow(BasicTest):
+
+    def __init__(self, name, exploitation_test: ExploitationTest, performance_config: PerformanceTestConfig):
+        self.name = name
+        self.exploitation_test = exploitation_test
+        self.island_client = exploitation_test.island_client
+        self.config_parser = exploitation_test.config_parser
+        self.performance_config = performance_config
+
+    def run(self):
+        self.island_client.import_config(self.config_parser.config_raw)
+        self.exploitation_test.print_test_starting_info()
+        try:
+            self.island_client.run_monkey_local()
+            self.exploitation_test.test_until_timeout()
+        finally:
+            self.island_client.kill_all_monkeys()
+            self.exploitation_test.wait_until_monkeys_die()
+            self.exploitation_test.wait_for_monkey_process_to_finish()
+            performance_test = EndpointPerformanceTest(self.name, self.performance_config, self.island_client)
+            performance_test.run()
+            self.exploitation_test.parse_logs()
+            self.island_client.reset_env()

envs/monkey_zoo/blackbox/tests/performance/report_generation.py

@@ -1,43 +1,38 @@
 from datetime import timedelta
 
-from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
 from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
 from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_workflow import PerformanceTestWorkflow
 from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
 
 MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
 MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
 
-MAP_RESOURCES = [
+REPORT_RESOURCES = [
-    "api/report/security"
+    "api/report/security",
+    "api/attack/report",
+    "api/report/zero_trust/findings",
+    "api/report/zero_trust/principles",
+    "api/report/zero_trust/pillars"
 ]
 
 
-class ReportGenerationTest(BasicTest):
+class ReportGenerationTest(PerformanceTest):
+    TEST_NAME = "Report generation performance test"
 
-    def __init__(self, name, island_client, config_parser, analyzers,
+    def __init__(self, island_client, config_parser, analyzers,
-                 timeout, log_handler, break_on_timeout=False):
+                 timeout, log_handler, break_on_timeout):
-        self.name = name
         self.island_client = island_client
         self.config_parser = config_parser
-        self.exploitation_test = ExploitationTest(name, island_client, config_parser, analyzers, timeout, log_handler)
+        exploitation_test = ExploitationTest(ReportGenerationTest.TEST_NAME, island_client,
-        self.break_on_timeout = break_on_timeout
+                                             config_parser, analyzers, timeout, log_handler)
-
-    def run(self):
-        self.island_client.import_config(self.config_parser.config_raw)
-        self.exploitation_test.print_test_starting_info()
-        try:
-            self.island_client.run_monkey_local()
-            self.exploitation_test.test_until_timeout()
-        finally:
-            self.island_client.kill_all_monkeys()
-            self.exploitation_test.wait_until_monkeys_die()
-            self.exploitation_test.wait_for_monkey_process_to_finish()
         performance_config = PerformanceTestConfig(max_allowed_single_page_time=MAX_ALLOWED_SINGLE_PAGE_TIME,
                                                    max_allowed_total_time=MAX_ALLOWED_TOTAL_TIME,
                                                    endpoints_to_test=REPORT_RESOURCES,
-                                                       break_on_timeout=self.break_on_timeout)
+                                                   break_on_timeout=break_on_timeout)
-            performance_test = PerformanceTest("Report generation test", performance_config, self.exploitation_test)
+        self.performance_test_workflow = PerformanceTestWorkflow(ReportGenerationTest.TEST_NAME,
-            performance_test.run()
+                                                                 exploitation_test,
-            self.exploitation_test.parse_logs()
+                                                                 performance_config)
-            self.island_client.reset_env()
+
+    def run(self):
+        self.performance_test_workflow.run()