From 5b64ea5151e715012fbd406fed41296e37fc002e Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 11:41:49 +0530 Subject: [PATCH 01/20] agent: ransomware: Iterate through files in directory and get list of files to encrypt --- monkey/infection_monkey/ransomware/utils.py | 94 +++++++++++++++++++++ vulture_allowlist.py | 1 + 2 files changed, 95 insertions(+) create mode 100644 monkey/infection_monkey/ransomware/utils.py diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py new file mode 100644 index 000000000..d9cc7d5ce --- /dev/null +++ b/monkey/infection_monkey/ransomware/utils.py @@ -0,0 +1,94 @@ +import os +from typing import Iterator, List + +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", +} + + +def get_files_to_encrypt(dir_path: str) -> List[str]: + all_files = get_all_files_in_directory(dir_path) + + files_to_encrypt = [] + for file in all_files: + if os.path.splitext(file)[1] in VALID_FILE_EXTENSIONS_FOR_ENCRYPTION: + files_to_encrypt.append(file) + + return files_to_encrypt + + +def get_all_files_in_directory(dir_path: str) -> Iterator: + return filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)]) 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) From 3192355f5a61c1e7074e05efe88a01383423ba65 Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 12:47:38 +0530 Subject: [PATCH 02/20] agent: ransomware: Return list instead of filter object when getting files in directory --- monkey/infection_monkey/ransomware/utils.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py index d9cc7d5ce..2fd0052fb 100644 --- a/monkey/infection_monkey/ransomware/utils.py +++ b/monkey/infection_monkey/ransomware/utils.py @@ -1,5 +1,5 @@ import os -from typing import Iterator, List +from typing import List VALID_FILE_EXTENSIONS_FOR_ENCRYPTION = { ".3ds", @@ -90,5 +90,7 @@ def get_files_to_encrypt(dir_path: str) -> List[str]: return files_to_encrypt -def get_all_files_in_directory(dir_path: str) -> Iterator: - return filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)]) +def get_all_files_in_directory(dir_path: str) -> List: + return list( + filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)]) + ) From f8c1886977cbc5adbcc68ec30f3a92b3741e95f4 Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 12:50:28 +0530 Subject: [PATCH 03/20] tests: Add unit tests for ransomware utils - get_files_to_encrypt(), get_all_files_in_directory() --- .../infection_monkey/ransomware/test_utils.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py new file mode 100644 index 000000000..579b7e2c9 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py @@ -0,0 +1,91 @@ +import os + +import infection_monkey.ransomware.utils + +VALID_FILE_EXTENSION_1 = "file.3ds" +VALID_FILE_EXTENSION_2 = "file.zip" +INVALID_FILE_EXTENSION_1 = "file.jpe" +INVALID_FILE_EXTENSION_2 = "file.xyz" +SUBDIR_1 = "subdir1" +SUBDIR_2 = "subdir2" + + +def test_get_files_to_encrypt__no_files(monkeypatch): + all_files = [] + monkeypatch.setattr( + "infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files + ) + + expected_return_value = [] + assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value + + +def test_get_files_to_encrypt__no_valid_files(monkeypatch): + all_files = [INVALID_FILE_EXTENSION_1, INVALID_FILE_EXTENSION_2] + monkeypatch.setattr( + "infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files + ) + + expected_return_value = [] + assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value + + +def test_get_files_to_encrypt__valid_files(monkeypatch): + all_files = [ + VALID_FILE_EXTENSION_1, + INVALID_FILE_EXTENSION_1, + VALID_FILE_EXTENSION_2, + INVALID_FILE_EXTENSION_2, + ] + monkeypatch.setattr( + "infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files + ) + + expected_return_value = [VALID_FILE_EXTENSION_1, VALID_FILE_EXTENSION_2] + assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value + + +def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch): + subdir1 = os.path.join(tmpdir, SUBDIR_1) + subdir2 = os.path.join(tmpdir, SUBDIR_2) + subdirs = [subdir1, subdir2] + + for subdir in subdirs: + os.mkdir(subdir) + + all_items_in_dir = subdirs + monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) + + expected_return_value = [] + assert ( + infection_monkey.ransomware.utils.get_all_files_in_directory(tmpdir) + == expected_return_value + ) + + +def test_get_all_files_in_directory__has_files(tmpdir, monkeypatch): + subdir1 = os.path.join(tmpdir, SUBDIR_1) + subdir2 = os.path.join(tmpdir, SUBDIR_2) + subdirs = [subdir1, subdir2] + + file1 = os.path.join(tmpdir, VALID_FILE_EXTENSION_1) + file2 = os.path.join(tmpdir, INVALID_FILE_EXTENSION_1) + file3 = os.path.join(tmpdir, VALID_FILE_EXTENSION_2) + file4 = os.path.join(tmpdir, INVALID_FILE_EXTENSION_2) + files = [file1, file2, file3, file4] + + for subdir in subdirs: + os.mkdir(subdir) + + for file in files: + with open(file, "w") as _: + pass + + all_items_in_dir = subdirs + files + monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) + + expected_return_value = files + assert ( + infection_monkey.ransomware.utils.get_all_files_in_directory(tmpdir) + == expected_return_value + ) From 8abb9d95cea438b202ba11766a4ecffe6c7960ba Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 18:06:14 +0530 Subject: [PATCH 04/20] agent: ransomware: Move VALID_FILE_EXTENSIONS_FOR_ENCRYPTION to separate file --- monkey/infection_monkey/ransomware/utils.py | 77 +------------------ .../ransomware/valid_file_extensions.py | 76 ++++++++++++++++++ 2 files changed, 77 insertions(+), 76 deletions(-) create mode 100644 monkey/infection_monkey/ransomware/valid_file_extensions.py diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py index 2fd0052fb..8d4cbcc8a 100644 --- a/monkey/infection_monkey/ransomware/utils.py +++ b/monkey/infection_monkey/ransomware/utils.py @@ -1,82 +1,7 @@ import os from typing import List -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", -} +from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION def get_files_to_encrypt(dir_path: str) -> List[str]: 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", +} From 48967a2e5b4ce1188845f86d7c2ac58a6d4c878a Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 19:03:54 +0530 Subject: [PATCH 05/20] tests: Modify unit tests for ransomware utils to test for a file like file.jpg.zip --- .../unit_tests/infection_monkey/ransomware/test_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py index 579b7e2c9..f3c0bea39 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py @@ -3,8 +3,8 @@ import os import infection_monkey.ransomware.utils VALID_FILE_EXTENSION_1 = "file.3ds" -VALID_FILE_EXTENSION_2 = "file.zip" -INVALID_FILE_EXTENSION_1 = "file.jpe" +VALID_FILE_EXTENSION_2 = "file.jpg.zip" +INVALID_FILE_EXTENSION_1 = "file.pqr" INVALID_FILE_EXTENSION_2 = "file.xyz" SUBDIR_1 = "subdir1" SUBDIR_2 = "subdir2" From c6da73ad798923998c5b25426a4595aac0c3dff0 Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 19:42:50 +0530 Subject: [PATCH 06/20] agent: ransomware: Incorporate changes into stub --- .../infection_monkey/ransomware/ransomware_payload.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 941055062..779181992 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -1,5 +1,8 @@ import logging +from infection_monkey.ransomware.utils import get_files_to_encrypt +from infection_monkey.utils.environment import is_windows_os + LOG = logging.getLogger(__name__) @@ -17,7 +20,12 @@ class RansomewarePayload: self._encrypt_files(file_list) def _find_files(self): - return [] + dir_path = ( + self.config["windows_dir_ransom"] + if is_windows_os() + else self.config["linux_dir_ransom"] + ) + return get_files_to_encrypt(dir_path) def _encrypt_files(self, file_list): for file in file_list: From 5f714e0d9846f84474d1ba2011a0e8da4f9672c5 Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 19:45:50 +0530 Subject: [PATCH 07/20] agent: Move get_all_files_in_directory() from ransomware/utils.py to utils/file_utils.py --- monkey/infection_monkey/ransomware/utils.py | 7 +------ monkey/infection_monkey/utils/file_utils.py | 8 ++++++++ 2 files changed, 9 insertions(+), 6 deletions(-) create mode 100644 monkey/infection_monkey/utils/file_utils.py diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py index 8d4cbcc8a..50e49683c 100644 --- a/monkey/infection_monkey/ransomware/utils.py +++ b/monkey/infection_monkey/ransomware/utils.py @@ -2,6 +2,7 @@ import os from typing import List from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION +from infection_monkey.utils.file_utils import get_all_files_in_directory def get_files_to_encrypt(dir_path: str) -> List[str]: @@ -13,9 +14,3 @@ def get_files_to_encrypt(dir_path: str) -> List[str]: files_to_encrypt.append(file) return files_to_encrypt - - -def get_all_files_in_directory(dir_path: str) -> List: - return list( - filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)]) - ) diff --git a/monkey/infection_monkey/utils/file_utils.py b/monkey/infection_monkey/utils/file_utils.py new file mode 100644 index 000000000..5562dd6ec --- /dev/null +++ b/monkey/infection_monkey/utils/file_utils.py @@ -0,0 +1,8 @@ +import os +from typing import List + + +def get_all_files_in_directory(dir_path: str) -> List: + return list( + filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)]) + ) From a2ebe3386f32314c6ce13a5803b0b8cbbc8f6d6d Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 19:57:56 +0530 Subject: [PATCH 08/20] agent: Rename utils/file_utils.py to utils/dir_utils.py --- monkey/infection_monkey/ransomware/utils.py | 2 +- monkey/infection_monkey/utils/{file_utils.py => dir_utils.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename monkey/infection_monkey/utils/{file_utils.py => dir_utils.py} (100%) diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py index 50e49683c..8ad1279c7 100644 --- a/monkey/infection_monkey/ransomware/utils.py +++ b/monkey/infection_monkey/ransomware/utils.py @@ -2,7 +2,7 @@ import os from typing import List from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION -from infection_monkey.utils.file_utils import get_all_files_in_directory +from infection_monkey.utils.dir_utils import get_all_files_in_directory def get_files_to_encrypt(dir_path: str) -> List[str]: diff --git a/monkey/infection_monkey/utils/file_utils.py b/monkey/infection_monkey/utils/dir_utils.py similarity index 100% rename from monkey/infection_monkey/utils/file_utils.py rename to monkey/infection_monkey/utils/dir_utils.py From efef40edf911573a3958e54a9620a580619293fe Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 22 Jun 2021 20:07:38 +0530 Subject: [PATCH 09/20] tests: Fix ransomware tests and move tests for get_all_files_in_directory from ransomware/test_utils.py to utils/test_dir_utils.py --- .../infection_monkey/ransomware/test_utils.py | 48 ------------------- .../infection_monkey/utils/test_dir_utils.py | 46 ++++++++++++++++++ 2 files changed, 46 insertions(+), 48 deletions(-) create mode 100644 monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py index f3c0bea39..37a7f059d 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py @@ -1,5 +1,3 @@ -import os - import infection_monkey.ransomware.utils VALID_FILE_EXTENSION_1 = "file.3ds" @@ -43,49 +41,3 @@ def test_get_files_to_encrypt__valid_files(monkeypatch): expected_return_value = [VALID_FILE_EXTENSION_1, VALID_FILE_EXTENSION_2] assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value - - -def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch): - subdir1 = os.path.join(tmpdir, SUBDIR_1) - subdir2 = os.path.join(tmpdir, SUBDIR_2) - subdirs = [subdir1, subdir2] - - for subdir in subdirs: - os.mkdir(subdir) - - all_items_in_dir = subdirs - monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) - - expected_return_value = [] - assert ( - infection_monkey.ransomware.utils.get_all_files_in_directory(tmpdir) - == expected_return_value - ) - - -def test_get_all_files_in_directory__has_files(tmpdir, monkeypatch): - subdir1 = os.path.join(tmpdir, SUBDIR_1) - subdir2 = os.path.join(tmpdir, SUBDIR_2) - subdirs = [subdir1, subdir2] - - file1 = os.path.join(tmpdir, VALID_FILE_EXTENSION_1) - file2 = os.path.join(tmpdir, INVALID_FILE_EXTENSION_1) - file3 = os.path.join(tmpdir, VALID_FILE_EXTENSION_2) - file4 = os.path.join(tmpdir, INVALID_FILE_EXTENSION_2) - files = [file1, file2, file3, file4] - - for subdir in subdirs: - os.mkdir(subdir) - - for file in files: - with open(file, "w") as _: - pass - - all_items_in_dir = subdirs + files - monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) - - expected_return_value = files - assert ( - infection_monkey.ransomware.utils.get_all_files_in_directory(tmpdir) - == expected_return_value - ) 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..22b2ba0c9 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -0,0 +1,46 @@ +import os + +from infection_monkey.utils.dir_utils import get_all_files_in_directory + +FILE_1 = "file.jpg.zip" +FILE_2 = "file.xyz" +SUBDIR_1 = "subdir1" +SUBDIR_2 = "subdir2" + + +def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch): + subdir1 = os.path.join(tmpdir, SUBDIR_1) + subdir2 = os.path.join(tmpdir, SUBDIR_2) + subdirs = [subdir1, subdir2] + + for subdir in subdirs: + os.mkdir(subdir) + + all_items_in_dir = subdirs + monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) + + expected_return_value = [] + assert get_all_files_in_directory(tmpdir) == expected_return_value + + +def test_get_all_files_in_directory__has_files(tmpdir, monkeypatch): + subdir1 = os.path.join(tmpdir, SUBDIR_1) + subdir2 = os.path.join(tmpdir, SUBDIR_2) + subdirs = [subdir1, subdir2] + + file1 = os.path.join(tmpdir, FILE_1) + file2 = os.path.join(tmpdir, FILE_2) + files = [file1, file2] + + for subdir in subdirs: + os.mkdir(subdir) + + for file in files: + with open(file, "w") as _: + pass + + all_items_in_dir = subdirs + files + monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) + + expected_return_value = files + assert get_all_files_in_directory(tmpdir) == expected_return_value From e2dfd6a5e3c7d0028985a990982fc040371c6b17 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 12:15:03 -0400 Subject: [PATCH 10/20] agent: Simplify get_all_files_in_directory() with list comprehension --- monkey/infection_monkey/utils/dir_utils.py | 8 +-- .../infection_monkey/utils/test_dir_utils.py | 50 +++++++++++-------- 2 files changed, 34 insertions(+), 24 deletions(-) diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 5562dd6ec..62f592c5e 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -1,8 +1,8 @@ -import os +from pathlib import Path from typing import List def get_all_files_in_directory(dir_path: str) -> List: - return list( - filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)]) - ) + path = Path(dir_path) + + return [str(f) for f in path.iterdir() if f.is_file()] 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 index 22b2ba0c9..7156662ba 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -1,4 +1,5 @@ import os +from pathlib import Path from infection_monkey.utils.dir_utils import get_all_files_in_directory @@ -8,39 +9,48 @@ SUBDIR_1 = "subdir1" SUBDIR_2 = "subdir2" -def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch): - subdir1 = os.path.join(tmpdir, SUBDIR_1) - subdir2 = os.path.join(tmpdir, SUBDIR_2) +def add_subdirs_to_dir(parent_dir): + subdir1 = os.path.join(parent_dir, SUBDIR_1) + subdir2 = os.path.join(parent_dir, SUBDIR_2) subdirs = [subdir1, subdir2] for subdir in subdirs: os.mkdir(subdir) - all_items_in_dir = subdirs - monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) + return subdirs + + +def add_files_to_dir(parent_dir): + file1 = os.path.join(parent_dir, FILE_1) + file2 = os.path.join(parent_dir, FILE_2) + files = [file1, file2] + + for f in files: + Path(f).touch() + + return files + + +def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch): + add_subdirs_to_dir(tmpdir) expected_return_value = [] assert get_all_files_in_directory(tmpdir) == expected_return_value def test_get_all_files_in_directory__has_files(tmpdir, monkeypatch): - subdir1 = os.path.join(tmpdir, SUBDIR_1) - subdir2 = os.path.join(tmpdir, SUBDIR_2) - subdirs = [subdir1, subdir2] + add_subdirs_to_dir(tmpdir) + files = add_files_to_dir(tmpdir) - file1 = os.path.join(tmpdir, FILE_1) - file2 = os.path.join(tmpdir, FILE_2) - files = [file1, file2] + expected_return_value = sorted(files) + assert sorted(get_all_files_in_directory(tmpdir)) == expected_return_value - for subdir in subdirs: - os.mkdir(subdir) - for file in files: - with open(file, "w") as _: - pass +def test_get_all_files_in_directory__subdir_has_files(tmpdir, monkeypatch): + subdirs = add_subdirs_to_dir(tmpdir) + add_files_to_dir(subdirs[0]) - all_items_in_dir = subdirs + files - monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) + files = add_files_to_dir(tmpdir) - expected_return_value = files - assert get_all_files_in_directory(tmpdir) == expected_return_value + expected_return_value = sorted(files) + assert sorted(get_all_files_in_directory(tmpdir)) == expected_return_value From bfc86041abff595090f260a70d64bf046d1c5784 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 12:47:51 -0400 Subject: [PATCH 11/20] agent: Switch get_all_files_in_directory from str to Path Path and PurePath objects provide a lot of syntactic sugar to file handling that makes the code clearer and more concise. --- monkey/infection_monkey/utils/dir_utils.py | 6 +-- .../infection_monkey/utils/test_dir_utils.py | 37 +++++++++---------- 2 files changed, 19 insertions(+), 24 deletions(-) diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 62f592c5e..4350a279d 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -2,7 +2,5 @@ from pathlib import Path from typing import List -def get_all_files_in_directory(dir_path: str) -> List: - path = Path(dir_path) - - return [str(f) for f in path.iterdir() if f.is_file()] +def get_all_files_in_directory(dir_path: Path) -> List[Path]: + return [f for f in dir_path.iterdir() if f.is_file()] 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 index 7156662ba..7369c5f41 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -1,6 +1,3 @@ -import os -from pathlib import Path - from infection_monkey.utils.dir_utils import get_all_files_in_directory FILE_1 = "file.jpg.zip" @@ -10,47 +7,47 @@ SUBDIR_2 = "subdir2" def add_subdirs_to_dir(parent_dir): - subdir1 = os.path.join(parent_dir, SUBDIR_1) - subdir2 = os.path.join(parent_dir, SUBDIR_2) + subdir1 = parent_dir / SUBDIR_1 + subdir2 = parent_dir / SUBDIR_2 subdirs = [subdir1, subdir2] for subdir in subdirs: - os.mkdir(subdir) + subdir.mkdir() return subdirs def add_files_to_dir(parent_dir): - file1 = os.path.join(parent_dir, FILE_1) - file2 = os.path.join(parent_dir, FILE_2) + file1 = parent_dir / FILE_1 + file2 = parent_dir / FILE_2 files = [file1, file2] for f in files: - Path(f).touch() + f.touch() return files -def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch): - add_subdirs_to_dir(tmpdir) +def test_get_all_files_in_directory__no_files(tmp_path, monkeypatch): + add_subdirs_to_dir(tmp_path) expected_return_value = [] - assert get_all_files_in_directory(tmpdir) == expected_return_value + assert get_all_files_in_directory(tmp_path) == expected_return_value -def test_get_all_files_in_directory__has_files(tmpdir, monkeypatch): - add_subdirs_to_dir(tmpdir) - files = add_files_to_dir(tmpdir) +def test_get_all_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_files_in_directory(tmpdir)) == expected_return_value + assert sorted(get_all_files_in_directory(tmp_path)) == expected_return_value -def test_get_all_files_in_directory__subdir_has_files(tmpdir, monkeypatch): - subdirs = add_subdirs_to_dir(tmpdir) +def test_get_all_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(tmpdir) + files = add_files_to_dir(tmp_path) expected_return_value = sorted(files) - assert sorted(get_all_files_in_directory(tmpdir)) == expected_return_value + assert sorted(get_all_files_in_directory(tmp_path)) == expected_return_value From cf2cdc4ab8b6a468886f56bd24ff205bd362a680 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 13:16:53 -0400 Subject: [PATCH 12/20] agent: Add filter_files() function to dir_utils --- monkey/infection_monkey/utils/dir_utils.py | 6 ++++- .../infection_monkey/utils/test_dir_utils.py | 24 ++++++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 4350a279d..b64e4e13f 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -1,6 +1,10 @@ from pathlib import Path -from typing import List +from typing import Callable, List def get_all_files_in_directory(dir_path: Path) -> List[Path]: return [f for f in dir_path.iterdir() if f.is_file()] + + +def filter_files(files: List[Path], file_filter: Callable[[Path], bool]): + return [f for f in files if file_filter(f)] 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 index 7369c5f41..604efc9cb 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -1,4 +1,7 @@ -from infection_monkey.utils.dir_utils import get_all_files_in_directory +from infection_monkey.utils.dir_utils import ( + filter_files, + get_all_files_in_directory, +) FILE_1 = "file.jpg.zip" FILE_2 = "file.xyz" @@ -51,3 +54,22 @@ def test_get_all_files_in_directory__subdir_has_files(tmp_path, monkeypatch): expected_return_value = sorted(files) assert sorted(get_all_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_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_files_in_directory(tmp_path) + filtered_files = filter_files(files_in_dir, lambda _: True) + + assert sorted(filtered_files) == expected_return_value From 5c1902ca73433adbfdc1e7b6c75583f1330442a8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 13:17:11 -0400 Subject: [PATCH 13/20] agent: Add file_extension_filter to dir_utils --- monkey/infection_monkey/ransomware/utils.py | 20 +++++++++---------- monkey/infection_monkey/utils/dir_utils.py | 9 ++++++++- .../infection_monkey/utils/test_dir_utils.py | 12 +++++++++++ 3 files changed, 29 insertions(+), 12 deletions(-) diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py index 8ad1279c7..dde0ff973 100644 --- a/monkey/infection_monkey/ransomware/utils.py +++ b/monkey/infection_monkey/ransomware/utils.py @@ -1,16 +1,14 @@ -import os +from pathlib import Path from typing import List from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION -from infection_monkey.utils.dir_utils import get_all_files_in_directory +from infection_monkey.utils.dir_utils import ( + file_extension_filter, + filter_files, + get_all_files_in_directory, +) -def get_files_to_encrypt(dir_path: str) -> List[str]: - all_files = get_all_files_in_directory(dir_path) - - files_to_encrypt = [] - for file in all_files: - if os.path.splitext(file)[1] in VALID_FILE_EXTENSIONS_FOR_ENCRYPTION: - files_to_encrypt.append(file) - - return files_to_encrypt +def get_files_to_encrypt(dir_path: str) -> List[Path]: + all_files = get_all_files_in_directory(Path(dir_path)) + return filter_files(all_files, file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)) diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index b64e4e13f..23b4408a9 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Callable, List +from typing import Callable, List, Set def get_all_files_in_directory(dir_path: Path) -> List[Path]: @@ -8,3 +8,10 @@ def get_all_files_in_directory(dir_path: Path) -> List[Path]: def filter_files(files: List[Path], file_filter: Callable[[Path], bool]): return [f for f in files if file_filter(f)] + + +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 index 604efc9cb..ac8312bd8 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -1,4 +1,5 @@ from infection_monkey.utils.dir_utils import ( + file_extension_filter, filter_files, get_all_files_in_directory, ) @@ -73,3 +74,14 @@ def test_filter_files__all_true(tmp_path): filtered_files = filter_files(files_in_dir, lambda _: True) assert sorted(filtered_files) == expected_return_value + + +def test_file_extension_filter(tmp_path): + valid_extensions = {".zip", ".tar"} + + files = add_files_to_dir(tmp_path) + + files_in_dir = get_all_files_in_directory(tmp_path) + filtered_files = filter_files(files_in_dir, file_extension_filter(valid_extensions)) + + assert files[0:1] == filtered_files From 0b953c8cff6455a4d137795f59d7068409be311f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 13:23:14 -0400 Subject: [PATCH 14/20] agent: Remove ransomware/utils.py The code for getting files to encrypt has become so trivial that it no longer warrants a separate function outside of _find_files(). --- .../ransomware/ransomware_payload.py | 12 +++++- monkey/infection_monkey/ransomware/utils.py | 14 ------ .../infection_monkey/ransomware/test_utils.py | 43 ------------------- 3 files changed, 10 insertions(+), 59 deletions(-) delete mode 100644 monkey/infection_monkey/ransomware/utils.py delete mode 100644 monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 779181992..a47e3ba18 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -1,6 +1,12 @@ import logging +from pathlib import Path -from infection_monkey.ransomware.utils import get_files_to_encrypt +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_files_in_directory, +) from infection_monkey.utils.environment import is_windows_os LOG = logging.getLogger(__name__) @@ -25,7 +31,9 @@ class RansomewarePayload: if is_windows_os() else self.config["linux_dir_ransom"] ) - return get_files_to_encrypt(dir_path) + + all_files = get_all_files_in_directory(Path(dir_path)) + return filter_files(all_files, file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)) def _encrypt_files(self, file_list): for file in file_list: diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py deleted file mode 100644 index dde0ff973..000000000 --- a/monkey/infection_monkey/ransomware/utils.py +++ /dev/null @@ -1,14 +0,0 @@ -from pathlib import Path -from typing import List - -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_files_in_directory, -) - - -def get_files_to_encrypt(dir_path: str) -> List[Path]: - all_files = get_all_files_in_directory(Path(dir_path)) - return filter_files(all_files, file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)) diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py deleted file mode 100644 index 37a7f059d..000000000 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py +++ /dev/null @@ -1,43 +0,0 @@ -import infection_monkey.ransomware.utils - -VALID_FILE_EXTENSION_1 = "file.3ds" -VALID_FILE_EXTENSION_2 = "file.jpg.zip" -INVALID_FILE_EXTENSION_1 = "file.pqr" -INVALID_FILE_EXTENSION_2 = "file.xyz" -SUBDIR_1 = "subdir1" -SUBDIR_2 = "subdir2" - - -def test_get_files_to_encrypt__no_files(monkeypatch): - all_files = [] - monkeypatch.setattr( - "infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files - ) - - expected_return_value = [] - assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value - - -def test_get_files_to_encrypt__no_valid_files(monkeypatch): - all_files = [INVALID_FILE_EXTENSION_1, INVALID_FILE_EXTENSION_2] - monkeypatch.setattr( - "infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files - ) - - expected_return_value = [] - assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value - - -def test_get_files_to_encrypt__valid_files(monkeypatch): - all_files = [ - VALID_FILE_EXTENSION_1, - INVALID_FILE_EXTENSION_1, - VALID_FILE_EXTENSION_2, - INVALID_FILE_EXTENSION_2, - ] - monkeypatch.setattr( - "infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files - ) - - expected_return_value = [VALID_FILE_EXTENSION_1, VALID_FILE_EXTENSION_2] - assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value From a8ebe6ae769d877b238340de0770f4b14e541dfb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 13:31:20 -0400 Subject: [PATCH 15/20] agent: Replace self.config with self.target_dir in RansomewarePayload --- .../ransomware/ransomware_payload.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index a47e3ba18..c09c35ab9 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -14,25 +14,17 @@ 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['windows_dir_ransom']}" - ) - LOG.info(f"Linux dir configured for encryption is " f"{self.config['linux_dir_ransom']}") - file_list = self._find_files() self._encrypt_files(file_list) def _find_files(self): - dir_path = ( - self.config["windows_dir_ransom"] - if is_windows_os() - else self.config["linux_dir_ransom"] - ) - - all_files = get_all_files_in_directory(Path(dir_path)) + all_files = get_all_files_in_directory(self.target_dir) return filter_files(all_files, file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)) def _encrypt_files(self, file_list): From db8dfd9f179cd0fcdd7d4aedfab698e79a53b5eb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 13:43:27 -0400 Subject: [PATCH 16/20] agent: Refactor filter_files to accept a list of filters --- .../ransomware/ransomware_payload.py | 4 +++- monkey/infection_monkey/utils/dir_utils.py | 8 ++++++-- .../infection_monkey/utils/test_dir_utils.py | 17 +++++++---------- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index c09c35ab9..e89dc2625 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -24,8 +24,10 @@ class RansomewarePayload: self._encrypt_files(file_list) def _find_files(self): + file_filters = [file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)] + all_files = get_all_files_in_directory(self.target_dir) - return filter_files(all_files, file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)) + return filter_files(all_files, file_filters) def _encrypt_files(self, file_list): for file in file_list: diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 23b4408a9..2a7797a7b 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -6,8 +6,12 @@ def get_all_files_in_directory(dir_path: Path) -> List[Path]: return [f for f in dir_path.iterdir() if f.is_file()] -def filter_files(files: List[Path], file_filter: Callable[[Path], bool]): - return [f for f in files if file_filter(f)] +def filter_files(files: List[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): 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 index ac8312bd8..4657155c0 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -4,8 +4,7 @@ from infection_monkey.utils.dir_utils import ( get_all_files_in_directory, ) -FILE_1 = "file.jpg.zip" -FILE_2 = "file.xyz" +FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz"] SUBDIR_1 = "subdir1" SUBDIR_2 = "subdir2" @@ -22,9 +21,7 @@ def add_subdirs_to_dir(parent_dir): def add_files_to_dir(parent_dir): - file1 = parent_dir / FILE_1 - file2 = parent_dir / FILE_2 - files = [file1, file2] + files = [parent_dir / f for f in FILES] for f in files: f.touch() @@ -61,7 +58,7 @@ def test_filter_files__no_results(tmp_path): add_files_to_dir(tmp_path) files_in_dir = get_all_files_in_directory(tmp_path) - filtered_files = filter_files(files_in_dir, lambda _: False) + filtered_files = filter_files(files_in_dir, [lambda _: False]) assert len(filtered_files) == 0 @@ -71,17 +68,17 @@ def test_filter_files__all_true(tmp_path): expected_return_value = sorted(files) files_in_dir = get_all_files_in_directory(tmp_path) - filtered_files = filter_files(files_in_dir, lambda _: True) + filtered_files = filter_files(files_in_dir, [lambda _: True]) assert sorted(filtered_files) == expected_return_value def test_file_extension_filter(tmp_path): - valid_extensions = {".zip", ".tar"} + valid_extensions = {".zip", ".xyz"} files = add_files_to_dir(tmp_path) files_in_dir = get_all_files_in_directory(tmp_path) - filtered_files = filter_files(files_in_dir, file_extension_filter(valid_extensions)) + filtered_files = filter_files(files_in_dir, [file_extension_filter(valid_extensions)]) - assert files[0:1] == filtered_files + assert sorted(files[0:2]) == sorted(filtered_files) From f33772060fbb5f004158b12576c27eb12314bc8b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 13:46:03 -0400 Subject: [PATCH 17/20] agent: Rename get_all_files_in_directory() Rename get_all_files_in_directory() -> get_all_regular_files_in_directory(), as this name is more explicit about exactly which files will be included in the function's output. --- .../ransomware/ransomware_payload.py | 4 ++-- monkey/infection_monkey/utils/dir_utils.py | 2 +- .../infection_monkey/utils/test_dir_utils.py | 20 +++++++++---------- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index e89dc2625..d6a36ab5b 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -5,7 +5,7 @@ from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSI from infection_monkey.utils.dir_utils import ( file_extension_filter, filter_files, - get_all_files_in_directory, + get_all_regular_files_in_directory, ) from infection_monkey.utils.environment import is_windows_os @@ -26,7 +26,7 @@ class RansomewarePayload: def _find_files(self): file_filters = [file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION)] - all_files = get_all_files_in_directory(self.target_dir) + all_files = get_all_regular_files_in_directory(self.target_dir) return filter_files(all_files, file_filters) def _encrypt_files(self, file_list): diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 2a7797a7b..8b026d36d 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Callable, List, Set -def get_all_files_in_directory(dir_path: Path) -> List[Path]: +def get_all_regular_files_in_directory(dir_path: Path) -> List[Path]: return [f for f in dir_path.iterdir() if f.is_file()] 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 index 4657155c0..d2f30a311 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -1,7 +1,7 @@ from infection_monkey.utils.dir_utils import ( file_extension_filter, filter_files, - get_all_files_in_directory, + get_all_regular_files_in_directory, ) FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz"] @@ -29,35 +29,35 @@ def add_files_to_dir(parent_dir): return files -def test_get_all_files_in_directory__no_files(tmp_path, monkeypatch): +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_files_in_directory(tmp_path) == expected_return_value + assert get_all_regular_files_in_directory(tmp_path) == expected_return_value -def test_get_all_files_in_directory__has_files(tmp_path, monkeypatch): +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_files_in_directory(tmp_path)) == expected_return_value + assert sorted(get_all_regular_files_in_directory(tmp_path)) == expected_return_value -def test_get_all_files_in_directory__subdir_has_files(tmp_path, monkeypatch): +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_files_in_directory(tmp_path)) == expected_return_value + 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_files_in_directory(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 @@ -67,7 +67,7 @@ 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_files_in_directory(tmp_path) + 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 @@ -78,7 +78,7 @@ def test_file_extension_filter(tmp_path): files = add_files_to_dir(tmp_path) - files_in_dir = get_all_files_in_directory(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) From b643cd1edd1c168d92aba0016b7a46498f820ac0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 13:48:59 -0400 Subject: [PATCH 18/20] agent: Reimplement get_all_regular_files_in_directory using filter_files --- monkey/infection_monkey/utils/dir_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 8b026d36d..3cc839d0e 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -1,12 +1,12 @@ from pathlib import Path -from typing import Callable, List, Set +from typing import Callable, Iterable, List, Set def get_all_regular_files_in_directory(dir_path: Path) -> List[Path]: - return [f for f in dir_path.iterdir() if f.is_file()] + return filter_files(dir_path.iterdir(), [lambda f: f.is_file()]) -def filter_files(files: List[Path], file_filters: List[Callable[[Path], bool]]): +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)] From 30f88ca319124baa10528768463752e3dd15b720 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 14:07:31 -0400 Subject: [PATCH 19/20] agent: Refactor add_subdirs_to_dir in test_dir_utils.py --- .../unit_tests/infection_monkey/utils/test_dir_utils.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) 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 index d2f30a311..1b193fc3e 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -5,14 +5,11 @@ from infection_monkey.utils.dir_utils import ( ) FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz"] -SUBDIR_1 = "subdir1" -SUBDIR_2 = "subdir2" +SUBDIRS = ["subdir1", "subdir2"] def add_subdirs_to_dir(parent_dir): - subdir1 = parent_dir / SUBDIR_1 - subdir2 = parent_dir / SUBDIR_2 - subdirs = [subdir1, subdir2] + subdirs = [parent_dir / s for s in SUBDIRS] for subdir in subdirs: subdir.mkdir() From bfa640444e09f53cfc0d1042f0e562cab8064b0f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 15:11:32 -0400 Subject: [PATCH 20/20] agent: Test filter_files() with multiple filters --- .../infection_monkey/utils/test_dir_utils.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) 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 index 1b193fc3e..b8ff47a65 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -4,7 +4,7 @@ from infection_monkey.utils.dir_utils import ( get_all_regular_files_in_directory, ) -FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz"] +FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz", "2.png", "2.mpg"] SUBDIRS = ["subdir1", "subdir2"] @@ -70,6 +70,18 @@ def test_filter_files__all_true(tmp_path): 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"}