diff --git a/.gitattributes b/.gitattributes index 1cc8cc472ea..807ae6822d5 100644 --- a/.gitattributes +++ b/.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 diff --git a/monkey/common/utils/file_utils.py b/monkey/common/utils/file_utils.py index a4cff2b482a..fd2c85ec1d2 100644 --- a/monkey/common/utils/file_utils.py +++ b/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() diff --git a/monkey/infection_monkey/ransomware/consts.py b/monkey/infection_monkey/ransomware/consts.py new file mode 100644 index 00000000000..8ff02fe9940 --- /dev/null +++ b/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" diff --git a/monkey/infection_monkey/ransomware/file_selectors.py b/monkey/infection_monkey/ransomware/file_selectors.py index 167c547e81f..33b73dd0645 100644 --- a/monkey/infection_monkey/ransomware/file_selectors.py +++ b/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 @@ def __call__(self, target_dir: Path) -> List[Path]: 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 diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 27cd6dca10f..3c8b3677026 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -2,15 +2,13 @@ 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 @@ def run_payload(self): 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}") diff --git a/monkey/tests/conftest.py b/monkey/tests/conftest.py index 23cc840a31e..fc44af01486 100644 --- a/monkey/tests/conftest.py +++ b/monkey/tests/conftest.py @@ -11,3 +11,13 @@ @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" diff --git a/monkey/tests/data_for_tests/stable_file.txt b/monkey/tests/data_for_tests/stable_file.txt new file mode 100644 index 00000000000..ffe82625b37 --- /dev/null +++ b/monkey/tests/data_for_tests/stable_file.txt @@ -0,0 +1 @@ +Don't change me! diff --git a/monkey/tests/unit_tests/common/utils/test_common_file_utils.py b/monkey/tests/unit_tests/common/utils/test_common_file_utils.py index 3fe981a749b..79d00d027c8 100644 --- a/monkey/tests/unit_tests/common/utils/test_common_file_utils.py +++ b/monkey/tests/unit_tests/common/utils/test_common_file_utils.py @@ -2,7 +2,7 @@ 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 diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py index fd948983795..42e852b953d 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py +++ b/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.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 diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py index 3003311d0ff..eb263322655 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py @@ -9,8 +9,8 @@ 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 diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py index 118ba8a27ef..6c73cfb8d6b 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py @@ -7,12 +7,9 @@ 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( diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_dropper.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_dropper.py index 17d0d953cd4..516e0393526 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_dropper.py +++ b/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 diff --git a/monkey/tests/utils.py b/monkey/tests/utils.py index 8aea2d00767..9b57a9cc70a 100644 --- a/monkey/tests/utils.py +++ b/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