From dd480d1703bfd6d8542d73bb1b0f934efed52ab6 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 20 Oct 2021 16:10:01 +0530 Subject: [PATCH] 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 | 20 ++++++++++ monkey/monkey_island/cc/setup/data_dir.py | 46 ++++++++++++----------- 2 files changed, 45 insertions(+), 21 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index e3e8462a8..38e7de5d8 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/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): diff --git a/monkey/monkey_island/cc/setup/data_dir.py b/monkey/monkey_island/cc/setup/data_dir.py index 590e966ba..cc810f9da 100644 --- a/monkey/monkey_island/cc/setup/data_dir.py +++ b/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)