monkey/envs/monkey_zoo/blackbox/tests/exploitation.py

import logging
from time import sleep

from envs.monkey_zoo.blackbox.island_client.island_config_parser import IslandConfigParser
from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
from envs.monkey_zoo.blackbox.utils.test_timer import TestTimer

MAX_TIME_FOR_MONKEYS_TO_DIE = 2 * 60
WAIT_TIME_BETWEEN_REQUESTS = 1
TIME_FOR_MONKEY_PROCESS_TO_FINISH = 5
DELAY_BETWEEN_ANALYSIS = 1
LOGGER = logging.getLogger(__name__)


class ExploitationTest(BasicTest):
    def __init__(self, name, island_client, raw_config, analyzers, timeout, log_handler):
        self.name = name
        self.island_client = island_client
        self.raw_config = raw_config
        self.analyzers = analyzers
        self.timeout = timeout
        self.log_handler = log_handler

    def run(self):
        self.island_client.import_config(self.raw_config)
        self.print_test_starting_info()
        try:
            self.island_client.run_monkey_local()
            self.test_until_timeout()
        finally:
            self.island_client.kill_all_monkeys()
            self.wait_until_monkeys_die()
            self.wait_for_monkey_process_to_finish()
            self.parse_logs()
            self.island_client.reset_env()

    def print_test_starting_info(self):
        LOGGER.info("Started {} test".format(self.name))
        machine_list = ", ".join(IslandConfigParser.get_ips_of_targets(self.raw_config))
        LOGGER.info(f"Machines participating in test: {machine_list}")
        print("")

    def test_until_timeout(self):
        timer = TestTimer(self.timeout)
        while not timer.is_timed_out():
            if self.all_analyzers_pass():
                self.log_success(timer)
                return
            sleep(DELAY_BETWEEN_ANALYSIS)
            LOGGER.debug(
                "Waiting until all analyzers passed. Time passed: {}".format(timer.get_time_taken())
            )
        self.log_failure(timer)
        assert False

    def log_success(self, timer):
        LOGGER.info(self.get_analyzer_logs())
        LOGGER.info(
            "{} test passed, time taken: {:.1f} seconds.".format(self.name, timer.get_time_taken())
        )

    def log_failure(self, timer):
        LOGGER.info(self.get_analyzer_logs())
        LOGGER.error(
            "{} test failed because of timeout. Time taken: {:.1f} seconds.".format(
                self.name, timer.get_time_taken()
            )
        )

    def all_analyzers_pass(self):
        analyzers_results = [analyzer.analyze_test_results() for analyzer in self.analyzers]
        return all(analyzers_results)

    def get_analyzer_logs(self):
        log = ""
        for analyzer in self.analyzers:
            log += "\n" + analyzer.log.get_contents()
        return log

    def wait_until_monkeys_die(self):
        time_passed = 0
        while (
            not self.island_client.is_all_monkeys_dead()
            and time_passed < MAX_TIME_FOR_MONKEYS_TO_DIE
        ):
            sleep(WAIT_TIME_BETWEEN_REQUESTS)
            time_passed += WAIT_TIME_BETWEEN_REQUESTS
            LOGGER.debug("Waiting for all monkeys to die. Time passed: {}".format(time_passed))
        if time_passed > MAX_TIME_FOR_MONKEYS_TO_DIE:
            LOGGER.error("Some monkeys didn't die after the test, failing")
            assert False
        LOGGER.info(f"After {time_passed} seconds all monkeys have died")

    def parse_logs(self):
        LOGGER.info("Parsing test logs:")
        self.log_handler.parse_test_logs()

    @staticmethod
    def wait_for_monkey_process_to_finish():
        """
        There is a time period when monkey is set to dead, but the process is still closing.
        If we try to launch monkey during that time window monkey will fail to start, that's
        why test needs to wait a bit even after all monkeys are dead.
        """
        LOGGER.debug("Waiting for Monkey process to close...")
        sleep(TIME_FOR_MONKEY_PROCESS_TO_FINISH)