Island: Extract start and connect mongodb functions

This commit is contained in:
Mike Salvatore 2021-08-03 10:06:26 -04:00
parent ba1e9be8f8
commit dbf1ad86bc
1 changed files with 28 additions and 25 deletions

View File

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