diff --git a/monkey/monkey_island.py b/monkey/monkey_island.py index 76399b6ff..e3d8348d1 100644 --- a/monkey/monkey_island.py +++ b/monkey/monkey_island.py @@ -1,15 +1,12 @@ from gevent import monkey as gevent_monkey from monkey_island.cc.arg_parser import parse_cli_args +from monkey_island.setup.config_setup import setup_config_by_cmd_arg, setup_default_config gevent_monkey.patch_all() import json # noqa: E402 -import os # noqa: E402 -import monkey_island.cc.environment.server_config_generator as server_config_generator # noqa: E402 -from monkey_island import config_loader # noqa: E402 -from monkey_island.cc.environment.data_dir_generator import create_data_dir # noqa: E402 from monkey_island.cc.server_utils.island_logger import setup_logging # noqa: E402 if "__main__" == __name__: @@ -19,13 +16,9 @@ if "__main__" == __name__: # imports, so the log init needs to be first. try: if island_args.server_config: - server_config_path = os.path.expanduser(island_args.server_config) + config, server_config_path = setup_config_by_cmd_arg(island_args.server_config) else: - server_config_path = server_config_generator.create_default_server_config_file() - - config = config_loader.load_server_config_from_file(server_config_path) - - create_data_dir(config["data_dir"], True) + config, server_config_path = setup_default_config() setup_logging(config["data_dir"], config["log_level"]) diff --git a/monkey/monkey_island/cc/environment/server_config_generator.py b/monkey/monkey_island/cc/environment/server_config_generator.py deleted file mode 100644 index 17ee4a50c..000000000 --- a/monkey/monkey_island/cc/environment/server_config_generator.py +++ /dev/null @@ -1,21 +0,0 @@ -import os -from pathlib import Path - -from monkey_island.cc.environment.data_dir_generator import create_data_dir -from monkey_island.cc.server_utils.consts import ( - DEFAULT_DATA_DIR, - DEFAULT_DEVELOP_SERVER_CONFIG_PATH, - DEFAULT_SERVER_CONFIG_PATH, -) - - -def create_default_server_config_file() -> str: - if not os.path.isfile(DEFAULT_SERVER_CONFIG_PATH): - create_data_dir(DEFAULT_DATA_DIR, False) - write_default_server_config_to_file(DEFAULT_SERVER_CONFIG_PATH) - return DEFAULT_SERVER_CONFIG_PATH - - -def write_default_server_config_to_file(path: str) -> None: - default_config = Path(DEFAULT_DEVELOP_SERVER_CONFIG_PATH).read_text() - Path(path).write_text(default_config) diff --git a/monkey/monkey_island/cc/environment/server_config_handler.py b/monkey/monkey_island/cc/environment/server_config_handler.py new file mode 100644 index 000000000..35cc91f3e --- /dev/null +++ b/monkey/monkey_island/cc/environment/server_config_handler.py @@ -0,0 +1,39 @@ +import json +import os +from pathlib import Path + +from monkey_island.cc.server_utils.consts import ( + DEFAULT_DATA_DIR, + DEFAULT_DEVELOP_SERVER_CONFIG_PATH, + DEFAULT_LOG_LEVEL, + DEFAULT_SERVER_CONFIG_PATH, +) + + +def create_default_server_config_file() -> None: + if not os.path.isfile(DEFAULT_SERVER_CONFIG_PATH): + write_default_server_config_to_file(DEFAULT_SERVER_CONFIG_PATH) + + +def write_default_server_config_to_file(path: str) -> None: + default_config = Path(DEFAULT_DEVELOP_SERVER_CONFIG_PATH).read_text() + Path(path).write_text(default_config) + + +def load_server_config_from_file(server_config_path): + with open(server_config_path, "r") as f: + config_content = f.read() + config = json.loads(config_content) + add_default_values_to_config(config) + + return config + + +def add_default_values_to_config(config): + config["data_dir"] = os.path.abspath( + os.path.expanduser(os.path.expandvars(config.get("data_dir", DEFAULT_DATA_DIR))) + ) + + config.setdefault("log_level", DEFAULT_LOG_LEVEL) + + return config diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index a46e5bd76..9b244a974 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -2,11 +2,12 @@ from mongoengine import connect # Needed so that a server_config.json file exists at the default path, # otherwise, unit tests will error while importing `env_singleton` below. -from monkey_island.cc.environment.server_config_generator import ( # noqa: E402 - create_default_server_config_file, -) +from monkey_island.cc.environment import data_dir_generator, server_config_handler # noqa: E402 -create_default_server_config_file() +from ..server_utils.consts import DEFAULT_DATA_DIR + +data_dir_generator.create_data_dir(DEFAULT_DATA_DIR, False) +server_config_handler.create_default_server_config_file() import monkey_island.cc.environment.environment_singleton as env_singleton # noqa: E402 diff --git a/monkey/monkey_island/cc/server_utils/consts.py b/monkey/monkey_island/cc/server_utils/consts.py index 8b88b000b..a62bb98ab 100644 --- a/monkey/monkey_island/cc/server_utils/consts.py +++ b/monkey/monkey_island/cc/server_utils/consts.py @@ -14,6 +14,8 @@ def get_default_data_dir() -> str: SERVER_CONFIG_FILENAME = "server_config.json" +DEFAULT_LOG_LEVEL = "INFO" + MONKEY_ISLAND_ABS_PATH = os.path.join(os.getcwd(), "monkey_island") DEFAULT_DATA_DIR = os.path.expandvars(get_default_data_dir()) diff --git a/monkey/monkey_island/config_loader.py b/monkey/monkey_island/config_loader.py deleted file mode 100644 index aaa9185d7..000000000 --- a/monkey/monkey_island/config_loader.py +++ /dev/null @@ -1,25 +0,0 @@ -import json -import os - -from monkey_island.cc.server_utils.consts import DEFAULT_DATA_DIR - -DEFAULT_LOG_LEVEL = "INFO" - - -def load_server_config_from_file(server_config_path): - with open(server_config_path, "r") as f: - config_content = f.read() - config = json.loads(config_content) - add_default_values_to_config(config) - - return config - - -def add_default_values_to_config(config): - config["data_dir"] = os.path.abspath( - os.path.expanduser(os.path.expandvars(config.get("data_dir", DEFAULT_DATA_DIR))) - ) - - config.setdefault("log_level", DEFAULT_LOG_LEVEL) - - return config diff --git a/monkey/monkey_island/setup/config_setup.py b/monkey/monkey_island/setup/config_setup.py new file mode 100644 index 000000000..05f410fa8 --- /dev/null +++ b/monkey/monkey_island/setup/config_setup.py @@ -0,0 +1,21 @@ +import os +from typing import Tuple + +from monkey_island.cc.environment import server_config_handler +from monkey_island.cc.environment.data_dir_generator import create_data_dir # noqa: E402 +from monkey_island.cc.server_utils.consts import DEFAULT_DATA_DIR, DEFAULT_SERVER_CONFIG_PATH + + +def setup_config_by_cmd_arg(server_config_path) -> Tuple[dict, str]: + server_config_path = os.path.expandvars(os.path.expanduser(server_config_path)) + config = server_config_handler.load_server_config_from_file(server_config_path) + create_data_dir(config["data_dir"], create_parent_dirs=True) + return config, server_config_path + + +def setup_default_config() -> Tuple[dict, str]: + server_config_path = DEFAULT_SERVER_CONFIG_PATH + create_data_dir(DEFAULT_DATA_DIR, create_parent_dirs=False) + server_config_handler.create_default_server_config_file() + config = server_config_handler.load_server_config_from_file(server_config_path) + return config, server_config_path diff --git a/monkey/tests/unit_tests/monkey_island/cc/environment/__init__.py b/monkey/tests/unit_tests/monkey_island/cc/environment/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/tests/unit_tests/monkey_island/test_config_loader.py b/monkey/tests/unit_tests/monkey_island/cc/environment/test_server_config_handler.py similarity index 64% rename from monkey/tests/unit_tests/monkey_island/test_config_loader.py rename to monkey/tests/unit_tests/monkey_island/cc/environment/test_server_config_handler.py index 20c330f6a..acd89d84f 100644 --- a/monkey/tests/unit_tests/monkey_island/test_config_loader.py +++ b/monkey/tests/unit_tests/monkey_island/cc/environment/test_server_config_handler.py @@ -1,11 +1,11 @@ import os -from monkey_island import config_loader +from monkey_island.cc.environment import server_config_handler from monkey_island.cc.server_utils.consts import DEFAULT_DATA_DIR def test_load_server_config_from_file(test_server_config, mock_home_env): - config = config_loader.load_server_config_from_file(test_server_config) + config = server_config_handler.load_server_config_from_file(test_server_config) assert config["data_dir"] == os.path.join(mock_home_env, ".monkey_island") assert config["log_level"] == "NOTICE" @@ -13,7 +13,7 @@ def test_load_server_config_from_file(test_server_config, mock_home_env): def test_default_log_level(): test_config = {} - config = config_loader.add_default_values_to_config(test_config) + config = server_config_handler.add_default_values_to_config(test_config) assert "log_level" in config assert config["log_level"] == "INFO" @@ -21,7 +21,7 @@ def test_default_log_level(): def test_default_data_dir(mock_home_env): test_config = {} - config = config_loader.add_default_values_to_config(test_config) + config = server_config_handler.add_default_values_to_config(test_config) assert "data_dir" in config assert config["data_dir"] == DEFAULT_DATA_DIR