island: Simplify remove_PBA_files()

This commit is contained in:
Mike Salvatore 2021-05-11 10:56:51 -04:00
parent 5742e85ff5
commit 4364a48561
3 changed files with 33 additions and 11 deletions

View File

@ -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):

View File

@ -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

View File

@ -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)