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

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

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"
]


class MapGenerationTest(BasicTest):

    def __init__(self, name, island_client, config_parser, analyzers,
                 timeout, log_handler, break_on_timeout=False):
        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)
        self.break_on_timeout = break_on_timeout

    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()