diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 2d13fae99..3f57b4fc0 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,4 +1,5 @@ -from .authentication.authentication_service import AuthenticationService -from .authentication.json_file_user_datastore import JsonFileUserDatastore from .i_file_storage_service import IFileStorageService from .directory_file_storage_service import DirectoryFileStorageService + +from .authentication.authentication_service import AuthenticationService +from .authentication.json_file_user_datastore import JsonFileUserDatastore diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 64797eff0..f3ff10395 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -1,3 +1,4 @@ +from monkey_island.cc.services import DirectoryFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService @@ -5,6 +6,8 @@ from . import AuthenticationService, JsonFileUserDatastore def initialize_services(data_dir): - PostBreachFilesService.initialize(data_dir) + # This is temporary until we get DI all worked out. + PostBreachFilesService.initialize(DirectoryFileStorageService(data_dir / "custom_pbas")) + LocalMonkeyRunService.initialize(data_dir) AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir)) diff --git a/monkey/monkey_island/cc/services/post_breach_files.py b/monkey/monkey_island/cc/services/post_breach_files.py index 8268265a9..9b24ce391 100644 --- a/monkey/monkey_island/cc/services/post_breach_files.py +++ b/monkey/monkey_island/cc/services/post_breach_files.py @@ -1,46 +1,24 @@ import logging -import os -from monkey_island.cc.server_utils.file_utils import create_secure_directory +from monkey_island.cc.services import IFileStorageService logger = logging.getLogger(__name__) +# TODO: This service wraps an IFileStorageService for the sole purpose of making the +# `remove_PBA_files()` method available to the ConfigService. This whole service can be +# removed once ConfigService is refactored to be stateful (it already is but everything is +# still statically/globally scoped) and use dependency injection. class PostBreachFilesService: - DATA_DIR = None - CUSTOM_PBA_DIRNAME = "custom_pbas" + _file_storage_service = None # TODO: A number of these services should be instance objects instead of # static/singleton hybrids. At the moment, this requires invasive refactoring that's # not a priority. @classmethod - def initialize(cls, data_dir): - cls.DATA_DIR = data_dir - custom_pba_dir = cls.get_custom_pba_directory() - create_secure_directory(custom_pba_dir) + def initialize(cls, file_storage_service: IFileStorageService): + cls._file_storage_service = file_storage_service - @staticmethod - def save_file(filename: str, file_contents: bytes): - file_path = os.path.join(PostBreachFilesService.get_custom_pba_directory(), filename) - with open(file_path, "wb") as f: - f.write(file_contents) - - @staticmethod - def remove_PBA_files(): - for f in os.listdir(PostBreachFilesService.get_custom_pba_directory()): - PostBreachFilesService.remove_file(f) - - @staticmethod - def remove_file(file_name): - file_path = os.path.join(PostBreachFilesService.get_custom_pba_directory(), file_name) - try: - if os.path.exists(file_path): - os.remove(file_path) - except OSError as e: - logger.error("Can't remove previously uploaded post breach files: %s" % e) - - @staticmethod - def get_custom_pba_directory(): - return os.path.join( - PostBreachFilesService.DATA_DIR, PostBreachFilesService.CUSTOM_PBA_DIRNAME - ) + @classmethod + def remove_PBA_files(cls): + cls._file_storage_service.delete_all_files() diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py index a0b8352e2..ca5f0f175 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py @@ -2,7 +2,6 @@ import pytest from tests.utils import raise_ from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE -from monkey_island.cc.services.post_breach_files import PostBreachFilesService TEST_FILE = b"""-----------------------------1 Content-Disposition: form-data; name="filepond" @@ -16,11 +15,6 @@ m0nk3y -----------------------------1--""" -@pytest.fixture(autouse=True) -def custom_pba_directory(tmpdir): - PostBreachFilesService.initialize(tmpdir) - - @pytest.fixture def mock_set_config_value(monkeypatch): monkeypatch.setattr( diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py index 90a649a39..23f42cf79 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py @@ -1,29 +1,31 @@ +import io import os import pytest -from tests.monkey_island.utils import assert_windows_permissions from tests.utils import raise_ -from monkey_island.cc.server_utils.file_utils import is_windows_os +from monkey_island.cc.services import DirectoryFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService +@pytest.fixture +def file_storage_service(tmp_path): + return DirectoryFileStorageService(tmp_path) + + @pytest.fixture(autouse=True) -def custom_pba_directory(tmpdir): - PostBreachFilesService.initialize(tmpdir) +def post_breach_files_service(file_storage_service): + PostBreachFilesService.initialize(file_storage_service) -def create_custom_pba_file(filename): - PostBreachFilesService.save_file(filename, b"") +def test_remove_pba_files(file_storage_service, tmp_path): + file_storage_service.save_file("linux_file", io.BytesIO(b"")) + file_storage_service.save_file("windows_file", io.BytesIO(b"")) + assert not dir_is_empty(tmp_path) - -def test_remove_pba_files(): - create_custom_pba_file("linux_file") - create_custom_pba_file("windows_file") - - assert not dir_is_empty(PostBreachFilesService.get_custom_pba_directory()) PostBreachFilesService.remove_PBA_files() - assert dir_is_empty(PostBreachFilesService.get_custom_pba_directory()) + + assert dir_is_empty(tmp_path) def dir_is_empty(dir_path): @@ -31,45 +33,11 @@ def dir_is_empty(dir_path): return len(dir_contents) == 0 -@pytest.mark.skipif(is_windows_os(), reason="Tests Posix (not Windows) permissions.") -def test_custom_pba_dir_permissions_linux(): - st = os.stat(PostBreachFilesService.get_custom_pba_directory()) - - assert st.st_mode == 0o40700 - - -@pytest.mark.skipif(not is_windows_os(), reason="Tests Windows (not Posix) permissions.") -def test_custom_pba_dir_permissions_windows(): - pba_dir = PostBreachFilesService.get_custom_pba_directory() - - assert_windows_permissions(pba_dir) - - -def test_remove_failure(monkeypatch): +def test_remove_failure(file_storage_service, monkeypatch): monkeypatch.setattr(os, "remove", lambda x: raise_(OSError("Permission denied"))) try: - create_custom_pba_file("windows_file") + file_storage_service.save_file("windows_file", io.BytesIO(b"")) PostBreachFilesService.remove_PBA_files() except Exception as ex: pytest.fail(f"Unxepected exception: {ex}") - - -def test_remove_nonexistant_file(monkeypatch): - monkeypatch.setattr(os, "remove", lambda x: raise_(FileNotFoundError("FileNotFound"))) - - try: - PostBreachFilesService.remove_file("/nonexistant/file") - except Exception as ex: - pytest.fail(f"Unxepected exception: {ex}") - - -def test_save_file(): - FILE_NAME = "test_file" - FILE_CONTENTS = b"hello" - PostBreachFilesService.save_file(FILE_NAME, FILE_CONTENTS) - - expected_file_path = os.path.join(PostBreachFilesService.get_custom_pba_directory(), FILE_NAME) - - assert os.path.isfile(expected_file_path) - assert FILE_CONTENTS == open(expected_file_path, "rb").read()