diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index f9141f77d..d6a36ab5b 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -1,26 +1,33 @@ import logging +from pathlib import Path + +from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION +from infection_monkey.utils.dir_utils import ( + file_extension_filter, + filter_files, + get_all_regular_files_in_directory, +) +from infection_monkey.utils.environment import is_windows_os LOG = logging.getLogger(__name__) class RansomewarePayload: def __init__(self, config: dict): - self.config = config + LOG.info(f"Windows dir configured for encryption is " f"{config['windows_dir']}") + LOG.info(f"Linux dir configured for encryption is " f"{config['linux_dir']}") + + self.target_dir = Path(config["windows_dir"] if is_windows_os() else config["linux_dir"]) def run_payload(self): - LOG.info( - f"Windows dir configured for encryption is " - f"{self.config['directories']['windows_dir']}" - ) - LOG.info( - f"Linux dir configured for encryption is " f"{self.config['directories']['linux_dir']}" - ) - file_list = self._find_files() self._encrypt_files(file_list) def _find_files(self): - return [] + file_filters = [file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)] + + all_files = get_all_regular_files_in_directory(self.target_dir) + return filter_files(all_files, file_filters) def _encrypt_files(self, file_list): for file in file_list: diff --git a/monkey/infection_monkey/ransomware/valid_file_extensions.py b/monkey/infection_monkey/ransomware/valid_file_extensions.py new file mode 100644 index 000000000..f67a6b761 --- /dev/null +++ b/monkey/infection_monkey/ransomware/valid_file_extensions.py @@ -0,0 +1,76 @@ +VALID_FILE_EXTENSIONS_FOR_ENCRYPTION = { + ".3ds", + ".7z", + ".accdb", + ".ai", + ".asp", + ".aspx", + ".avhd", + ".avi", + ".back", + ".bak", + ".c", + ".cfg", + ".conf", + ".cpp", + ".cs", + ".ctl", + ".dbf", + ".disk", + ".djvu", + ".doc", + ".docx", + ".dwg", + ".eml", + ".fdb", + ".giff", + ".gz", + ".h", + ".hdd", + ".jpg", + ".jpeg", + ".kdbx", + ".mail", + ".mdb", + ".mpg", + ".mpeg", + ".msg", + ".nrg", + ".ora", + ".ost", + ".ova", + ".ovf", + ".pdf", + ".php", + ".pmf", + ".png", + ".ppt", + ".pptx", + ".pst", + ".pvi", + ".py", + ".pyc", + ".rar", + ".rtf", + ".sln", + ".sql", + ".tar", + ".tiff", + ".txt", + ".vbox", + ".vbs", + ".vcb", + ".vdi", + ".vfd", + ".vmc", + ".vmdk", + ".vmsd", + ".vmx", + ".vsdx", + ".vsv", + ".work", + ".xls", + ".xlsx", + ".xvd", + ".zip", +} diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py new file mode 100644 index 000000000..3cc839d0e --- /dev/null +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -0,0 +1,21 @@ +from pathlib import Path +from typing import Callable, Iterable, List, Set + + +def get_all_regular_files_in_directory(dir_path: Path) -> List[Path]: + return filter_files(dir_path.iterdir(), [lambda f: f.is_file()]) + + +def filter_files(files: Iterable[Path], file_filters: List[Callable[[Path], bool]]): + filtered_files = files + for file_filter in file_filters: + filtered_files = [f for f in filtered_files if file_filter(f)] + + return filtered_files + + +def file_extension_filter(file_extensions: Set): + def inner_filter(f: Path): + return f.suffix in file_extensions + + return inner_filter diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py new file mode 100644 index 000000000..b8ff47a65 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -0,0 +1,93 @@ +from infection_monkey.utils.dir_utils import ( + file_extension_filter, + filter_files, + get_all_regular_files_in_directory, +) + +FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz", "2.png", "2.mpg"] +SUBDIRS = ["subdir1", "subdir2"] + + +def add_subdirs_to_dir(parent_dir): + subdirs = [parent_dir / s for s in SUBDIRS] + + for subdir in subdirs: + subdir.mkdir() + + return subdirs + + +def add_files_to_dir(parent_dir): + files = [parent_dir / f for f in FILES] + + for f in files: + f.touch() + + return files + + +def test_get_all_regular_files_in_directory__no_files(tmp_path, monkeypatch): + add_subdirs_to_dir(tmp_path) + + expected_return_value = [] + assert get_all_regular_files_in_directory(tmp_path) == expected_return_value + + +def test_get_all_regular_files_in_directory__has_files(tmp_path, monkeypatch): + add_subdirs_to_dir(tmp_path) + files = add_files_to_dir(tmp_path) + + expected_return_value = sorted(files) + assert sorted(get_all_regular_files_in_directory(tmp_path)) == expected_return_value + + +def test_get_all_regular_files_in_directory__subdir_has_files(tmp_path, monkeypatch): + subdirs = add_subdirs_to_dir(tmp_path) + add_files_to_dir(subdirs[0]) + + files = add_files_to_dir(tmp_path) + + expected_return_value = sorted(files) + assert sorted(get_all_regular_files_in_directory(tmp_path)) == expected_return_value + + +def test_filter_files__no_results(tmp_path): + add_files_to_dir(tmp_path) + + files_in_dir = get_all_regular_files_in_directory(tmp_path) + filtered_files = filter_files(files_in_dir, [lambda _: False]) + + assert len(filtered_files) == 0 + + +def test_filter_files__all_true(tmp_path): + files = add_files_to_dir(tmp_path) + expected_return_value = sorted(files) + + files_in_dir = get_all_regular_files_in_directory(tmp_path) + filtered_files = filter_files(files_in_dir, [lambda _: True]) + + assert sorted(filtered_files) == expected_return_value + + +def test_filter_files__multiple_filters(tmp_path): + files = add_files_to_dir(tmp_path) + expected_return_value = sorted(files[4:6]) + + files_in_dir = get_all_regular_files_in_directory(tmp_path) + filtered_files = filter_files( + files_in_dir, [lambda f: f.name.startswith("2"), lambda f: f.name.endswith("g")] + ) + + assert sorted(filtered_files) == expected_return_value + + +def test_file_extension_filter(tmp_path): + valid_extensions = {".zip", ".xyz"} + + files = add_files_to_dir(tmp_path) + + files_in_dir = get_all_regular_files_in_directory(tmp_path) + filtered_files = filter_files(files_in_dir, [file_extension_filter(valid_extensions)]) + + assert sorted(files[0:2]) == sorted(filtered_files) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 2c937ee4f..304ff6f12 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -171,6 +171,7 @@ ISLAND # unused variable (monkey/monkey_island/cc/services/utils/node_states.py MONKEY_LINUX_RUNNING # unused variable (monkey/monkey_island/cc/services/utils/node_states.py:26) import_status # monkey_island\cc\resources\configuration_import.py:19 config_schema # monkey_island\cc\resources\configuration_import.py:25 +get_files_to_encrypt # monkey/infection_monkey/ransomware/utils.py:82 # these are not needed for it to work, but may be useful extra information to understand what's going on WINDOWS_PBA_TYPE # unused variable (monkey/monkey_island/cc/resources/pba_file_upload.py:23)