island: Prompt user for old data dir's deletion during Island setup if old data dir's and Island's versions mismatch

monkey/monkey_island/cc/server_setup.py

@@ -1,6 +1,7 @@
 import atexit
 import json
 import logging
+import shutil
 import sys
 from pathlib import Path
 from threading import Thread
@@ -32,6 +33,7 @@ from monkey_island.cc.services.initialize import initialize_services  # noqa: E4
 from monkey_island.cc.services.reporting.exporter_init import populate_exporter_list  # noqa: E402
 from monkey_island.cc.services.utils.network_utils import local_ip_addresses  # noqa: E402
 from monkey_island.cc.setup import island_config_options_validator  # noqa: E402
+from monkey_island.cc.setup.data_dir import OldDataError  # 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
@@ -68,6 +70,24 @@ def _setup_data_dir(island_args: IslandCmdArgs) -> Tuple[IslandConfigOptions, st
     except json.JSONDecodeError as ex:
         print(f"Error loading server config: {ex}")
         exit(1)
+    except OldDataError as ex:
+        user_response = input(
+            f"\nExisting data directory ({ex.old_data_dir}) needs to be deleted."
+            " All data from previous runs will be lost. Proceed to delete? (y/n) "
+        )
+        if user_response == "y":
+            shutil.rmtree(ex.old_data_dir)
+            print("\nOld data directory was deleted. Trying to set up again...\n")
+            return _setup_data_dir(island_args)
+        elif user_response == "n":
+            print(
+                "\nExiting. Please backup and delete the existing data directory. Then, try again."
+                "\nTo learn how to restore and use a backup, please refer to the documentation.\n"
+            )
+            exit(1)
+        else:
+            print("\nExiting. Unrecognized response, please try again.\n")
+            exit(1)
 
 
 def _exit_on_invalid_config_options(config_options: IslandConfigOptions):

monkey/monkey_island/cc/setup/data_dir.py

@@ -1,5 +1,4 @@
 import logging
-import shutil
 from pathlib import Path
 
 from common.version import get_version
@@ -7,27 +6,32 @@ from monkey_island.cc.server_utils.file_utils import create_secure_directory
 from monkey_island.cc.setup.version_file_setup import get_version_from_dir, write_version
 
 logger = logging.getLogger(__name__)
+
 _data_dir_backup_suffix = ".old"
 
 
-def setup_data_dir(data_dir_path: Path):
-    logger.info(f"Setting up data directory in {data_dir_path}.")
-    _backup_current_data_dir(data_dir_path)
+class OldDataError(Exception):
+    def __init__(self, old_data_dir: Path) -> None:
+        self.old_data_dir = old_data_dir
+
+
+def setup_data_dir(data_dir_path: Path) -> None:
+    logger.info(f"Setting up data directory at {data_dir_path}.")
+    if data_dir_path.exists():
+        logger.info(f"Data directory already exists at {data_dir_path}.")
+        _check_current_data_dir(data_dir_path)
     create_secure_directory(str(data_dir_path))
     write_version(data_dir_path)
     logger.info("Data directory set up.")
 
 
-def _backup_current_data_dir(data_dir_path: Path):
-    if _is_backup_needed(data_dir_path):
-        logger.debug("Data directory backup needed.")
-        try:
-            return _rename_data_dir(data_dir_path)
-        except FileNotFoundError:
-            logger.debug("No data directory found to backup, this is likely a first installation.")
+def _check_current_data_dir(data_dir_path: Path) -> None:
+    if _data_dir_version_mismatch_exists(data_dir_path):
+        logger.info("Version in data directory does not match the Island's version.")
+        raise OldDataError(data_dir_path)
 
 
-def _is_backup_needed(data_dir_path: Path) -> bool:
+def _data_dir_version_mismatch_exists(data_dir_path: Path) -> bool:
     try:
         data_dir_version = get_version_from_dir(data_dir_path)
     except FileNotFoundError:
@@ -39,14 +43,14 @@ def _is_backup_needed(data_dir_path: Path) -> bool:
     return island_version != data_dir_version
 
 
-def _rename_data_dir(data_dir_path: Path):
-    backup_path = _get_backup_path(data_dir_path)
-    if backup_path.is_dir():
-        shutil.rmtree(backup_path)
-    Path(data_dir_path).replace(backup_path)
-    logger.info(f"Old data directory renamed to {backup_path}.")
+# def _rename_data_dir(data_dir_path: Path):
+#     backup_path = _get_backup_path(data_dir_path)
+#     if backup_path.is_dir():
+#         shutil.rmtree(backup_path)
+#     Path(data_dir_path).replace(backup_path)
+#     logger.info(f"Old data directory renamed to {backup_path}.")
 
 
-def _get_backup_path(data_dir_path: Path) -> Path:
-    backup_dir_name = data_dir_path.name + _data_dir_backup_suffix
-    return Path(data_dir_path.parent, backup_dir_name)
+# def _get_backup_path(data_dir_path: Path) -> Path:
+#     backup_dir_name = data_dir_path.name + _data_dir_backup_suffix
+#     return Path(data_dir_path.parent, backup_dir_name)