From 0230c26f19012ed21298b25bdfd62d6d2f522d2e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 8 Feb 2021 11:18:18 -0500 Subject: [PATCH] cc: allow server_config.json to be specified at runtime --- monkey/monkey_island.py | 13 +++++++--- .../cc/environment/environment_config.py | 20 +++++++++------ .../cc/environment/environment_singleton.py | 25 +++++++++++-------- .../cc/environment/test_environment_config.py | 7 +++--- monkey/monkey_island/cc/main.py | 5 +++- 5 files changed, 44 insertions(+), 26 deletions(-) diff --git a/monkey/monkey_island.py b/monkey/monkey_island.py index cd452066c..2e410cd9f 100644 --- a/monkey/monkey_island.py +++ b/monkey/monkey_island.py @@ -3,19 +3,24 @@ from gevent import monkey as gevent_monkey gevent_monkey.patch_all() from monkey_island.cc.main import main +from monkey_island.cc.environment.environment_config import DEFAULT_SERVER_CONFIG_PATH def parse_cli_args(): import argparse - parser = argparse.ArgumentParser(description="Infection Monkey Island CnC Server. See https://infectionmonkey.com") + parser = argparse.ArgumentParser(description="Infection Monkey Island CnC Server. See https://infectionmonkey.com", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("-s", "--setup-only", action="store_true", help="Pass this flag to cause the Island to setup and exit without actually starting. " "This is useful for preparing Island to boot faster later-on, so for " "compiling/packaging Islands.") + parser.add_argument("-c", "--config", action="store", + help="The path to the server configuration file.", + default=DEFAULT_SERVER_CONFIG_PATH) args = parser.parse_args() - return args.setup_only + return (args.setup_only, args.config) if "__main__" == __name__: - is_setup_only = parse_cli_args() - main(is_setup_only) + (is_setup_only, config) = parse_cli_args() + main(is_setup_only, config) diff --git a/monkey/monkey_island/cc/environment/environment_config.py b/monkey/monkey_island/cc/environment/environment_config.py index 35dbafc8e..a680658a2 100644 --- a/monkey/monkey_island/cc/environment/environment_config.py +++ b/monkey/monkey_island/cc/environment/environment_config.py @@ -12,6 +12,7 @@ from monkey_island.cc.resources.auth.auth_user import User from monkey_island.cc.resources.auth.user_store import UserStore SERVER_CONFIG_FILENAME = "server_config.json" +DEFAULT_SERVER_CONFIG_PATH = os.path.join(MONKEY_ISLAND_ABS_PATH, 'cc', SERVER_CONFIG_FILENAME) class EnvironmentConfig: @@ -20,6 +21,7 @@ class EnvironmentConfig: deployment: str, user_creds: UserCreds, aws=None): + self.server_config_path = None self.server_config = server_config self.deployment = deployment self.user_creds = user_creds @@ -40,22 +42,24 @@ class EnvironmentConfig: aws=aws) def save_to_file(self): - file_path = EnvironmentConfig.get_config_file_path() - with open(file_path, 'w') as f: + with open(self.server_config_path, 'w') as f: f.write(json.dumps(self.to_dict(), indent=2)) @staticmethod - def get_from_file() -> EnvironmentConfig: - file_path = EnvironmentConfig.get_config_file_path() + def get_from_file(file_path=DEFAULT_SERVER_CONFIG_PATH) -> EnvironmentConfig: + file_path = os.path.expanduser(file_path) + if not Path(file_path).is_file(): server_config_generator.create_default_config_file(file_path) with open(file_path, 'r') as f: config_content = f.read() - return EnvironmentConfig.get_from_json(config_content) - @staticmethod - def get_config_file_path() -> str: - return os.path.join(MONKEY_ISLAND_ABS_PATH, 'cc', SERVER_CONFIG_FILENAME) + environment_config = EnvironmentConfig.get_from_json(config_content) + # TODO: Populating this property is not ideal. Revisit this when you + # make the logger config file configurable at runtime. + environment_config.server_config_path = file_path + + return environment_config def to_dict(self) -> Dict: config_dict = {'server_config': self.server_config, diff --git a/monkey/monkey_island/cc/environment/environment_singleton.py b/monkey/monkey_island/cc/environment/environment_singleton.py index 6b98d0b7c..aab06285f 100644 --- a/monkey/monkey_island/cc/environment/environment_singleton.py +++ b/monkey/monkey_island/cc/environment/environment_singleton.py @@ -1,7 +1,9 @@ import logging import monkey_island.cc.resources.auth.user_store as user_store -from monkey_island.cc.environment import EnvironmentConfig, aws, password, standard, testing +from monkey_island.cc.environment import (EnvironmentConfig, aws, password, + standard, testing) +from monkey_island.cc.environment.environment_config import DEFAULT_SERVER_CONFIG_PATH __author__ = 'itay.mizeretz' @@ -35,13 +37,16 @@ def set_to_standard(): env.save_config() user_store.UserStore.set_users(env.get_auth_users()) +def initialize_from_file(file_path): + try: + config = EnvironmentConfig.get_from_file(file_path) -try: - config = EnvironmentConfig.get_from_file() - __env_type = config.server_config - set_env(__env_type, config) - # noinspection PyUnresolvedReferences - logger.info('Monkey\'s env is: {0}'.format(env.__class__.__name__)) -except Exception: - logger.error('Failed initializing environment', exc_info=True) - raise + __env_type = config.server_config + set_env(__env_type, config) + # noinspection PyUnresolvedReferences + logger.info('Monkey\'s env is: {0}'.format(env.__class__.__name__)) + except Exception: + logger.error('Failed initializing environment', exc_info=True) + raise + +initialize_from_file(DEFAULT_SERVER_CONFIG_PATH) diff --git a/monkey/monkey_island/cc/environment/test_environment_config.py b/monkey/monkey_island/cc/environment/test_environment_config.py index ed9b0ef96..2f076a3a0 100644 --- a/monkey/monkey_island/cc/environment/test_environment_config.py +++ b/monkey/monkey_island/cc/environment/test_environment_config.py @@ -7,7 +7,7 @@ from unittest.mock import MagicMock, patch import monkey_island.cc.test_common.environment.server_config_mocks as config_mocks from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH -from monkey_island.cc.environment.environment_config import EnvironmentConfig +from monkey_island.cc.environment.environment_config import EnvironmentConfig, DEFAULT_SERVER_CONFIG from monkey_island.cc.environment.user_creds import UserCreds @@ -46,6 +46,7 @@ class TestEnvironmentConfig(TestCase): env_config = EnvironmentConfig(server_config=config['server_config'], deployment=config['deployment'], user_creds=user_creds) + env_config.server_config_path = get_server_config_file_path_test_version() env_config.save_to_file() file_path = get_server_config_file_path_test_version() @@ -55,12 +56,12 @@ class TestEnvironmentConfig(TestCase): self.assertDictEqual(config, json.loads(content_from_file)) - def test_get_server_config_file_path(self): + def test_default_server_config_file_path(self): if platform.system() == "Windows": server_file_path = MONKEY_ISLAND_ABS_PATH + r"\cc\server_config.json" else: server_file_path = MONKEY_ISLAND_ABS_PATH + "/cc/server_config.json" - self.assertEqual(EnvironmentConfig.get_config_file_path(), server_file_path) + self.assertEqual(DEFAULT_SERVER_CONFIG_PATH, server_file_path) def test_get_from_dict(self): config_dict = config_mocks.CONFIG_WITH_CREDENTIALS diff --git a/monkey/monkey_island/cc/main.py b/monkey/monkey_island/cc/main.py index ce142edcc..27d928134 100644 --- a/monkey/monkey_island/cc/main.py +++ b/monkey/monkey_island/cc/main.py @@ -21,6 +21,7 @@ json_setup_logging(default_path=Path(MONKEY_ISLAND_ABS_PATH, 'cc', 'island_logge logger = logging.getLogger(__name__) import monkey_island.cc.environment.environment_singleton as env_singleton # noqa: E402 +from monkey_island.cc.environment.environment_config import DEFAULT_SERVER_CONFIG_PATH from common.version import get_version # noqa: E402 from monkey_island.cc.app import init_app # noqa: E402 from monkey_island.cc.server_utils.bootloader_server import BootloaderHttpServer # noqa: E402 @@ -34,8 +35,10 @@ from monkey_island.cc.setup import setup # noqa: E402 MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0" -def main(should_setup_only=False): +def main(should_setup_only=False, server_config_filename=DEFAULT_SERVER_CONFIG_PATH): logger.info("Starting bootloader server") + env_singleton.initialize_from_file(server_config_filename) + mongo_url = os.environ.get('MONGO_URL', env_singleton.env.get_mongo_url()) bootloader_server_thread = Thread(target=BootloaderHttpServer(mongo_url).serve_forever, daemon=True)