Island: Modify PostBreachFilesService to wrap IFileStorageService

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

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

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

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(

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