Merge pull request #1395 from guardicore/1354/check-if-mongodb-has-started

Island: Check if mongodb is running in wait_for_mongodb.

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

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

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

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

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)