island: Simplify remove_PBA_files()

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():
+        for f in os.listdir(PostBreachFilesService.get_custom_pba_directory()):
-            windows_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
+            PostBreachFilesService.remove_file(f)
-                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)
 
     @staticmethod
     def remove_file(file_name):

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

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)