Merge pull request #1188 from guardicore/mongodb_connection_setup_refactoring

Mongodb connection setup refactoring

monkey/monkey_island.py

@@ -3,8 +3,8 @@ import monkey_island.setup.gevent_setup  # noqa: F401 isort:skip
 
 import json
 
-import monkey_island.cc.environment.environment_singleton as env_singleton
 from monkey_island.cc.arg_parser import parse_cli_args
+from monkey_island.cc.server_setup import setup_island
 from monkey_island.cc.server_utils.island_logger import setup_logging
 from monkey_island.setup.config_setup import setup_config_by_cmd_arg, setup_default_config
 
@@ -29,11 +29,4 @@ if "__main__" == __name__:
         print(f"Error loading server config: {ex}")
         exit(1)
 
-    # We need to initialize environment singleton before importing main,
-    # because main imports modules from monkey_island/cc/models and models need a connection to the
-    # mongodb. Mongodb connection parameters are initialized in environment singleton.
-    env_singleton.initialize_from_file(server_config_path)
-
-    from monkey_island.cc.main import main  # noqa: E402
-
-    main(island_args.setup_only, config)
+    setup_island(island_args.setup_only, config, server_config_path)

monkey/monkey_island/cc/environment/__init__.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from abc import ABCMeta, abstractmethod
 from datetime import timedelta
 
@@ -18,13 +17,6 @@ logger = logging.getLogger(__name__)
 
 class Environment(object, metaclass=ABCMeta):
     _ISLAND_PORT = 5000
-    _MONGO_DB_NAME = "monkeyisland"
-    _MONGO_DB_HOST = "localhost"
-    _MONGO_DB_PORT = 27017
-    _MONGO_URL = os.environ.get(
-        "MONKEY_MONGO_URL",
-        "mongodb://{0}:{1}/{2}".format(_MONGO_DB_HOST, _MONGO_DB_PORT, str(_MONGO_DB_NAME)),
-    )
     _DEBUG_SERVER = False
     _AUTH_EXPIRATION_TIME = timedelta(minutes=30)
 
@@ -96,9 +88,6 @@ class Environment(object, metaclass=ABCMeta):
     def get_island_port(self):
         return self._ISLAND_PORT
 
-    def get_mongo_url(self):
-        return self._MONGO_URL
-
     def is_debug(self):
         return self._DEBUG_SERVER
 
@@ -110,15 +99,3 @@ class Environment(object, metaclass=ABCMeta):
         if self._config and self._config.deployment:
             deployment = self._config.deployment
         return deployment
-
-    @property
-    def mongo_db_name(self):
-        return self._MONGO_DB_NAME
-
-    @property
-    def mongo_db_host(self):
-        return self._MONGO_DB_HOST
-
-    @property
-    def mongo_db_port(self):
-        return self._MONGO_DB_PORT

monkey/monkey_island/cc/models/__init__.py

@@ -1,7 +1,3 @@
-from mongoengine import connect
-
-import monkey_island.cc.environment.environment_singleton as env_singleton  # noqa: E402
-
 from .command_control_channel import CommandControlChannel  # noqa: F401, E402
 
 # Order of importing matters here, for registering the embedded and referenced documents before
@@ -11,10 +7,3 @@ from .creds import Creds  # noqa: F401, E402
 from .monkey import Monkey  # noqa: F401, E402
 from .monkey_ttl import MonkeyTtl  # noqa: F401, E402
 from .pba_results import PbaResults  # noqa: F401, E402
-
-# TODO refactor into explicit call when implementing mongodb startup
-connect(
-    db=env_singleton.env.mongo_db_name,
-    host=env_singleton.env.mongo_db_host,
-    port=env_singleton.env.mongo_db_port,
-)

monkey/monkey_island/cc/server_setup.py

@@ -1,66 +1,55 @@
 import logging
 import os
 import sys
-import time
 from pathlib import Path
 from threading import Thread
 
-# Add the monkey_island directory to the path, to make sure imports that don't start with
-# "monkey_island." work.
 from gevent.pywsgi import WSGIServer
 
-from monkey_island.cc.setup.database_initializer import init_collections
-from monkey_island.setup.island_config_options import IslandConfigOptions
-
+# Add the monkey_island directory to the path, to make sure imports that don't start with
+# "monkey_island." work.
 MONKEY_ISLAND_DIR_BASE_PATH = str(Path(__file__).parent.parent)
 if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path:
     sys.path.insert(0, MONKEY_ISLAND_DIR_BASE_PATH)
 
-from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH  # noqa: E402
-
-logger = logging.getLogger(__name__)
-
 import monkey_island.cc.environment.environment_singleton as env_singleton  # noqa: E402
 from common.version import get_version  # noqa: E402
 from monkey_island.cc.app import init_app  # noqa: E402
-from monkey_island.cc.database import get_db_version  # noqa: E402
-from monkey_island.cc.database import is_db_server_up  # noqa: E402
 from monkey_island.cc.resources.monkey_download import MonkeyDownload  # noqa: E402
 from monkey_island.cc.server_utils.bootloader_server import BootloaderHttpServer  # noqa: E402
+from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH  # noqa: E402
 from monkey_island.cc.server_utils.encryptor import initialize_encryptor  # noqa: E402
 from monkey_island.cc.services.initialize import initialize_services  # noqa: E402
 from monkey_island.cc.services.reporting.exporter_init import populate_exporter_list  # noqa: E402
 from monkey_island.cc.services.utils.network_utils import local_ip_addresses  # noqa: E402
-from monkey_island.cc.setup.mongo_process_runner import MongoDbRunner  # noqa: E402
+from monkey_island.cc.setup.mongo.database_initializer import init_collections  # noqa: E402
+from monkey_island.cc.setup.mongo.mongo_setup import MONGO_URL, setup_mongodb  # noqa: E402
+from monkey_island.setup.island_config_options import IslandConfigOptions  # noqa: E402
 
-MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0"
+logger = logging.getLogger(__name__)
 
 
-def main(setup_only: bool, config_options: IslandConfigOptions):
+def setup_island(setup_only: bool, config_options: IslandConfigOptions, server_config_path: str):
+    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(
-        target=BootloaderHttpServer(mongo_url).serve_forever, daemon=True
+        target=BootloaderHttpServer(MONGO_URL).serve_forever, daemon=True
     )
 
     bootloader_server_thread.start()
-    start_island_server(setup_only, config_options)
+    _start_island_server(setup_only, config_options)
     bootloader_server_thread.join()
 
 
-def start_island_server(should_setup_only, config_options: IslandConfigOptions):
-    if config_options.start_mongodb:
-        MongoDbRunner(
-            db_dir_parent_path=config_options.data_dir, logging_dir_path=config_options.data_dir
-        ).launch_mongodb()
-    mongo_url = os.environ.get("MONGO_URL", env_singleton.env.get_mongo_url())
-    wait_for_mongo_db_server(mongo_url)
-    assert_mongo_db_version(mongo_url)
+def _start_island_server(should_setup_only, config_options: IslandConfigOptions):
+
+    setup_mongodb(config_options)
 
     populate_exporter_list()
-    app = init_app(mongo_url)
+    app = init_app(MONGO_URL)
 
     crt_path = str(Path(MONKEY_ISLAND_ABS_PATH, "cc", "server.crt"))
     key_path = str(Path(MONKEY_ISLAND_ABS_PATH, "cc", "server.key"))
@@ -80,11 +69,11 @@ def start_island_server(should_setup_only, config_options: IslandConfigOptions):
             certfile=os.environ.get("SERVER_CRT", crt_path),
             keyfile=os.environ.get("SERVER_KEY", key_path),
         )
-        log_init_info()
+        _log_init_info()
         http_server.serve_forever()
 
 
-def log_init_info():
+def _log_init_info():
     logger.info("Monkey Island Server is running!")
     logger.info(f"version: {get_version()}")
     logger.info(
@@ -98,28 +87,3 @@ def log_init_info():
         )
     )
     MonkeyDownload.log_executable_hashes()
-
-
-def wait_for_mongo_db_server(mongo_url):
-    while not is_db_server_up(mongo_url):
-        logger.info("Waiting for MongoDB server on {0}".format(mongo_url))
-        time.sleep(1)
-
-
-def assert_mongo_db_version(mongo_url):
-    """
-    Checks if the mongodb version is new enough for running the app.
-    If the DB is too old, quits.
-    :param mongo_url: URL to the mongo the Island will use
-    """
-    required_version = tuple(MINIMUM_MONGO_DB_VERSION_REQUIRED.split("."))
-    server_version = get_db_version(mongo_url)
-    if server_version < required_version:
-        logger.error(
-            "Mongo DB version too old. {0} is required, but got {1}".format(
-                str(required_version), str(server_version)
-            )
-        )
-        sys.exit(-1)
-    else:
-        logger.info("Mongo DB version OK. Got {0}".format(str(server_version)))

monkey/monkey_island/cc/setup/mongo/mongo_connector.py

@@ -0,0 +1,9 @@
+from mongoengine import connect
+
+MONGO_DB_NAME = "monkeyisland"
+MONGO_DB_HOST = "localhost"
+MONGO_DB_PORT = 27017
+
+
+def connect_dal_to_mongodb(db=MONGO_DB_NAME, host=MONGO_DB_HOST, port=MONGO_DB_PORT):
+    connect(db=db, host=host, port=port)

monkey/monkey_island/cc/setup/mongo/mongo_setup.py

@@ -0,0 +1,53 @@
+import logging
+import os
+import sys
+import time
+
+from monkey_island.cc.database import get_db_version, is_db_server_up
+from monkey_island.cc.setup.mongo import mongo_connector
+from monkey_island.cc.setup.mongo.mongo_connector import MONGO_DB_HOST, MONGO_DB_NAME, MONGO_DB_PORT
+from monkey_island.cc.setup.mongo.mongo_process_runner import MongoDbRunner
+from monkey_island.setup.island_config_options import IslandConfigOptions
+
+MONGO_URL = os.environ.get(
+    "MONKEY_MONGO_URL",
+    "mongodb://{0}:{1}/{2}".format(MONGO_DB_HOST, MONGO_DB_PORT, MONGO_DB_NAME),
+)
+MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0"
+
+logger = logging.getLogger(__name__)
+
+
+def setup_mongodb(config_options: IslandConfigOptions):
+    if config_options.start_mongodb:
+        MongoDbRunner(
+            db_dir_parent_path=config_options.data_dir, logging_dir_path=config_options.data_dir
+        ).launch_mongodb()
+    wait_for_mongo_db_server(MONGO_URL)
+    assert_mongo_db_version(MONGO_URL)
+    mongo_connector.connect_dal_to_mongodb()
+
+
+def wait_for_mongo_db_server(mongo_url):
+    while not is_db_server_up(mongo_url):
+        logger.info("Waiting for MongoDB server on {0}".format(mongo_url))
+        time.sleep(1)
+
+
+def assert_mongo_db_version(mongo_url):
+    """
+    Checks if the mongodb version is new enough for running the app.
+    If the DB is too old, quits.
+    :param mongo_url: URL to the mongo the Island will use
+    """
+    required_version = tuple(MINIMUM_MONGO_DB_VERSION_REQUIRED.split("."))
+    server_version = get_db_version(mongo_url)
+    if server_version < required_version:
+        logger.error(
+            "Mongo DB version too old. {0} is required, but got {1}".format(
+                str(required_version), str(server_version)
+            )
+        )
+        sys.exit(-1)
+    else:
+        logger.info("Mongo DB version OK. Got {0}".format(str(server_version)))

monkey/tests/unit_tests/monkey_island/cc/conftest.py

@@ -1,11 +1,3 @@
-import monkey_island.cc.environment.environment_singleton as env_singleton
-from monkey_island.cc.environment.testing import TestingEnvironment
-
-# Mock environment singleton because it contains mongodb parameters
-# needed for model tests. See monkey/monkey_island/cc/models/__init__.py
-env_config = {}
-env_singleton.env = TestingEnvironment(env_config)
-
 # Without these imports pytests can't use fixtures,
 # because they are not found
 from tests.unit_tests.monkey_island.cc.mongomock_fixtures import *  # noqa: F401,F403,E402

monkey/tests/unit_tests/monkey_island/cc/services/test_config.py

@@ -19,6 +19,18 @@ def config(monkeypatch):
     return config
 
 
+class MockClass:
+    pass
+
+
+@pytest.fixture(scope="function", autouse=True)
+def mock_port_in_env_singleton(monkeypatch):
+    mock_singleton = MockClass()
+    mock_singleton.env = MockClass()
+    mock_singleton.env.get_island_port = lambda: PORT
+    monkeypatch.setattr("monkey_island.cc.services.config.env_singleton", mock_singleton)
+
+
 def test_set_server_ips_in_config_command_servers(config):
     ConfigService.set_server_ips_in_config(config)
     expected_config_command_servers = [f"{ip}:{PORT}" for ip in IPS]

whitelist.py

@@ -166,8 +166,7 @@ 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)
-fake_db_dir  # unused function 'fake_db_dir' (monkey/tests/unit_tests/monkey_island/cc/setup/test_process_runner.py:10)
-
+mock_port_in_env_singleton  # monkey\tests\unit_tests\monkey_island\cc\services\test_config.py:26:
 
 # 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)