From 4364a485612b803e5c20eda4b6bbff3b8577fc90 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 11 May 2021 10:56:51 -0400 Subject: [PATCH] island: Simplify remove_PBA_files() --- .../cc/services/post_breach_files.py | 13 ++------ .../cc/services/test_post_breach_files.py | 30 +++++++++++++++++++ whitelist.py | 1 + 3 files changed, 33 insertions(+), 11 deletions(-) create mode 100644 monkey/tests/monkey_island/cc/services/test_post_breach_files.py diff --git a/monkey/monkey_island/cc/services/post_breach_files.py b/monkey/monkey_island/cc/services/post_breach_files.py index d7eeb9306..7e2ce4e45 100644 --- a/monkey/monkey_island/cc/services/post_breach_files.py +++ b/monkey/monkey_island/cc/services/post_breach_files.py @@ -28,17 +28,8 @@ class PostBreachFilesService: @staticmethod def remove_PBA_files(): - if monkey_island.cc.services.config.ConfigService.get_config(): - windows_filename = monkey_island.cc.services.config.ConfigService.get_config_value( - PBA_WINDOWS_FILENAME_PATH - ) - linux_filename = monkey_island.cc.services.config.ConfigService.get_config_value( - PBA_LINUX_FILENAME_PATH - ) - if linux_filename: - PostBreachFilesService.remove_file(linux_filename) - if windows_filename: - PostBreachFilesService.remove_file(windows_filename) + for f in os.listdir(PostBreachFilesService.get_custom_pba_directory()): + PostBreachFilesService.remove_file(f) @staticmethod def remove_file(file_name): diff --git a/monkey/tests/monkey_island/cc/services/test_post_breach_files.py b/monkey/tests/monkey_island/cc/services/test_post_breach_files.py new file mode 100644 index 000000000..9eda4af55 --- /dev/null +++ b/monkey/tests/monkey_island/cc/services/test_post_breach_files.py @@ -0,0 +1,30 @@ +import os + +import pytest + +from monkey_island.cc.services.post_breach_files import PostBreachFilesService + + +@pytest.fixture(autouse=True) +def custom_pba_directory(tmpdir): + PostBreachFilesService.initialize(tmpdir) + + +def create_custom_pba_file(filename): + assert os.path.isdir(PostBreachFilesService.get_custom_pba_directory()) + + file_path = os.path.join(PostBreachFilesService.get_custom_pba_directory(), filename) + open(file_path, "a").close() + + +def test_remove_pba_files(): + create_custom_pba_file("linux_file") + create_custom_pba_file("windows_file") + + custom_pda_dir_contents = os.listdir(PostBreachFilesService.get_custom_pba_directory()) + assert len(custom_pda_dir_contents) == 2 + + PostBreachFilesService.remove_PBA_files() + + custom_pda_dir_contents = os.listdir(PostBreachFilesService.get_custom_pba_directory()) + assert len(custom_pda_dir_contents) == 0 diff --git a/whitelist.py b/whitelist.py index ad346ff0b..bd147220a 100644 --- a/whitelist.py +++ b/whitelist.py @@ -20,6 +20,7 @@ set_os_windows # unused variable (monkey/tests/infection_monkey/post_breach/act patch_new_user_classes # unused variable (monkey/tests/infection_monkey/utils/test_auto_new_user_factory.py:25) patch_new_user_classes # unused variable (monkey/tests/infection_monkey/utils/test_auto_new_user_factory.py:31) mock_home_env # unused variable (monkey/tests/monkey_island/cc/server_utils/test_island_logger.py:20) +custom_pba_directory # unused variable (monkey/tests/monkey_island/cc/services/test_post_breach_files.py:20) configure_resources # unused function (monkey/tests/monkey_island/cc/environment/test_environment.py:26) change_to_mongo_mock # unused function (monkey/monkey_island/cc/test_common/fixtures/mongomock_fixtures.py:9) uses_database # unused function (monkey/monkey_island/cc/test_common/fixtures/mongomock_fixtures.py:16)