From 559af47928e0edce27d194cf9e300c78d4ee4e00 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Wed, 2 Jun 2021 14:20:21 -0400
Subject: [PATCH] island: Clean up MongoDB subprocess when Monkey Island shuts
 down

---
 monkey/monkey_island/cc/server_setup.py        |  6 +++++-
 .../cc/setup/mongo/mongo_db_process.py         | 17 ++++++++++++++++-
 .../cc/setup/mongo/mongo_setup.py              | 18 +++++++++++++-----
 3 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py
index e6895ac59..ba0ed1411 100644
--- a/monkey/monkey_island/cc/server_setup.py
+++ b/monkey/monkey_island/cc/server_setup.py
@@ -28,6 +28,7 @@ from monkey_island.cc.setup.mongo.database_initializer import init_collections
 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.setup.config_setup import setup_data_dir  # noqa: E402
@@ -43,7 +44,10 @@ def run_monkey_island():
     _configure_logging(config_options)
     _initialize_global_resources(config_options, server_config_path)
 
-    start_mongodb(config_options)
+    if config_options.start_mongodb:
+        mongo_db_process = start_mongodb(config_options)
+        register_mongo_shutdown_callback(mongo_db_process)
+
     connect_to_mongodb()
 
     bootloader_server_thread = _start_bootloader_server(MONGO_URL)
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 c8b043249..f24406630 100644
--- a/monkey/monkey_island/cc/setup/mongo/mongo_db_process.py
+++ b/monkey/monkey_island/cc/setup/mongo/mongo_db_process.py
@@ -11,6 +11,7 @@ logger = logging.getLogger(__name__)
 DB_DIR_NAME = "db"
 DB_DIR_PARAM = "--dbpath"
 MONGO_LOG_FILENAME = "mongo_log.txt"
+TERMINATE_TIMEOUT = 10
 
 
 class MongoDbProcess:
@@ -21,11 +22,25 @@ class MongoDbProcess:
         """
         self.db_dir_parent_path = db_dir_parent_path
         self.logging_dir_path = logging_dir_path
+        self._process = None
 
     def start(self):
         db_path = self._create_db_dir()
         self._start_mongodb_process(db_path)
 
+    def stop(self):
+        if self._process:
+            logger.info("Terminating MongoDB process")
+            self._process.terminate()
+            try:
+                self._process.wait(timeout=TERMINATE_TIMEOUT)
+                logger.info("MongoDB process terminated successfully")
+            except subprocess.TimeoutExpired as te:
+                logger.warning(
+                    f"MongoDB did not terminate gracefully and will be forcefully killed: {te}"
+                )
+                self._process.kill()
+
     def _create_db_dir(self) -> str:
         db_path = os.path.join(self.db_dir_parent_path, DB_DIR_NAME)
         logger.info(f"Database content directory: {db_path}.")
@@ -42,7 +57,7 @@ class MongoDbProcess:
         logger.info(f"Mongodb log will be available at {mongo_log_path}.")
 
         with open(mongo_log_path, "w") as log:
-            subprocess.Popen(mongo_run_cmd, stderr=subprocess.STDOUT, stdout=log)
+            self._process = subprocess.Popen(mongo_run_cmd, stderr=subprocess.STDOUT, stdout=log)
         logger.info("MongoDb launched successfully!")
 
     @staticmethod
diff --git a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py
index 7d83ee5f4..412ff49d6 100644
--- a/monkey/monkey_island/cc/setup/mongo/mongo_setup.py
+++ b/monkey/monkey_island/cc/setup/mongo/mongo_setup.py
@@ -1,3 +1,4 @@
+import atexit
 import logging
 import os
 import sys
@@ -18,11 +19,18 @@ MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0"
 logger = logging.getLogger(__name__)
 
 
-def start_mongodb(config_options: IslandConfigOptions):
-    if config_options.start_mongodb:
-        MongoDbProcess(
-            db_dir_parent_path=config_options.data_dir, logging_dir_path=config_options.data_dir
-        ).start()
+def start_mongodb(config_options: IslandConfigOptions) -> MongoDbProcess:
+    mongo_db_process = MongoDbProcess(
+        db_dir_parent_path=config_options.data_dir, logging_dir_path=config_options.data_dir
+    )
+
+    mongo_db_process.start()
+
+    return mongo_db_process
+
+
+def register_mongo_shutdown_callback(mongo_db_process: MongoDbProcess):
+    atexit.register(mongo_db_process.stop)
 
 
 def connect_to_mongodb():