From f0e3d76501d835487f5a2ab85a467041da30e1ad Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 3 Aug 2021 09:59:03 +0200 Subject: [PATCH 1/7] Island: Check if mongodb is running. --- monkey/monkey_island/cc/server_setup.py | 12 +++++++++++- .../cc/setup/mongo/mongo_db_process.py | 18 +++++++++++++++++- .../cc/setup/mongo/mongo_setup.py | 13 +++++++++---- 3 files changed, 37 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index a72cea033..fe0dec126 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -33,6 +33,7 @@ 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.database_initializer import init_collections # noqa: E402 +from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcessException # noqa: E402 from monkey_island.cc.setup.mongo.mongo_setup import ( # noqa: E402 MONGO_URL, connect_to_mongodb, @@ -52,11 +53,20 @@ 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) - connect_to_mongodb() + try: + connect_to_mongodb(mongo_db_process) + except MongoDbProcessException: + logger.error( + f"MongoDB could not start. For details, check the MongoDB log at " + f"{mongo_db_process.get_log_file()}" + ) + sys.exit(-1) _configure_gevent_exception_handling(Path(config_options.data_dir)) _start_island_server(island_args.setup_only, config_options) 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..ffa531647 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,10 @@ class MongoDbProcess: self._mongo_run_cmd, stderr=subprocess.STDOUT, stdout=log ) - logger.info("MongoDB launched successfully!") + if not self.is_running(): + raise MongoDbProcessException + + logger.info("MongoDB has been launched!") def stop(self): if not self._process: @@ -47,3 +50,16 @@ 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 + + def get_log_file(self) -> str: + return self._log_file + + +class MongoDbProcessException(Exception): + pass diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index 02e097c68..449beeb72 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py @@ -8,7 +8,7 @@ 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 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_db_process import MongoDbProcess +from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcess, MongoDbProcessException DB_DIR_NAME = "db" MONGO_LOG_FILENAME = "mongodb.log" @@ -26,6 +26,7 @@ def start_mongodb(data_dir: str) -> MongoDbProcess: log_file = os.path.join(data_dir, MONGO_LOG_FILENAME) mongo_db_process = MongoDbProcess(db_dir=db_dir, log_file=log_file) + mongo_db_process.start() return mongo_db_process @@ -43,15 +44,19 @@ 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(mongo_db_process: MongoDbProcess): + _wait_for_mongo_db_server(MONGO_URL, mongo_db_process) _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, mongo_db_process): while not is_db_server_up(mongo_url): logger.info("Waiting for MongoDB server on {0}".format(mongo_url)) + + if mongo_db_process is not None and not mongo_db_process.is_running(): + raise MongoDbProcessException + time.sleep(1) From 8dc2905c715d7591f114bb3a9d63441040873a0f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 3 Aug 2021 14:37:12 +0200 Subject: [PATCH 2/7] Island: Add Timeout and Version exceptions to mongo setup. --- monkey/monkey_island/cc/server_setup.py | 22 +++++++++------ .../cc/setup/mongo/mongo_db_process.py | 10 ++----- .../cc/setup/mongo/mongo_setup.py | 27 ++++++++++++------- 3 files changed, 34 insertions(+), 25 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index fe0dec126..304ec9788 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -33,9 +33,11 @@ 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.database_initializer import init_collections # noqa: E402 -from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcessException # noqa: E402 from monkey_island.cc.setup.mongo.mongo_setup import ( # noqa: E402 MONGO_URL, + TIMEOUT, + MongoDBTimeOutException, + MongoDBVersionException, connect_to_mongodb, register_mongo_shutdown_callback, start_mongodb, @@ -60,13 +62,17 @@ def run_monkey_island(): register_mongo_shutdown_callback(mongo_db_process) try: - connect_to_mongodb(mongo_db_process) - except MongoDbProcessException: - logger.error( - f"MongoDB could not start. For details, check the MongoDB log at " - f"{mongo_db_process.get_log_file()}" - ) - sys.exit(-1) + connect_to_mongodb() + except MongoDBTimeOutException: + if config_options.start_mongodb 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(f"Failed to connect to MongoDB after {TIMEOUT} seconds. ") + sys.exit(1) + except MongoDBVersionException: + sys.exit(1) _configure_gevent_exception_handling(Path(config_options.data_dir)) _start_island_server(island_args.setup_only, config_options) 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 ffa531647..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,9 +29,6 @@ class MongoDbProcess: self._mongo_run_cmd, stderr=subprocess.STDOUT, stdout=log ) - if not self.is_running(): - raise MongoDbProcessException - logger.info("MongoDB has been launched!") def stop(self): @@ -57,9 +54,6 @@ class MongoDbProcess: return False - def get_log_file(self) -> str: + @property + def log_file(self) -> str: return self._log_file - - -class MongoDbProcessException(Exception): - pass diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index 449beeb72..df0470f3b 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py @@ -1,14 +1,13 @@ import atexit 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.server_utils.file_utils import create_secure_directory 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_db_process import MongoDbProcess, MongoDbProcessException +from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcess DB_DIR_NAME = "db" MONGO_LOG_FILENAME = "mongodb.log" @@ -17,6 +16,7 @@ MONGO_URL = os.environ.get( "mongodb://{0}:{1}/{2}".format(MONGO_DB_HOST, MONGO_DB_PORT, MONGO_DB_NAME), ) MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0" +TIMEOUT = 15 logger = logging.getLogger(__name__) @@ -26,7 +26,6 @@ def start_mongodb(data_dir: str) -> MongoDbProcess: log_file = os.path.join(data_dir, MONGO_LOG_FILENAME) mongo_db_process = MongoDbProcess(db_dir=db_dir, log_file=log_file) - mongo_db_process.start() return mongo_db_process @@ -44,18 +43,20 @@ def register_mongo_shutdown_callback(mongo_db_process: MongoDbProcess): atexit.register(mongo_db_process.stop) -def connect_to_mongodb(mongo_db_process: MongoDbProcess): - _wait_for_mongo_db_server(MONGO_URL, mongo_db_process) +def connect_to_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, mongo_db_process): +def _wait_for_mongo_db_server(mongo_url): + start_time = time.time() + while not is_db_server_up(mongo_url): logger.info("Waiting for MongoDB server on {0}".format(mongo_url)) - if mongo_db_process is not None and not mongo_db_process.is_running(): - raise MongoDbProcessException + if (time.time() - start_time) > TIMEOUT: + raise MongoDBTimeOutException time.sleep(1) @@ -74,6 +75,14 @@ def _assert_mongo_db_version(mongo_url): str(required_version), str(server_version) ) ) - sys.exit(-1) + raise MongoDBVersionException else: logger.info("Mongo DB version OK. Got {0}".format(str(server_version))) + + +class MongoDBTimeOutException(Exception): + pass + + +class MongoDBVersionException(Exception): + pass From 109a9922011c11e22880a8880f19421f5e24770f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 3 Aug 2021 09:40:51 -0400 Subject: [PATCH 3/7] Island: Add error messages to Mongo exceptions --- monkey/monkey_island/cc/server_setup.py | 8 ++++---- monkey/monkey_island/cc/setup/mongo/mongo_setup.py | 9 +++------ 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 304ec9788..ef44ca6fa 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -35,7 +35,6 @@ from monkey_island.cc.setup.island_config_options import IslandConfigOptions # 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, - TIMEOUT, MongoDBTimeOutException, MongoDBVersionException, connect_to_mongodb, @@ -63,15 +62,16 @@ def run_monkey_island(): try: connect_to_mongodb() - except MongoDBTimeOutException: + except MongoDBTimeOutException as ex: if config_options.start_mongodb 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(f"Failed to connect to MongoDB after {TIMEOUT} seconds. ") + logger.error(ex) sys.exit(1) - except MongoDBVersionException: + except MongoDBVersionException as ex: + logger.error(ex) sys.exit(1) _configure_gevent_exception_handling(Path(config_options.data_dir)) diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index df0470f3b..6ebd6e89b 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py @@ -56,7 +56,7 @@ def _wait_for_mongo_db_server(mongo_url): logger.info("Waiting for MongoDB server on {0}".format(mongo_url)) if (time.time() - start_time) > TIMEOUT: - raise MongoDBTimeOutException + raise MongoDBTimeOutException(f"Failed to connect to MongoDB after {TIMEOUT} seconds.") time.sleep(1) @@ -70,12 +70,9 @@ 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 MongoDBVersionException( + f"Mongo DB version too old. {required_version} is required, but got {server_version}." ) - raise MongoDBVersionException else: logger.info("Mongo DB version OK. Got {0}".format(str(server_version))) From ba1e9be8f89416f0a18dcf070553c414e065a291 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 3 Aug 2021 09:43:10 -0400 Subject: [PATCH 4/7] Island: Change .format() to format string in mongo_setup.py --- monkey/monkey_island/cc/setup/mongo/mongo_setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index 6ebd6e89b..f5b51e514 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py @@ -53,7 +53,7 @@ def _wait_for_mongo_db_server(mongo_url): 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 MongoDBTimeOutException(f"Failed to connect to MongoDB after {TIMEOUT} seconds.") @@ -74,7 +74,7 @@ def _assert_mongo_db_version(mongo_url): f"Mongo DB version too old. {required_version} is required, but got {server_version}." ) else: - logger.info("Mongo DB version OK. Got {0}".format(str(server_version))) + logger.info(f"Mongo DB version OK. Got {server_version}") class MongoDBTimeOutException(Exception): From dbf1ad86bca26d41d4dabe5fb25d1a9215c4de35 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 3 Aug 2021 10:06:26 -0400 Subject: [PATCH 5/7] Island: Extract start and connect mongodb functions --- monkey/monkey_island/cc/server_setup.py | 53 +++++++++++++------------ 1 file changed, 28 insertions(+), 25 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index ef44ca6fa..bab48b7ea 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -32,15 +32,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, - MongoDBTimeOutException, - MongoDBVersionException, - 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__) @@ -55,24 +49,10 @@ def run_monkey_island(): _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) - try: - connect_to_mongodb() - except MongoDBTimeOutException as ex: - if config_options.start_mongodb 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 MongoDBVersionException as ex: - logger.error(ex) - sys.exit(1) + _connect_to_mongodb(config_options.data_dir, mongo_db_process) _configure_gevent_exception_handling(Path(config_options.data_dir)) _start_island_server(island_args.setup_only, config_options) @@ -109,6 +89,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(data_dir: Path, mongo_db_process: MongoDbProcess): + try: + mongo_setup.connect_to_mongodb() + except mongo_setup.MongoDBTimeOutException 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.MongoDBVersionException as ex: + logger.error(ex) + sys.exit(1) + + def _configure_gevent_exception_handling(data_dir): hub = gevent.hub.get_hub() @@ -123,7 +126,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() From c368fd6a99d8fecb3dca3866d5b6f957bcd9358d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 3 Aug 2021 10:12:17 -0400 Subject: [PATCH 6/7] Island: Rename MongoDB*Exception -> MongoDB*Error --- monkey/monkey_island/cc/server_setup.py | 4 ++-- monkey/monkey_island/cc/setup/mongo/mongo_setup.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index bab48b7ea..65236509d 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -99,7 +99,7 @@ def _start_mongodb(data_dir: Path) -> MongoDbProcess: def _connect_to_mongodb(data_dir: Path, mongo_db_process: MongoDbProcess): try: mongo_setup.connect_to_mongodb() - except mongo_setup.MongoDBTimeOutException as ex: + 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}." @@ -107,7 +107,7 @@ def _connect_to_mongodb(data_dir: Path, mongo_db_process: MongoDbProcess): else: logger.error(ex) sys.exit(1) - except mongo_setup.MongoDBVersionException as ex: + except mongo_setup.MongoDBVersionError as ex: logger.error(ex) sys.exit(1) diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index f5b51e514..d9440abfe 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py @@ -56,7 +56,7 @@ def _wait_for_mongo_db_server(mongo_url): logger.info(f"Waiting for MongoDB server on {mongo_url}") if (time.time() - start_time) > TIMEOUT: - raise MongoDBTimeOutException(f"Failed to connect to MongoDB after {TIMEOUT} seconds.") + raise MongoDBTimeOutError(f"Failed to connect to MongoDB after {TIMEOUT} seconds.") time.sleep(1) @@ -70,16 +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: - raise MongoDBVersionException( + raise MongoDBVersionError( f"Mongo DB version too old. {required_version} is required, but got {server_version}." ) else: logger.info(f"Mongo DB version OK. Got {server_version}") -class MongoDBTimeOutException(Exception): +class MongoDBTimeOutError(Exception): pass -class MongoDBVersionException(Exception): +class MongoDBVersionError(Exception): pass From 374b1a7063a3eea98be9f60258496ca0fa76bfd2 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 3 Aug 2021 16:26:50 +0200 Subject: [PATCH 7/7] UT: Add test for mongo timeout and mongo version --- monkey/monkey_island/cc/server_setup.py | 11 +++++++---- monkey/monkey_island/cc/server_utils/consts.py | 1 + .../monkey_island/cc/setup/mongo/mongo_setup.py | 14 +++++++------- .../cc/setup/mongo/test_mongo_setup.py | 16 ++++++++++++++++ 4 files changed, 31 insertions(+), 11 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_mongo_setup.py diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 65236509d..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 @@ -52,7 +55,7 @@ def run_monkey_island(): if config_options.start_mongodb: mongo_db_process = _start_mongodb(config_options.data_dir) - _connect_to_mongodb(config_options.data_dir, mongo_db_process) + _connect_to_mongodb(mongo_db_process) _configure_gevent_exception_handling(Path(config_options.data_dir)) _start_island_server(island_args.setup_only, config_options) @@ -96,9 +99,9 @@ def _start_mongodb(data_dir: Path) -> MongoDbProcess: return mongo_db_process -def _connect_to_mongodb(data_dir: Path, mongo_db_process: MongoDbProcess): +def _connect_to_mongodb(mongo_db_process: MongoDbProcess): try: - mongo_setup.connect_to_mongodb() + 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( 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_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py index d9440abfe..196ad54bf 100644 --- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py +++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py @@ -2,6 +2,7 @@ import atexit import logging import os 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 @@ -16,12 +17,11 @@ MONGO_URL = os.environ.get( "mongodb://{0}:{1}/{2}".format(MONGO_DB_HOST, MONGO_DB_PORT, MONGO_DB_NAME), ) MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0" -TIMEOUT = 15 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,20 +43,20 @@ 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(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.") + if (time.time() - start_time) > timeout: + raise MongoDBTimeOutError(f"Failed to connect to MongoDB after {timeout} seconds.") time.sleep(1) 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)