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()