From 9cffb9e9d1d977cd8b5c6b60004e45668ca5e826 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 21 May 2021 11:26:02 +0300 Subject: [PATCH] Refactored `SetupParams` into IslandConfigOptions and altered setup workflow to use it --- monkey/monkey_island.py | 14 ++++---- monkey/monkey_island/cc/main.py | 13 ++++--- monkey/monkey_island/config_file_parser.py | 8 ++++- .../setup/island_config_options.py | 30 ++++++++++++++++ .../setup/setup_param_factory.py | 34 ------------------- monkey/monkey_island/setup/setup_params.py | 18 ---------- whitelist.py | 2 +- 7 files changed, 50 insertions(+), 69 deletions(-) create mode 100644 monkey/monkey_island/setup/island_config_options.py delete mode 100644 monkey/monkey_island/setup/setup_param_factory.py delete mode 100644 monkey/monkey_island/setup/setup_params.py diff --git a/monkey/monkey_island.py b/monkey/monkey_island.py index 7f21bb0b4..59ac8301a 100644 --- a/monkey/monkey_island.py +++ b/monkey/monkey_island.py @@ -1,8 +1,8 @@ from gevent import monkey as gevent_monkey from monkey_island.cc.arg_parser import parse_cli_args -from monkey_island.config_file_parser import load_server_config_from_file -from monkey_island.setup.setup_param_factory import SetupParamFactory +from monkey_island.config_file_parser import load_island_config_from_file +from monkey_island.setup.island_config_options import IslandConfigOptions gevent_monkey.patch_all() @@ -12,15 +12,13 @@ from monkey_island.cc.server_utils.island_logger import setup_logging # noqa: E if "__main__" == __name__: island_args = parse_cli_args() - config_contents = load_server_config_from_file(island_args.server_config_path) - - setup_params = SetupParamFactory().build(island_args, config_contents) + config_options = IslandConfigOptions() try: - # This is here in order to catch EVERYTHING, some functions are being called on # imports, so the log init needs to be first. - setup_logging(setup_params.data_dir, setup_params.log_level) + config_options = load_island_config_from_file(island_args.server_config_path) + setup_logging(config_options.data_dir, config_options.log_level) except OSError as ex: print(f"Error opening server config file: {ex}") @@ -32,4 +30,4 @@ if "__main__" == __name__: from monkey_island.cc.main import main # noqa: E402 - main(setup_params) + main(island_args.setup_only, island_args.server_config_path, config_options) diff --git a/monkey/monkey_island/cc/main.py b/monkey/monkey_island/cc/main.py index 549504218..b9dd9b197 100644 --- a/monkey/monkey_island/cc/main.py +++ b/monkey/monkey_island/cc/main.py @@ -9,7 +9,7 @@ from threading import Thread # "monkey_island." work. from gevent.pywsgi import WSGIServer -from monkey_island.setup.setup_params import SetupParams +from monkey_island.setup.island_config_options import IslandConfigOptions MONKEY_ISLAND_DIR_BASE_PATH = str(Path(__file__).parent.parent) if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path: @@ -35,12 +35,11 @@ from monkey_island.cc.setup import setup # noqa: E402 MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0" -def main(setup_params: SetupParams): - logger.info("Starting bootloader server") +def main(setup_only: bool, server_config_path: str, config_options: IslandConfigOptions): - env_singleton.initialize_from_file(setup_params.server_config_path) - initialize_encryptor(setup_params.data_dir) - initialize_services(setup_params.data_dir) + env_singleton.initialize_from_file(server_config_path) + initialize_encryptor(config_options.data_dir) + initialize_services(config_options.data_dir) mongo_url = os.environ.get("MONGO_URL", env_singleton.env.get_mongo_url()) bootloader_server_thread = Thread( @@ -48,7 +47,7 @@ def main(setup_params: SetupParams): ) bootloader_server_thread.start() - start_island_server(setup_params.setup_only) + start_island_server(setup_only) bootloader_server_thread.join() diff --git a/monkey/monkey_island/config_file_parser.py b/monkey/monkey_island/config_file_parser.py index 78a621ab2..c7aa4f146 100644 --- a/monkey/monkey_island/config_file_parser.py +++ b/monkey/monkey_island/config_file_parser.py @@ -2,9 +2,15 @@ import json from os.path import isfile from monkey_island.cc.server_utils.consts import DEFAULT_SERVER_CONFIG_PATH +from monkey_island.setup.island_config_options import IslandConfigOptions -def load_server_config_from_file(server_config_path: str) -> dict: +def load_island_config_from_file(server_config_path: str) -> IslandConfigOptions: + config_contents = read_config_file(server_config_path) + return IslandConfigOptions.build_from_config_file_contents(config_contents) + + +def read_config_file(server_config_path: str) -> dict: if not server_config_path or not isfile(server_config_path): server_config_path = DEFAULT_SERVER_CONFIG_PATH with open(server_config_path, "r") as f: diff --git a/monkey/monkey_island/setup/island_config_options.py b/monkey/monkey_island/setup/island_config_options.py new file mode 100644 index 000000000..ce4a96f7b --- /dev/null +++ b/monkey/monkey_island/setup/island_config_options.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from monkey_island.cc.server_utils.consts import ( + DEFAULT_DATA_DIR, + DEFAULT_LOG_LEVEL, + DEFAULT_START_MONGO_DB, +) + + +@dataclass +class IslandConfigOptions: + log_level = DEFAULT_LOG_LEVEL + data_dir = DEFAULT_DATA_DIR + start_mongodb = DEFAULT_START_MONGO_DB + + @staticmethod + def build_from_config_file_contents(config_contents: dict) -> IslandConfigOptions: + config = IslandConfigOptions() + if "data_dir" in config_contents: + config.data_dir = config_contents["data_dir"] + + if "log_level" in config_contents: + config.log_level = config_contents["log_level"] + + if "mongodb" in config_contents and "start_mongodb" in config_contents["mongodb"]: + config.start_mongodb = config_contents["mongodb"]["start_mongodb"] + + return config diff --git a/monkey/monkey_island/setup/setup_param_factory.py b/monkey/monkey_island/setup/setup_param_factory.py deleted file mode 100644 index 32b1b60db..000000000 --- a/monkey/monkey_island/setup/setup_param_factory.py +++ /dev/null @@ -1,34 +0,0 @@ -import os - -from monkey_island.cc.arg_parser import IslandArgs -from monkey_island.setup.setup_params import SetupParams - - -class SetupParamFactory: - def __init__(self): - self.setup_params = SetupParams() - - def build(self, cmd_args: IslandArgs, config_contents: dict) -> SetupParams: - - self._update_by_cmd_args(cmd_args) - self._update_by_config_file(config_contents) - - return self.setup_params - - def _update_by_cmd_args(self, cmd_args: IslandArgs): - if type(cmd_args.setup_only) == bool: - self.setup_params.setup_only = cmd_args.setup_only - - if cmd_args.server_config_path: - self.setup_params.server_config_path = os.path.expanduser(cmd_args.server_config_path) - - def _update_by_config_file(self, config_contents: dict): - - if "data_dir" in config_contents: - self.setup_params.data_dir = config_contents["data_dir"] - - if "log_level" in config_contents: - self.setup_params.log_level = config_contents["log_level"] - - if "mongodb" in config_contents and "start_mongodb" in config_contents["mongodb"]: - self.setup_params.start_mongodb = config_contents["mongodb"]["start_mongodb"] diff --git a/monkey/monkey_island/setup/setup_params.py b/monkey/monkey_island/setup/setup_params.py deleted file mode 100644 index c302c7fc0..000000000 --- a/monkey/monkey_island/setup/setup_params.py +++ /dev/null @@ -1,18 +0,0 @@ -from dataclasses import dataclass - -from monkey_island.cc.server_utils.consts import ( - DEFAULT_DATA_DIR, - DEFAULT_LOG_LEVEL, - DEFAULT_SERVER_CONFIG_PATH, - DEFAULT_SHOULD_SETUP_ONLY, - DEFAULT_START_MONGO_DB, -) - - -@dataclass -class SetupParams: - server_config_path = DEFAULT_SERVER_CONFIG_PATH - log_level = DEFAULT_LOG_LEVEL - data_dir = DEFAULT_DATA_DIR - start_mongodb = DEFAULT_START_MONGO_DB - setup_only = DEFAULT_SHOULD_SETUP_ONLY diff --git a/whitelist.py b/whitelist.py index bd147220a..51d4c22b8 100644 --- a/whitelist.py +++ b/whitelist.py @@ -165,7 +165,7 @@ ALIBABA # unused variable (monkey/common/cloud/environment_names.py:10) IBM # unused variable (monkey/common/cloud/environment_names.py:11) DigitalOcean # unused variable (monkey/common/cloud/environment_names.py:12) _.aws_info # unused attribute (monkey/monkey_island/cc/environment/aws.py:13) - +build_from_config_file_contents # unused method 'build_from_config_file_contents' (\monkey_island\setup\island_config_options.py:18) # these are not needed for it to work, but may be useful extra information to understand what's going on WINDOWS_PBA_TYPE # unused variable (monkey/monkey_island/cc/resources/pba_file_upload.py:23)