From 367017a6b6d7c5cd085df8a5d87d5adbe1e26d97 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 10 Apr 2020 14:32:39 +0300 Subject: [PATCH] Refactored and further improved performance test code structure --- envs/monkey_zoo/blackbox/test_blackbox.py | 121 +++++++++--------- .../performance/endpoint_performance_test.py | 42 ++++++ .../tests/performance/map_generation.py | 48 +++---- .../tests/performance/performance_test.py | 46 ++----- .../performance/performance_test_workflow.py | 29 +++++ .../tests/performance/report_generation.py | 47 +++---- 6 files changed, 180 insertions(+), 153 deletions(-) create mode 100644 envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py create mode 100644 envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py diff --git a/envs/monkey_zoo/blackbox/test_blackbox.py b/envs/monkey_zoo/blackbox/test_blackbox.py index e2da6a992..2c7ade076 100644 --- a/envs/monkey_zoo/blackbox/test_blackbox.py +++ b/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()) - 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 - )], - log_handler=log_handler).run() + log_handler = TestLogsHandler(performance_test_class.TEST_NAME, + island_client, + TestMonkeyBlackbox.get_log_dir_path()) + analyzer = CommunicationAnalyzer(island_client, config_parser.get_ips_of_targets()) + 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): - TestMonkeyBlackbox.run_basic_test(island_client, "SSH.conf", "SSH_exploiter_and_keys") + #def test_ssh_exploiter(self, island_client): + # 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): - 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): + def test_report_generation_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( - island_client, - "PERFORMANCE.conf", - "test_report_performance", - timeout_in_seconds=10*60) + TestMonkeyBlackbox.run_performance_test(ReportGenerationTest, + island_client, + "PERFORMANCE.conf", + timeout_in_seconds=10*60) + + diff --git a/envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py b/envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py new file mode 100644 index 000000000..dc2b6b0c7 --- /dev/null +++ b/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() diff --git a/envs/monkey_zoo/blackbox/tests/performance/map_generation.py b/envs/monkey_zoo/blackbox/tests/performance/map_generation.py index b8f1b416e..782bf82fe 100644 --- a/envs/monkey_zoo/blackbox/tests/performance/map_generation.py +++ b/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 = [ - "api/report/security", - "api/attack/report", - "api/report/zero_trust/findings", - "api/report/zero_trust/principles", - "api/report/zero_trust/pillars" +MAP_RESOURCES = [ + "api/map", ] -class MapGenerationTest(BasicTest): +class MapGenerationTest(PerformanceTest): - def __init__(self, name, island_client, config_parser, analyzers, - timeout, log_handler, break_on_timeout=False): - self.name = name + TEST_NAME = "Map generation performance test" + + 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) - self.break_on_timeout = break_on_timeout + exploitation_test = ExploitationTest(MapGenerationTest.TEST_NAME, island_client, + config_parser, analyzers, timeout, log_handler) + performance_config = PerformanceTestConfig(max_allowed_single_page_time=MAX_ALLOWED_SINGLE_PAGE_TIME, + max_allowed_total_time=MAX_ALLOWED_TOTAL_TIME, + endpoints_to_test=MAP_RESOURCES, + break_on_timeout=break_on_timeout) + self.performance_test_workflow = PerformanceTestWorkflow(MapGenerationTest.TEST_NAME, + exploitation_test, + 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_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) - performance_test = PerformanceTest("Report generation test", performance_config, self.exploitation_test) - performance_test.run() - self.exploitation_test.parse_logs() - self.island_client.reset_env() + self.performance_test_workflow.run() diff --git a/envs/monkey_zoo/blackbox/tests/performance/performance_test.py b/envs/monkey_zoo/blackbox/tests/performance/performance_test.py index 294626986..b26c20f93 100644 --- a/envs/monkey_zoo/blackbox/tests/performance/performance_test.py +++ b/envs/monkey_zoo/blackbox/tests/performance/performance_test.py @@ -1,42 +1,16 @@ -import logging -from datetime import timedelta +from abc import ABCMeta, abstractmethod 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): - - def __init__(self, name, test_config: PerformanceTestConfig, exploitation_test: ExploitationTest): - self.name = name - 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() + @property + @abstractmethod + def TEST_NAME(self): + pass diff --git a/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py b/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py new file mode 100644 index 000000000..e496c09a5 --- /dev/null +++ b/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() diff --git a/envs/monkey_zoo/blackbox/tests/performance/report_generation.py b/envs/monkey_zoo/blackbox/tests/performance/report_generation.py index 8c98a8289..52fe76288 100644 --- a/envs/monkey_zoo/blackbox/tests/performance/report_generation.py +++ b/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 = [ - "api/report/security" +REPORT_RESOURCES = [ + "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, - 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) - self.break_on_timeout = break_on_timeout + exploitation_test = ExploitationTest(ReportGenerationTest.TEST_NAME, island_client, + config_parser, analyzers, timeout, log_handler) + 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=break_on_timeout) + self.performance_test_workflow = PerformanceTestWorkflow(ReportGenerationTest.TEST_NAME, + exploitation_test, + 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_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) - performance_test = PerformanceTest("Report generation test", performance_config, self.exploitation_test) - performance_test.run() - self.exploitation_test.parse_logs() - self.island_client.reset_env() + self.performance_test_workflow.run()