forked from p15670423/monkey
parent
3912b85d08
commit
8879dae276
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue