p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

agent: Refactor filter_files to accept a list of filters

This commit is contained in:
Mike Salvatore 2021-06-22 13:43:27 -04:00
parent a8ebe6ae76
commit db8dfd9f17
3 changed files with 16 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
monkey/infection_monkey/ransomware/ransomware_payload.py
View File

@ -24,8 +24,10 @@ class RansomewarePayload:
self._encrypt_files(file_list) self._encrypt_files(file_list)
def _find_files(self): 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_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): def _encrypt_files(self, file_list):
for file in file_list: for file in file_list:

8
monkey/infection_monkey/utils/dir_utils.py
View File

@ -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()] return [f for f in dir_path.iterdir() if f.is_file()]
def filter_files(files: List[Path], file_filter: Callable[[Path], bool]): def filter_files(files: List[Path], file_filters: List[Callable[[Path], bool]]):
return [f for f in files if file_filter(f)] 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 file_extension_filter(file_extensions: Set):

17
monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py
View File

@ -4,8 +4,7 @@ from infection_monkey.utils.dir_utils import (
get_all_files_in_directory, get_all_files_in_directory,
) )
FILE_1 = "file.jpg.zip" FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz"]
FILE_2 = "file.xyz"
SUBDIR_1 = "subdir1" SUBDIR_1 = "subdir1"
SUBDIR_2 = "subdir2" SUBDIR_2 = "subdir2"
@ -22,9 +21,7 @@ def add_subdirs_to_dir(parent_dir):
def add_files_to_dir(parent_dir): def add_files_to_dir(parent_dir):
file1 = parent_dir / FILE_1 files = [parent_dir / f for f in FILES]
file2 = parent_dir / FILE_2
files = [file1, file2]
for f in files: for f in files:
f.touch() f.touch()
@ -61,7 +58,7 @@ def test_filter_files__no_results(tmp_path):
add_files_to_dir(tmp_path) add_files_to_dir(tmp_path)
files_in_dir = get_all_files_in_directory(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 assert len(filtered_files) == 0
@ -71,17 +68,17 @@ def test_filter_files__all_true(tmp_path):
expected_return_value = sorted(files) expected_return_value = sorted(files)
files_in_dir = get_all_files_in_directory(tmp_path) 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 assert sorted(filtered_files) == expected_return_value
def test_file_extension_filter(tmp_path): def test_file_extension_filter(tmp_path):
valid_extensions = {".zip", ".tar"} valid_extensions = {".zip", ".xyz"}
files = add_files_to_dir(tmp_path) files = add_files_to_dir(tmp_path)
files_in_dir = get_all_files_in_directory(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)