Merge pull request #1530 from guardicore/version_file_in_data_dir

Version file in data dir

CHANGELOG.md

@@ -56,6 +56,8 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - ATT&CK report bug that said that the technique "`.bash_profile` and `.bashrc`" was not attempted
   when it actually was attempted but failed. #1511
 - Bug that periodically cleared the telemetry table's filter. #1392
+- Crashes, stack traces, and other malfunctions when data from older versions of Infection Monkey is
+  present in the data directory. #1114
 
 ### Security
 - Generate a random password when creating a new user for CommunicateAsNewUser

monkey/monkey_island/cc/server_setup.py

@@ -32,6 +32,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 IncompatibleDataDirectory  # 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
@@ -61,13 +62,15 @@ def run_monkey_island():
 
 def _setup_data_dir(island_args: IslandCmdArgs) -> Tuple[IslandConfigOptions, str]:
     try:
-        return config_setup.setup_data_dir(island_args)
+        return config_setup.setup_server_config(island_args)
     except OSError as ex:
         print(f"Error opening server config file: {ex}")
         exit(1)
     except json.JSONDecodeError as ex:
         print(f"Error loading server config: {ex}")
         exit(1)
+    except IncompatibleDataDirectory:
+        exit(1)
 
 
 def _exit_on_invalid_config_options(config_options: IslandConfigOptions):

monkey/monkey_island/cc/setup/config_setup.py

@@ -4,11 +4,11 @@ from common.utils.file_utils import expand_path
 from monkey_island.cc.arg_parser import IslandCmdArgs
 from monkey_island.cc.environment import server_config_handler
 from monkey_island.cc.server_utils.consts import DEFAULT_SERVER_CONFIG_PATH
-from monkey_island.cc.server_utils.file_utils import create_secure_directory
+from monkey_island.cc.setup.data_dir import setup_data_dir
 from monkey_island.cc.setup.island_config_options import IslandConfigOptions
 
 
-def setup_data_dir(island_args: IslandCmdArgs) -> Tuple[IslandConfigOptions, str]:
+def setup_server_config(island_args: IslandCmdArgs) -> Tuple[IslandConfigOptions, str]:
     if island_args.server_config_path:
         return _setup_config_by_cmd_arg(island_args.server_config_path)
 
@@ -18,7 +18,12 @@ def setup_data_dir(island_args: IslandCmdArgs) -> Tuple[IslandConfigOptions, str
 def _setup_config_by_cmd_arg(server_config_path) -> Tuple[IslandConfigOptions, str]:
     server_config_path = expand_path(server_config_path)
     config = server_config_handler.load_server_config_from_file(server_config_path)
-    create_secure_directory(str(config.data_dir))
+
+    # TODO refactor like in https://github.com/guardicore/monkey/pull/1528 because
+    # there's absolutely no reason to be exposed to IslandConfigOptions extraction logic
+    # if you want to modify data directory related code.
+    setup_data_dir(config.data_dir)
+
     return config, server_config_path
 
 
@@ -26,7 +31,10 @@ def _setup_default_config() -> Tuple[IslandConfigOptions, str]:
     default_config = server_config_handler.load_server_config_from_file(DEFAULT_SERVER_CONFIG_PATH)
     default_data_dir = default_config.data_dir
 
-    create_secure_directory(str(default_data_dir))
+    # TODO refactor like in https://github.com/guardicore/monkey/pull/1528 because
+    # there's absolutely no reason to be exposed to IslandConfigOptions extraction logic
+    # if you want to modify data directory related code.
+    setup_data_dir(default_data_dir)
 
     server_config_path = server_config_handler.create_default_server_config_file(default_data_dir)
     config = server_config_handler.load_server_config_from_file(server_config_path)

monkey/monkey_island/cc/setup/data_dir.py

@@ -0,0 +1,60 @@
+import logging
+import shutil
+from pathlib import Path
+
+from common.version import get_version
+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__)
+
+
+class IncompatibleDataDirectory(Exception):
+    pass
+
+
+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() and _data_dir_version_mismatch_exists(data_dir_path):
+        logger.info("Version in data directory does not match the Island's version.")
+        _handle_old_data_directory(data_dir_path)
+    create_secure_directory(str(data_dir_path))
+    write_version(data_dir_path)
+    logger.info("Data directory set up.")
+
+
+def _handle_old_data_directory(data_dir_path: Path) -> None:
+    should_delete_data_directory = _prompt_user_to_delete_data_directory(data_dir_path)
+    if should_delete_data_directory:
+        shutil.rmtree(data_dir_path)
+        logger.info(f"{data_dir_path} was deleted.")
+    else:
+        logger.error(
+            "Unable to set up data directory. Please backup and delete the existing data directory"
+            f" ({data_dir_path}). Then, try again. To learn how to restore and use a backup, please"
+            " refer to the documentation."
+        )
+        raise IncompatibleDataDirectory()
+
+
+def _prompt_user_to_delete_data_directory(data_dir_path: Path) -> bool:
+    user_response = input(
+        f"\nExisting data directory ({data_dir_path}) needs to be deleted."
+        " All data from previous runs will be lost. Proceed to delete? (y/n) "
+    )
+    print()
+    if user_response.lower() in {"y", "yes"}:
+        return True
+    return False
+
+
+def _data_dir_version_mismatch_exists(data_dir_path: Path) -> bool:
+    try:
+        data_dir_version = get_version_from_dir(data_dir_path)
+    except FileNotFoundError:
+        logger.debug("Version file not found.")
+        return True
+
+    island_version = get_version()
+
+    return island_version != data_dir_version

monkey/monkey_island/cc/setup/version_file_setup.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+
+from common.version import get_version
+
+_version_filename = "VERSION"
+
+
+def get_version_from_dir(dir_path: Path) -> str:
+    version_file_path = dir_path / _version_filename
+    return version_file_path.read_text()
+
+
+def write_version(dir_path: Path):
+    version_file_path = dir_path / _version_filename
+    version_file_path.write_text(get_version())

monkey/tests/unit_tests/monkey_island/cc/setup/test_data_dir.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+
+import pytest
+
+from monkey_island.cc.setup.data_dir import IncompatibleDataDirectory, setup_data_dir
+from monkey_island.cc.setup.version_file_setup import _version_filename
+
+current_version = "1.1.1"
+old_version = "1.1.0"
+
+
+@pytest.fixture(autouse=True)
+def mock_version(monkeypatch):
+    monkeypatch.setattr("monkey_island.cc.setup.data_dir.get_version", lambda: current_version)
+    monkeypatch.setattr(
+        "monkey_island.cc.setup.version_file_setup.get_version", lambda: current_version
+    )
+
+
+@pytest.fixture
+def temp_data_dir_path(tmp_path) -> Path:
+    return tmp_path / "data_dir"
+
+
+@pytest.fixture
+def temp_version_file_path(temp_data_dir_path) -> Path:
+    return temp_data_dir_path / _version_filename
+
+
+def test_setup_data_dir(temp_data_dir_path, temp_version_file_path):
+    data_dir_path = temp_data_dir_path
+    setup_data_dir(data_dir_path)
+    assert data_dir_path.is_dir()
+
+    version_file_path = temp_version_file_path
+    assert version_file_path.read_text() == current_version
+
+
+def test_old_version_removed(monkeypatch, temp_data_dir_path, temp_version_file_path):
+    monkeypatch.setattr("builtins.input", lambda _: "y")
+
+    temp_data_dir_path.mkdir()
+    temp_version_file_path.write_text(old_version)
+    bogus_file_path = temp_data_dir_path / "test.txt"
+    bogus_file_path.touch()
+
+    setup_data_dir(temp_data_dir_path)
+
+    assert temp_version_file_path.read_text() == current_version
+    assert not bogus_file_path.is_file()
+
+
+@pytest.mark.parametrize("input_value", ["n", "x"])
+def test_old_version_not_removed(
+    monkeypatch, temp_data_dir_path, temp_version_file_path, input_value
+):
+    monkeypatch.setattr("builtins.input", lambda _: input_value)
+
+    temp_data_dir_path.mkdir()
+    temp_version_file_path.write_text(old_version)
+    bogus_file_path = temp_data_dir_path / "test.txt"
+    bogus_file_path.touch()
+
+    with pytest.raises(IncompatibleDataDirectory):
+        setup_data_dir(temp_data_dir_path)
+
+    assert temp_version_file_path.read_text() == old_version
+    assert bogus_file_path.is_file()
+
+
+def test_data_dir_setup_not_needed(temp_data_dir_path, temp_version_file_path):
+    temp_data_dir_path.mkdir()
+    temp_version_file_path.write_text(current_version)
+    bogus_file_path = temp_data_dir_path / "test.txt"
+    bogus_file_path.touch()
+
+    setup_data_dir(temp_data_dir_path)
+    assert temp_version_file_path.read_text() == current_version
+    assert bogus_file_path.is_file()