From 8dc2905c715d7591f114bb3a9d63441040873a0f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 3 Aug 2021 14:37:12 +0200 Subject: [PATCH] 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