Merge pull request #1333 from guardicore/ransomware-skip-encrypt-readme

Ransomware skip encrypt readme

.gitattributes

@@ -1 +1,4 @@
 monkey/tests/data_for_tests/ransomware_targets/** -text
+monkey/tests/data_for_tests/test_readme.txt -text
+monkey/tests/data_for_tests/stable_file.txt -text
+monkey/infection_monkey/ransomware/ransomware_readme.txt -text

monkey/common/utils/file_utils.py

@@ -1,3 +1,4 @@
+import hashlib
 import os
 from pathlib import Path
 
@@ -11,3 +12,12 @@ def expand_path(path: str) -> Path:
         raise InvalidPath("Empty path provided")
 
     return Path(os.path.expandvars(os.path.expanduser(path)))
+
+
+def get_file_sha256_hash(filepath: Path):
+    sha256 = hashlib.sha256()
+    with open(filepath, "rb") as f:
+        for block in iter(lambda: f.read(65536), b""):
+            sha256.update(block)
+
+    return sha256.hexdigest()

monkey/infection_monkey/ransomware/consts.py

@@ -0,0 +1,5 @@
+from pathlib import Path
+
+README_SRC = Path(__file__).parent / "ransomware_readme.txt"
+README_FILE_NAME = "README.txt"
+README_SHA256_HASH = "e3d9343cbcce6097c83044327b00ead14b6e8e6aa0d411160610033a856032fc"

monkey/infection_monkey/ransomware/file_selectors.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import List, Set
 
+from common.utils.file_utils import get_file_sha256_hash
+from infection_monkey.ransomware.consts import README_FILE_NAME, README_SHA256_HASH
 from infection_monkey.utils.dir_utils import (
     file_extension_filter,
     filter_files,
@@ -19,7 +21,15 @@ class ProductionSafeTargetFileSelector:
             file_extension_filter(self._targeted_file_extensions),
             is_not_shortcut_filter,
             is_not_symlink_filter,
+            _is_not_ransomware_readme_filter,
         ]
 
         all_files = get_all_regular_files_in_directory(target_dir)
         return filter_files(all_files, file_filters)
+
+
+def _is_not_ransomware_readme_filter(filepath: Path) -> bool:
+    if filepath.name != README_FILE_NAME:
+        return True
+
+    return get_file_sha256_hash(filepath) != README_SHA256_HASH

monkey/infection_monkey/ransomware/ransomware_payload.py

@@ -2,15 +2,13 @@ import logging
 from pathlib import Path
 from typing import Callable, List
 
+from infection_monkey.ransomware.consts import README_FILE_NAME, README_SRC
 from infection_monkey.ransomware.ransomware_config import RansomwareConfig
 from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 LOG = logging.getLogger(__name__)
 
-README_SRC = Path(__file__).parent / "ransomware_readme.txt"
-README_DEST = "README.txt"
-
 
 class RansomwarePayload:
     def __init__(
@@ -39,7 +37,7 @@ class RansomwarePayload:
             self._encrypt_files(file_list)
 
         if self._config.readme_enabled:
-            self._leave_readme(README_SRC, self._config.target_directory / README_DEST)
+            self._leave_readme(README_SRC, self._config.target_directory / README_FILE_NAME)
 
     def _find_files(self) -> List[Path]:
         LOG.info(f"Collecting files in {self._config.target_directory}")

monkey/tests/conftest.py

@@ -11,3 +11,13 @@ sys.path.insert(0, MONKEY_BASE_PATH)
 @pytest.fixture(scope="session")
 def data_for_tests_dir(pytestconfig):
     return Path(os.path.join(pytestconfig.rootdir, "monkey", "tests", "data_for_tests"))
+
+
+@pytest.fixture(scope="session")
+def stable_file(data_for_tests_dir) -> Path:
+    return data_for_tests_dir / "stable_file.txt"
+
+
+@pytest.fixture(scope="session")
+def stable_file_sha256_hash() -> str:
+    return "d9dcaadc91261692dafa86e7275b1bf39bb7e19d2efcfacd6fe2bfc9a1ae1062"

monkey/tests/data_for_tests/stable_file.txt

@@ -0,0 +1 @@
+Don't change me!

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

@@ -2,7 +2,7 @@ import os
 
 import pytest
 
-from common.utils.file_utils import InvalidPath, expand_path
+from common.utils.file_utils import InvalidPath, expand_path, get_file_sha256_hash
 
 
 def test_expand_user(patched_home_env):
@@ -22,3 +22,7 @@ def test_expand_vars(patched_home_env):
 def test_expand_path__empty_path_provided():
     with pytest.raises(InvalidPath):
         expand_path("")
+
+
+def test_get_file_sha256_hash(stable_file, stable_file_sha256_hash):
+    assert get_file_sha256_hash(stable_file) == stable_file_sha256_hash

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

@@ -1,4 +1,5 @@
 import os
+import shutil
 
 import pytest
 from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
@@ -12,6 +13,7 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
 from tests.utils import is_user_admin
 
 from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector
+from infection_monkey.ransomware.ransomware_payload import README_SRC
 
 TARGETED_FILE_EXTENSIONS = [".pdf", ".txt"]
 
@@ -53,3 +55,21 @@ def test_directories_not_selected(ransomware_test_data, file_selector):
     selected_files = file_selector(ransomware_test_data)
 
     assert (ransomware_test_data / SUBDIR / HELLO_TXT) not in selected_files
+
+
+def test_ransomware_readme_not_selected(ransomware_target, file_selector):
+    readme_file = ransomware_target / "README.txt"
+    shutil.copyfile(README_SRC, readme_file)
+
+    selected_files = file_selector(ransomware_target)
+
+    assert readme_file not in selected_files
+
+
+def test_pre_existing_readme_is_selected(ransomware_target, stable_file, file_selector):
+    readme_file = ransomware_target / "README.txt"
+    shutil.copyfile(stable_file, readme_file)
+
+    selected_files = file_selector(ransomware_target)
+
+    assert readme_file in selected_files

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

@@ -9,8 +9,8 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
     TEST_KEYBOARD_TXT_CLEARTEXT_SHA256,
     TEST_KEYBOARD_TXT_ENCRYPTED_SHA256,
 )
-from tests.utils import hash_file
 
+from common.utils.file_utils import get_file_sha256_hash
 from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor
 from infection_monkey.utils.bit_manipulators import flip_bits
 
@@ -44,11 +44,11 @@ def test_file_encrypted(
 ):
     test_keyboard = ransomware_target / file_name
 
-    assert hash_file(test_keyboard) == cleartext_hash
+    assert get_file_sha256_hash(test_keyboard) == cleartext_hash
 
     in_place_bitflip_file_encryptor(test_keyboard)
 
-    assert hash_file(test_keyboard) == encrypted_hash
+    assert get_file_sha256_hash(test_keyboard) == encrypted_hash
 
 
 def test_file_encrypted_in_place(in_place_bitflip_file_encryptor, ransomware_target):
@@ -70,4 +70,4 @@ def test_encrypted_file_has_new_extension(ransomware_target):
 
     assert not test_keyboard.exists()
     assert encrypted_test_keyboard.exists()
-    assert hash_file(encrypted_test_keyboard) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256
+    assert get_file_sha256_hash(encrypted_test_keyboard) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256

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

@@ -7,12 +7,9 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
     TEST_KEYBOARD_TXT,
 )
 
+from infection_monkey.ransomware.consts import README_FILE_NAME, README_SRC
 from infection_monkey.ransomware.ransomware_config import RansomwareConfig
-from infection_monkey.ransomware.ransomware_payload import (
-    README_DEST,
-    README_SRC,
-    RansomwarePayload,
-)
+from infection_monkey.ransomware.ransomware_payload import RansomwarePayload
 
 
 @pytest.fixture
@@ -162,7 +159,7 @@ def test_readme_true(
     ransomware_payload = build_ransomware_payload(ransomware_payload_config)
 
     ransomware_payload.run_payload()
-    mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_DEST)
+    mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_FILE_NAME)
 
 
 def test_no_readme_if_no_directory(

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

@@ -1,6 +1,6 @@
 import pytest
-from tests.utils import hash_file
 
+from common.utils.file_utils import get_file_sha256_hash
 from infection_monkey.ransomware.readme_dropper import leave_readme
 
 DEST_FILE = "README.TXT"
@@ -23,10 +23,10 @@ def test_readme_already_exists(src_readme, dest_readme):
 
     leave_readme(src_readme, dest_readme)
 
-    assert hash_file(dest_readme) == EMPTY_FILE_HASH
+    assert get_file_sha256_hash(dest_readme) == EMPTY_FILE_HASH
 
 
 def test_leave_readme(src_readme, dest_readme):
     leave_readme(src_readme, dest_readme)
 
-    assert hash_file(dest_readme) == README_HASH
+    assert get_file_sha256_hash(dest_readme) == README_HASH

monkey/tests/utils.py

@@ -1,7 +1,5 @@
 import ctypes
-import hashlib
 import os
-from pathlib import Path
 
 
 def is_user_admin():
@@ -11,14 +9,5 @@ def is_user_admin():
     return ctypes.windll.shell32.IsUserAnAdmin()
 
 
-def hash_file(filepath: Path):
-    sha256 = hashlib.sha256()
-    with open(filepath, "rb") as f:
-        for block in iter(lambda: f.read(65536), b""):
-            sha256.update(block)
-
-    return sha256.hexdigest()
-
-
 def raise_(ex):
     raise ex