Select server config at runtime

.travis.yml

@@ -18,7 +18,7 @@ os: linux
 
 before_install:
 # Init server_config.json to default
-- cp monkey/monkey_island/cc/server_config.json.default monkey/monkey_island/cc/server_config.json
+- cp monkey/monkey_island/cc/server_config.json.develop monkey/monkey_island/cc/server_config.json
 
 install:
 # Python

monkey/monkey_island.py

@@ -2,20 +2,25 @@
 
 gevent_monkey.patch_all()
 
-from monkey_island.cc.main import main
+from monkey_island.cc.main import main  # noqa: E402
+from monkey_island.cc.environment.environment_config import DEFAULT_SERVER_CONFIG_PATH  # noqa: E402
 
 
 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)

monkey/monkey_island/cc/consts.py

@@ -1,6 +1,14 @@
 import os
 
-__author__ = 'itay.mizeretz'
+__author__ = "itay.mizeretz"
 
-MONKEY_ISLAND_ABS_PATH = os.path.join(os.getcwd(), 'monkey_island')
+MONKEY_ISLAND_ABS_PATH = os.path.join(os.getcwd(), "monkey_island")
 DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS = 60 * 5
+
+DEFAULT_SERVER_CONFIG_PATH = os.path.join(
+    MONKEY_ISLAND_ABS_PATH, "cc", "server_config.json"
+)
+
+DEFAULT_DEVELOP_SERVER_CONFIG_PATH = os.path.join(
+    MONKEY_ISLAND_ABS_PATH, "cc", "server_config.json.develop"
+)

monkey/monkey_island/cc/environment/environment_config.py

@@ -6,56 +6,47 @@
 from typing import Dict, List
 
 import monkey_island.cc.environment.server_config_generator as server_config_generator
-from monkey_island.cc.consts import MONKEY_ISLAND_ABS_PATH
 from monkey_island.cc.environment.user_creds import UserCreds
 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"
-
 
 class EnvironmentConfig:
-    def __init__(self,
-                 server_config: str,
-                 deployment: str,
-                 user_creds: UserCreds,
-                 aws=None):
-        self.server_config = server_config
-        self.deployment = deployment
-        self.user_creds = user_creds
-        self.aws = aws
-
-    @staticmethod
-    def get_from_json(config_json: str) -> EnvironmentConfig:
-        data = json.loads(config_json)
-        return EnvironmentConfig.get_from_dict(data)
+    def __init__(self, file_path):
+        self._server_config_path = os.path.expanduser(file_path)
+        self.server_config = None
+        self.deployment = None
+        self.user_creds = None
+        self.aws = None
 
-    @staticmethod
-    def get_from_dict(dict_data: Dict) -> EnvironmentConfig:
-        user_creds = UserCreds.get_from_dict(dict_data)
-        aws = dict_data['aws'] if 'aws' in dict_data else None
-        return EnvironmentConfig(server_config=dict_data['server_config'],
-                                 deployment=dict_data['deployment'],
-                                 user_creds=user_creds,
-                                 aws=aws)
+        self._load_from_file(self._server_config_path)
 
-    def save_to_file(self):
-        file_path = EnvironmentConfig.get_config_file_path()
-        with open(file_path, 'w') as f:
-            f.write(json.dumps(self.to_dict(), indent=2))
+    def _load_from_file(self, file_path):
+        file_path = os.path.expanduser(file_path)
 
-    @staticmethod
-    def get_from_file() -> EnvironmentConfig:
-        file_path = EnvironmentConfig.get_config_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)
+        self._load_from_json(config_content)
+
+    def _load_from_json(self, config_json: str) -> EnvironmentConfig:
+        data = json.loads(config_json)
+        self._load_from_dict(data)
+
+    def _load_from_dict(self, dict_data: Dict):
+        user_creds = UserCreds.get_from_dict(dict_data)
+        aws = dict_data['aws'] if 'aws' in dict_data else None
+
+        self.server_config = dict_data['server_config']
+        self.deployment = dict_data['deployment']
+        self.user_creds = user_creds
+        self.aws = aws
+
+    def save_to_file(self):
+        with open(self._server_config_path, 'w') as f:
+            f.write(json.dumps(self.to_dict(), indent=2))
 
     def to_dict(self) -> Dict:
         config_dict = {'server_config': self.server_config,

monkey/monkey_island/cc/environment/environment_singleton.py

@@ -3,6 +3,7 @@
 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.consts import DEFAULT_SERVER_CONFIG_PATH
 
 __author__ = 'itay.mizeretz'
 
@@ -39,12 +40,17 @@ def set_to_standard():
         user_store.UserStore.set_users(env.get_auth_users())
 
 
-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
+def initialize_from_file(file_path):
+    try:
+        config = EnvironmentConfig(file_path)
+
+        __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)

monkey/monkey_island/cc/environment/server_config_generator.py

@@ -1,7 +1,8 @@
 from pathlib import Path
 
+from monkey_island.cc.consts import DEFAULT_DEVELOP_SERVER_CONFIG_PATH
+
 
 def create_default_config_file(path):
-    default_config_path = f"{path}.default"
-    default_config = Path(default_config_path).read_text()
+    default_config = Path(DEFAULT_DEVELOP_SERVER_CONFIG_PATH).read_text()
     Path(path).write_text(default_config)

monkey/monkey_island/cc/environment/set_server_config.py

@@ -14,7 +14,7 @@ def add_monkey_dir_to_sys_path():
 
 add_monkey_dir_to_sys_path()
 
-from monkey_island.cc.environment.environment_config import EnvironmentConfig  # noqa: E402 isort:skip
+from monkey_island.cc.consts import DEFAULT_SERVER_CONFIG_PATH  # noqa: E402 isort:skip
 
 SERVER_CONFIG = "server_config"
 BACKUP_CONFIG_FILENAME = "./server_config.backup"
@@ -26,7 +26,7 @@ def add_monkey_dir_to_sys_path():
 
 def main():
     args = parse_args()
-    file_path = EnvironmentConfig.get_config_file_path()
+    file_path = DEFAULT_SERVER_CONFIG_PATH
 
     if args.server_config == "restore":
         restore_previous_config(file_path)

monkey/monkey_island/cc/environment/test__init__.py

@@ -1,17 +1,43 @@
-import json
 import os
+import tempfile
 from typing import Dict
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
-import monkey_island.cc.testing.environment.server_config_mocks as config_mocks
+from monkey_island.cc.consts import MONKEY_ISLAND_ABS_PATH
 from common.utils.exceptions import (AlreadyRegisteredError,
                                      CredentialsNotRequiredError,
                                      InvalidRegistrationCredentialsError,
                                      RegistrationNotNeededError)
 from monkey_island.cc.environment import (Environment, EnvironmentConfig,
                                           UserCreds)
 
+TEST_RESOURCES_DIR = os.path.join(MONKEY_ISLAND_ABS_PATH, "cc", "testing", "environment")
+
+WITH_CREDENTIALS = os.path.join(TEST_RESOURCES_DIR, "server_config_with_credentials.json")
+NO_CREDENTIALS = os.path.join(TEST_RESOURCES_DIR, "server_config_no_credentials.json")
+PARTIAL_CREDENTIALS = os.path.join(TEST_RESOURCES_DIR, "server_config_partial_credentials.json")
+STANDARD_WITH_CREDENTIALS = os.path.join(TEST_RESOURCES_DIR,
+                                         "server_config_standard_with_credentials.json")
+STANDARD_ENV = os.path.join(TEST_RESOURCES_DIR,
+                            "server_config_standard_env.json")
+
+
+def get_tmp_file():
+    with tempfile.NamedTemporaryFile(delete=False) as f:
+        return f.name
+
+
+class StubEnvironmentConfig(EnvironmentConfig):
+    def __init__(self, server_config, deployment, user_creds):
+        self.server_config = server_config
+        self.deployment = deployment
+        self.user_creds = user_creds
+        self.server_config_path = get_tmp_file()
+
+    def __del__(self):
+        os.remove(self.server_config_path)
+
 
 def get_server_config_file_path_test_version():
     return os.path.join(os.getcwd(), 'test_config.json')
@@ -21,7 +47,7 @@ class TestEnvironment(TestCase):
 
     class EnvironmentCredentialsNotRequired(Environment):
         def __init__(self):
-            config = EnvironmentConfig('test', 'test', UserCreds())
+            config = StubEnvironmentConfig('test', 'test', UserCreds())
             super().__init__(config)
 
         _credentials_required = False
@@ -31,7 +57,7 @@ def get_auth_users(self):
 
     class EnvironmentCredentialsRequired(Environment):
         def __init__(self):
-            config = EnvironmentConfig('test', 'test', UserCreds())
+            config = StubEnvironmentConfig('test', 'test', UserCreds())
             super().__init__(config)
 
         _credentials_required = True
@@ -41,7 +67,7 @@ def get_auth_users(self):
 
     class EnvironmentAlreadyRegistered(Environment):
         def __init__(self):
-            config = EnvironmentConfig('test', 'test', UserCreds('test_user', 'test_secret'))
+            config = StubEnvironmentConfig('test', 'test', UserCreds('test_user', 'test_secret'))
             super().__init__(config)
 
         _credentials_required = True
@@ -77,36 +103,37 @@ def test_try_needs_registration(self):
         self.assertTrue(env._try_needs_registration())
 
     def test_needs_registration(self):
+
         env = TestEnvironment.EnvironmentCredentialsRequired()
-        self._test_bool_env_method("needs_registration", env, config_mocks.CONFIG_WITH_CREDENTIALS, False)
-        self._test_bool_env_method("needs_registration", env, config_mocks.CONFIG_NO_CREDENTIALS, True)
-        self._test_bool_env_method("needs_registration", env, config_mocks.CONFIG_PARTIAL_CREDENTIALS, True)
+        self._test_bool_env_method("needs_registration", env, WITH_CREDENTIALS, False)
+        self._test_bool_env_method("needs_registration", env, NO_CREDENTIALS, True)
+        self._test_bool_env_method("needs_registration", env, PARTIAL_CREDENTIALS, True)
 
         env = TestEnvironment.EnvironmentCredentialsNotRequired()
-        self._test_bool_env_method("needs_registration", env, config_mocks.CONFIG_STANDARD_ENV, False)
-        self._test_bool_env_method("needs_registration", env, config_mocks.CONFIG_STANDARD_WITH_CREDENTIALS, False)
+        self._test_bool_env_method("needs_registration", env, STANDARD_ENV, False)
+        self._test_bool_env_method("needs_registration", env, STANDARD_WITH_CREDENTIALS, False)
 
     def test_is_registered(self):
         env = TestEnvironment.EnvironmentCredentialsRequired()
-        self._test_bool_env_method("_is_registered", env, config_mocks.CONFIG_WITH_CREDENTIALS, True)
-        self._test_bool_env_method("_is_registered", env, config_mocks.CONFIG_NO_CREDENTIALS, False)
-        self._test_bool_env_method("_is_registered", env, config_mocks.CONFIG_PARTIAL_CREDENTIALS, False)
+        self._test_bool_env_method("_is_registered", env, WITH_CREDENTIALS, True)
+        self._test_bool_env_method("_is_registered", env, NO_CREDENTIALS, False)
+        self._test_bool_env_method("_is_registered", env, PARTIAL_CREDENTIALS, False)
 
         env = TestEnvironment.EnvironmentCredentialsNotRequired()
-        self._test_bool_env_method("_is_registered", env, config_mocks.CONFIG_STANDARD_ENV, False)
-        self._test_bool_env_method("_is_registered", env, config_mocks.CONFIG_STANDARD_WITH_CREDENTIALS, False)
+        self._test_bool_env_method("_is_registered", env, STANDARD_ENV, False)
+        self._test_bool_env_method("_is_registered", env, STANDARD_WITH_CREDENTIALS, False)
 
     def test_is_credentials_set_up(self):
         env = TestEnvironment.EnvironmentCredentialsRequired()
-        self._test_bool_env_method("_is_credentials_set_up", env, config_mocks.CONFIG_NO_CREDENTIALS, False)
-        self._test_bool_env_method("_is_credentials_set_up", env, config_mocks.CONFIG_WITH_CREDENTIALS, True)
-        self._test_bool_env_method("_is_credentials_set_up", env, config_mocks.CONFIG_PARTIAL_CREDENTIALS, False)
+        self._test_bool_env_method("_is_credentials_set_up", env, NO_CREDENTIALS, False)
+        self._test_bool_env_method("_is_credentials_set_up", env, WITH_CREDENTIALS, True)
+        self._test_bool_env_method("_is_credentials_set_up", env, PARTIAL_CREDENTIALS, False)
 
         env = TestEnvironment.EnvironmentCredentialsNotRequired()
-        self._test_bool_env_method("_is_credentials_set_up", env, config_mocks.CONFIG_STANDARD_ENV, False)
+        self._test_bool_env_method("_is_credentials_set_up", env, STANDARD_ENV, False)
 
     def _test_bool_env_method(self, method_name: str, env: Environment, config: Dict, expected_result: bool):
-        env._config = EnvironmentConfig.get_from_json(json.dumps(config))
+        env._config = EnvironmentConfig(config)
         method = getattr(env, method_name)
         if expected_result:
             self.assertTrue(method())