diff --git a/.gitattributes b/.gitattributes index 1cc8cc472..807ae6822 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 a4cff2b48..fd2c85ec1 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 000000000..8ff02fe99 --- /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 167c547e8..33b73dd06 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 @@ 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 diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 27cd6dca1..3c8b36770 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/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}") diff --git a/monkey/tests/conftest.py b/monkey/tests/conftest.py index 23cc840a3..fc44af014 100644 --- a/monkey/tests/conftest.py +++ b/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" 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 000000000..ffe82625b --- /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 3fe981a74..79d00d027 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 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 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 fd9489837..42e852b95 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.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 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 3003311d0..eb2633226 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 @@ 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 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 118ba8a27..6c73cfb8d 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 @@ 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( 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 17d0d953c..516e03935 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 8aea2d007..9b57a9cc7 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