From 41bf137ee4817c731928bda9dc81b887d70af2d8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Jun 2021 15:24:36 -0400 Subject: [PATCH] agent: Add is_not_shortcut_filter() Adds a filter that can be used with filter_files() to return only files that are not Windows shortcuts. --- monkey/infection_monkey/utils/dir_utils.py | 4 ++++ .../infection_monkey/utils/test_dir_utils.py | 14 +++++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 31455535e..704556335 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -23,3 +23,7 @@ def file_extension_filter(file_extensions: Set): def is_not_symlink_filter(f: Path): return not f.is_symlink() + + +def is_not_shortcut_filter(f: Path): + return f.suffix != ".lnk" 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 089563cee..8ebddf280 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 @@ -7,10 +7,12 @@ from infection_monkey.utils.dir_utils import ( file_extension_filter, filter_files, get_all_regular_files_in_directory, + is_not_shortcut_filter, is_not_symlink_filter, ) -FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz", "2.png", "2.mpg"] +SHORTCUT = "shortcut.lnk" +FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz", "2.png", "2.mpg", SHORTCUT] SUBDIRS = ["subdir1", "subdir2"] @@ -113,3 +115,13 @@ def test_is_not_symlink_filter(tmp_path): assert link_path in files_in_dir assert len(filtered_files) == len(FILES) assert link_path not in filtered_files + + +def test_is_not_shortcut_filter(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, [is_not_shortcut_filter]) + + assert len(filtered_files) == len(FILES) - 1 + assert SHORTCUT not in [f.name for f in filtered_files]