Merge pull request #1184 from guardicore/data-dir-on-island-refactor

Data dir on island refactor

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

monkey/monkey_island/cc/environment/server_config_generator.py

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

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

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
 

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

monkey/monkey_island/config_loader.py

@@ -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

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

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