Merge pull request #1292 from guardicore/ransomware_dir_fix

Ransomware: bugfix for directories

monkey/common/utils/file_utils.py

@@ -0,0 +1,5 @@
+import os
+
+
+def expand_path(path: str) -> str:
+    return os.path.expandvars(os.path.expanduser(path))

monkey/infection_monkey/ransomware/ransomware_payload.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from pprint import pformat
 from typing import List, Optional, Tuple
 
+from common.utils.file_utils import expand_path
 from infection_monkey.ransomware.bitflip_encryptor import BitflipEncryptor
 from infection_monkey.ransomware.file_selectors import select_production_safe_target_files
 from infection_monkey.ransomware.targeted_file_extensions import TARGETED_FILE_EXTENSIONS
@@ -28,10 +29,12 @@ class RansomwarePayload:
         self._readme_enabled = config["other_behaviors"]["readme"]
 
         target_directories = config["encryption"]["directories"]
-        self._target_dir = (
-            target_directories["windows_target_dir"]
-            if is_windows_os()
-            else target_directories["linux_target_dir"]
+        self._target_dir = Path(
+            expand_path(
+                target_directories["windows_target_dir"]
+                if is_windows_os()
+                else target_directories["linux_target_dir"]
+            )
         )
 
         self._new_file_extension = EXTENSION
@@ -55,7 +58,7 @@ class RansomwarePayload:
             return []
 
         return select_production_safe_target_files(
-            Path(self._target_dir), self._valid_file_extensions_for_encryption
+            self._target_dir, self._valid_file_extensions_for_encryption
         )
 
     def _encrypt_files(self, file_list: List[Path]) -> List[Tuple[Path, Optional[Exception]]]:
@@ -84,7 +87,7 @@ class RansomwarePayload:
 
     def _leave_readme(self):
         if self._readme_enabled:
-            readme_dest_path = Path(self._target_dir) / README_DEST
+            readme_dest_path = self._target_dir / README_DEST
             LOG.info(f"Leaving a ransomware README file at {readme_dest_path}")
 
             try:

monkey/monkey_island/cc/server_utils/consts.py

@@ -1,7 +1,7 @@
 import os
 from pathlib import Path
 
-from monkey_island.cc.server_utils import file_utils
+from common.utils.file_utils import expand_path
 from monkey_island.cc.server_utils.file_utils import is_windows_os
 
 __author__ = "itay.mizeretz"
@@ -26,7 +26,7 @@ SERVER_CONFIG_FILENAME = "server_config.json"
 
 MONKEY_ISLAND_ABS_PATH = _get_monkey_island_abs_path()
 
-DEFAULT_DATA_DIR = file_utils.expand_path(get_default_data_dir())
+DEFAULT_DATA_DIR = expand_path(get_default_data_dir())
 
 DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS = 60 * 5
 
@@ -37,7 +37,7 @@ MONGO_EXECUTABLE_PATH = (
     _MONGO_EXECUTABLE_PATH_WIN if is_windows_os() else _MONGO_EXECUTABLE_PATH_LINUX
 )
 
-DEFAULT_SERVER_CONFIG_PATH = file_utils.expand_path(
+DEFAULT_SERVER_CONFIG_PATH = expand_path(
     os.path.join(MONKEY_ISLAND_ABS_PATH, "cc", SERVER_CONFIG_FILENAME)
 )
 

monkey/monkey_island/cc/server_utils/file_utils.py

@@ -20,10 +20,6 @@ if is_windows_os():
     import monkey_island.cc.server_utils.windows_permissions as windows_permissions
 
 
-def expand_path(path: str) -> str:
-    return os.path.expandvars(os.path.expanduser(path))
-
-
 def create_secure_directory(path: str):
     if not os.path.isdir(path):
         if is_windows_os():

monkey/monkey_island/cc/setup/config_setup.py

@@ -1,8 +1,8 @@
 from typing import Tuple
 
+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 import file_utils
 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.island_config_options import IslandConfigOptions
@@ -16,7 +16,7 @@ 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 = file_utils.expand_path(server_config_path)
+    server_config_path = expand_path(server_config_path)
     config = server_config_handler.load_server_config_from_file(server_config_path)
     create_secure_directory(config.data_dir)
     return config, server_config_path

monkey/monkey_island/cc/setup/island_config_options.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from common.utils.file_utils import expand_path
 from monkey_island.cc.server_utils.consts import (
     DEFAULT_CERTIFICATE_PATHS,
     DEFAULT_CRT_PATH,
@@ -8,7 +9,6 @@ from monkey_island.cc.server_utils.consts import (
     DEFAULT_LOG_LEVEL,
     DEFAULT_START_MONGO_DB,
 )
-from monkey_island.cc.server_utils.file_utils import expand_path
 
 
 class IslandConfigOptions:

monkey/tests/unit_tests/common/utils/test_common_file_utils.py

@@ -0,0 +1,17 @@
+import os
+
+from common.utils.file_utils import expand_path
+
+
+def test_expand_user(mock_home_env):
+    input_path = os.path.join("~", "test")
+    expected_path = os.path.join(mock_home_env, "test")
+
+    assert expand_path(input_path) == expected_path
+
+
+def test_expand_vars(mock_home_env):
+    input_path = os.path.join("$HOME", "test")
+    expected_path = os.path.join(mock_home_env, "test")
+
+    assert expand_path(input_path) == expected_path

monkey/tests/unit_tests/infection_monkey/ransomware/conftest.py

@@ -4,6 +4,13 @@ from pathlib import Path
 import pytest
 
 
+@pytest.fixture
+def patched_home_env(monkeypatch, tmp_path):
+    monkeypatch.setenv("HOME", str(tmp_path))
+
+    return tmp_path
+
+
 @pytest.fixture
 def ransomware_target(tmp_path, data_for_tests_dir):
     ransomware_test_data = Path(data_for_tests_dir) / "ransomware_targets"

monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py

@@ -48,6 +48,24 @@ def ransomware_payload(ransomware_payload_config, telemetry_messenger_spy):
     return RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy)
 
 
+def test_env_variables_in_target_dir_resolved_linux(
+    ransomware_payload_config, ransomware_target, telemetry_messenger_spy, patched_home_env
+):
+    path_with_env_variable = "$HOME/ransomware_target"
+
+    ransomware_payload_config["encryption"]["directories"][
+        "linux_target_dir"
+    ] = ransomware_payload_config["encryption"]["directories"][
+        "windows_target_dir"
+    ] = path_with_env_variable
+    RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy).run_payload()
+
+    assert (
+        hash_file(ransomware_target / with_extension(ALL_ZEROS_PDF))
+        == ALL_ZEROS_PDF_ENCRYPTED_SHA256
+    )
+
+
 def test_file_with_excluded_extension_not_encrypted(ransomware_target, ransomware_payload):
     ransomware_payload.run_payload()
 

monkey/tests/unit_tests/monkey_island/cc/server_utils/test_file_utils.py

@@ -6,26 +6,11 @@ from tests.monkey_island.utils import assert_windows_permissions
 
 from monkey_island.cc.server_utils.file_utils import (
     create_secure_directory,
-    expand_path,
     is_windows_os,
     open_new_securely_permissioned_file,
 )
 
 
-def test_expand_user(patched_home_env):
-    input_path = os.path.join("~", "test")
-    expected_path = os.path.join(patched_home_env, "test")
-
-    assert expand_path(input_path) == expected_path
-
-
-def test_expand_vars(patched_home_env):
-    input_path = os.path.join("$HOME", "test")
-    expected_path = os.path.join(patched_home_env, "test")
-
-    assert expand_path(input_path) == expected_path
-
-
 @pytest.fixture
 def test_path_nested(tmpdir):
     path = os.path.join(tmpdir, "test1", "test2", "test3")