diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index a72cea033..35879a1d4 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -23,7 +23,10 @@ from monkey_island.cc.arg_parser import IslandCmdArgs # noqa: E402 from monkey_island.cc.arg_parser import parse_cli_args # 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 GEVENT_EXCEPTION_LOG # noqa: E402 +from monkey_island.cc.server_utils.consts import ( # noqa: E402 + GEVENT_EXCEPTION_LOG, + MONGO_CONNECTION_TIMEOUT, +) from monkey_island.cc.server_utils.encryptor import initialize_encryptor # noqa: E402 from monkey_island.cc.server_utils.island_logger import reset_logger, setup_logging # noqa: E402 from monkey_island.cc.services.initialize import initialize_services # noqa: E402 @@ -32,13 +35,9 @@ from monkey_island.cc.services.utils.network_utils import local_ip_addresses # from monkey_island.cc.setup import island_config_options_validator # noqa: E402 from monkey_island.cc.setup.gevent_hub_error_handler import GeventHubErrorHandler # noqa: E402 from monkey_island.cc.setup.island_config_options import IslandConfigOptions # noqa: E402 +from monkey_island.cc.setup.mongo import mongo_setup # noqa: E402 from monkey_island.cc.setup.mongo.database_initializer import init_collections # noqa: E402 -from monkey_island.cc.setup.mongo.mongo_setup import ( # noqa: E402 - MONGO_URL, - connect_to_mongodb, - register_mongo_shutdown_callback, - start_mongodb, -) +from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcess # noqa: E402 logger = logging.getLogger(__name__) @@ -52,11 +51,11 @@ def run_monkey_island(): _configure_logging(config_options) _initialize_globals(config_options, server_config_path) + mongo_db_process = None if config_options.start_mongodb: - mongo_db_process = start_mongodb(config_options.data_dir) - register_mongo_shutdown_callback(mongo_db_process) + mongo_db_process = _start_mongodb(config_options.data_dir) - connect_to_mongodb() + _connect_to_mongodb(mongo_db_process) _configure_gevent_exception_handling(Path(config_options.data_dir)) _start_island_server(island_args.setup_only, config_options) @@ -93,6 +92,29 @@ def _initialize_globals(config_options: IslandConfigOptions, server_config_path: initialize_services(config_options.data_dir) +def _start_mongodb(data_dir: Path) -> MongoDbProcess: + mongo_db_process = mongo_setup.start_mongodb(data_dir) + mongo_setup.register_mongo_shutdown_callback(mongo_db_process) + + return mongo_db_process + + +def _connect_to_mongodb(mongo_db_process: MongoDbProcess): + try: + mongo_setup.connect_to_mongodb(MONGO_CONNECTION_TIMEOUT) + except mongo_setup.MongoDBTimeOutError as ex: + if mongo_db_process and not mongo_db_process.is_running(): + logger.error( + f"Failed to start MongoDB process. Check log at {mongo_db_process.log_file}." + ) + else: + logger.error(ex) + sys.exit(1) + except mongo_setup.MongoDBVersionError as ex: + logger.error(ex) + sys.exit(1) + + def _configure_gevent_exception_handling(data_dir): hub = gevent.hub.get_hub() @@ -107,7 +129,7 @@ def _configure_gevent_exception_handling(data_dir): def _start_island_server(should_setup_only, config_options: IslandConfigOptions): populate_exporter_list() - app = init_app(MONGO_URL) + app = init_app(mongo_setup.MONGO_URL) init_collections() diff --git a/monkey/monkey_island/cc/server_utils/consts.py b/monkey/monkey_island/cc/server_utils/consts.py index fb5edbfab..30749cb3e 100644 --- a/monkey/monkey_island/cc/server_utils/consts.py +++ b/monkey/monkey_island/cc/server_utils/consts.py @@ -34,6 +34,7 @@ _MONGO_EXECUTABLE_PATH_LINUX = os.path.join(_MONGO_BINARY_DIR, "bin", "mongod") MONGO_EXECUTABLE_PATH = ( _MONGO_EXECUTABLE_PATH_WIN if is_windows_os() else _MONGO_EXECUTABLE_PATH_LINUX ) +MONGO_CONNECTION_TIMEOUT = 15 DEFAULT_SERVER_CONFIG_PATH = str(Path(MONKEY_ISLAND_ABS_PATH, "cc", SERVER_CONFIG_FILENAME)) diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_db_process.py b/monkey/monkey_island/cc/setup/mongo/mongo_db_process.py index 0a3b55d19..db3f5c0ca 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_db_process.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_db_process.py @@ -29,7 +29,7 @@ class MongoDbProcess: self._mongo_run_cmd, stderr=subprocess.STDOUT, stdout=log ) - logger.info("MongoDB launched successfully!") + logger.info("MongoDB has been launched!") def stop(self): if not self._process: @@ -47,3 +47,13 @@ class MongoDbProcess: f"MongoDB did not terminate gracefully and will be forcefully killed: {te}" ) self._process.kill() + + def is_running(self) -> bool: + if self._process.poll() is None: + return True + + return False + + @property + def log_file(self) -> str: + return self._log_file diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index 02e097c68..196ad54bf 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py @@ -1,8 +1,8 @@ import atexit import logging import os -import sys import time +from pathlib import Path from monkey_island.cc.database import get_db_version, is_db_server_up from monkey_island.cc.server_utils.file_utils import create_secure_directory @@ -21,7 +21,7 @@ MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0" logger = logging.getLogger(__name__) -def start_mongodb(data_dir: str) -> MongoDbProcess: +def start_mongodb(data_dir: Path) -> MongoDbProcess: db_dir = _create_db_dir(data_dir) log_file = os.path.join(data_dir, MONGO_LOG_FILENAME) @@ -43,15 +43,21 @@ def register_mongo_shutdown_callback(mongo_db_process: MongoDbProcess): atexit.register(mongo_db_process.stop) -def connect_to_mongodb(): - _wait_for_mongo_db_server(MONGO_URL) +def connect_to_mongodb(timeout: float): + _wait_for_mongo_db_server(MONGO_URL, timeout) _assert_mongo_db_version(MONGO_URL) mongo_connector.connect_dal_to_mongodb() -def _wait_for_mongo_db_server(mongo_url): +def _wait_for_mongo_db_server(mongo_url, timeout): + start_time = time.time() + while not is_db_server_up(mongo_url): - logger.info("Waiting for MongoDB server on {0}".format(mongo_url)) + logger.info(f"Waiting for MongoDB server on {mongo_url}") + + if (time.time() - start_time) > timeout: + raise MongoDBTimeOutError(f"Failed to connect to MongoDB after {timeout} seconds.") + time.sleep(1) @@ -64,11 +70,16 @@ def _assert_mongo_db_version(mongo_url): 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) - ) + raise MongoDBVersionError( + f"Mongo DB version too old. {required_version} is required, but got {server_version}." ) - sys.exit(-1) else: - logger.info("Mongo DB version OK. Got {0}".format(str(server_version))) + logger.info(f"Mongo DB version OK. Got {server_version}") + + +class MongoDBTimeOutError(Exception): + pass + + +class MongoDBVersionError(Exception): + pass diff --git a/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py b/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py new file mode 100644 index 000000000..502e7dbfe --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py @@ -0,0 +1,16 @@ +import pytest + +from monkey_island.cc.setup.mongo import mongo_setup + + +def test_connect_to_mongodb_timeout(monkeypatch): + monkeypatch.setattr(mongo_setup, "is_db_server_up", lambda _: False) + with pytest.raises(mongo_setup.MongoDBTimeOutError): + mongo_setup.connect_to_mongodb(0.0000000001) + + +def test_connect_to_mongodb_version_too_old(monkeypatch): + monkeypatch.setattr(mongo_setup, "is_db_server_up", lambda _: True) + monkeypatch.setattr(mongo_setup, "get_db_version", lambda _: ("1", "0", "0")) + with pytest.raises(mongo_setup.MongoDBVersionError): + mongo_setup.connect_to_mongodb(0)