Island: Use IFileStorageService in FileUpload resource

monkey/monkey_island/cc/app.py

@@ -150,17 +150,22 @@ def init_api_resources(api, data_dir: Path):
     api.add_resource(TelemetryFeed, "/api/telemetry-feed")
     api.add_resource(Log, "/api/log")
     api.add_resource(IslandLog, "/api/log/island/download")
+
+    # This is temporary until we get DI all worked out.
+    file_storage_service = DirectoryFileStorageService(data_dir / "custom_pbas")
     api.add_resource(
         PBAFileDownload,
         "/api/pba/download/<string:filename>",
-        resource_class_kwargs={"file_storage_service": DirectoryFileStorageService(data_dir)},
+        resource_class_kwargs={"file_storage_service": file_storage_service},
     )
     api.add_resource(
         FileUpload,
         "/api/file-upload/<string:file_type>",
         "/api/file-upload/<string:file_type>?load=<string:filename>",
         "/api/file-upload/<string:file_type>?restore=<string:filename>",
+        resource_class_kwargs={"file_storage_service": file_storage_service},
     )
+
     api.add_resource(PropagationCredentials, "/api/propagation-credentials/<string:guid>")
     api.add_resource(RemoteRun, "/api/remote-monkey")
     api.add_resource(VersionUpdate, "/api/version-update")

monkey/monkey_island/cc/resources/pba_file_upload.py

@@ -1,15 +1,19 @@
 import copy
+import logging
 from http import HTTPStatus
 
 import flask_restful
-from flask import Response, request, send_from_directory
+from flask import Response, make_response, request, send_file
 from werkzeug.datastructures import FileStorage
 from werkzeug.utils import secure_filename
 
 from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH
 from monkey_island.cc.resources.auth.auth import jwt_required
+from monkey_island.cc.services import IFileStorageService
 from monkey_island.cc.services.config import ConfigService
-from monkey_island.cc.services.post_breach_files import PostBreachFilesService
+
+logger = logging.getLogger(__file__)
+
 
 # Front end uses these strings to identify which files to work with (linux or windows)
 LINUX_PBA_TYPE = "PBAlinux"
@@ -21,6 +25,14 @@ class FileUpload(flask_restful.Resource):
     File upload endpoint used to exchange files with filepond component on the front-end
     """
 
+    def __init__(self, file_storage_service: IFileStorageService):
+        self._file_storage_service = file_storage_service
+
+    # TODO: Fix references/coupling to filepond
+    # TODO: Replace "file_type" with "target_os" or similar
+    # TODO: Prefix private functions with "_"
+    # TODO: Add comment explaining why this is basically a duplicate of the endpoint in the
+    #       PBAFileDownload resource.
     @jwt_required
     def get(self, file_type):
         """
@@ -33,10 +45,20 @@ class FileUpload(flask_restful.Resource):
 
         # Verify that file_name is indeed a file from config
         if file_type == LINUX_PBA_TYPE:
+            # TODO: Make these paths Tuples so we don't need to copy them
             filename = ConfigService.get_config_value(copy.deepcopy(PBA_LINUX_FILENAME_PATH))
         else:
             filename = ConfigService.get_config_value(copy.deepcopy(PBA_WINDOWS_FILENAME_PATH))
-        return send_from_directory(PostBreachFilesService.get_custom_pba_directory(), filename)
+
+        try:
+            file = self._file_storage_service.open_file(filename)
+
+            # `send_file()` handles the closing of the open file.
+            return send_file(file, mimetype="application/octet-stream")
+        except OSError as ex:
+            error_msg = f"Failed to open file {filename}: {ex}"
+            logger.error(error_msg)
+            return make_response({"error": error_msg}, 404)
 
     @jwt_required
     def post(self, file_type):
@@ -48,15 +70,17 @@ class FileUpload(flask_restful.Resource):
         if self.is_pba_file_type_supported(file_type):
             return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain")
 
-        filename = FileUpload.upload_pba_file(
-            request.files["filepond"], (file_type == LINUX_PBA_TYPE)
+        filename = self.upload_pba_file(
+            # TODO: This "filepond" string can be changed to be more generic in the `react-filepond`
+            # component.
+            request.files["filepond"],
+            (file_type == LINUX_PBA_TYPE),
         )
 
         response = Response(response=filename, status=200, mimetype="text/plain")
         return response
 
-    @staticmethod
-    def upload_pba_file(file_storage: FileStorage, is_linux=True):
+    def upload_pba_file(self, file_storage: FileStorage, is_linux=True):
         """
         Uploads PBA file to island's file system
         :param request_: Request object containing PBA file
@@ -64,9 +88,7 @@ class FileUpload(flask_restful.Resource):
         :return: filename string
         """
         filename = secure_filename(file_storage.filename)
-        file_contents = file_storage.read()
-
-        PostBreachFilesService.save_file(filename, file_contents)
+        self._file_storage_service.save_file(filename, file_storage.stream)
 
         ConfigService.set_config_value(
             (PBA_LINUX_FILENAME_PATH if is_linux else PBA_WINDOWS_FILENAME_PATH), filename
@@ -89,10 +111,10 @@ class FileUpload(flask_restful.Resource):
         )
         filename = ConfigService.get_config_value(filename_path)
         if filename:
-            PostBreachFilesService.remove_file(filename)
+            self._file_storage_service.delete_file(filename)
             ConfigService.set_config_value(filename_path, "")
 
-        return {}
+        return make_response({}, 200)
 
     @staticmethod
     def is_pba_file_type_supported(file_type: str) -> bool:

monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py

@@ -1,7 +1,7 @@
 def test_file_download_endpoint(tmp_path, flask_client):
     file_contents = "HelloWorld!"
     file_name = "test_file"
-    (tmp_path / file_name).write_text(file_contents)
+    (tmp_path / "custom_pbas" / file_name).write_text(file_contents)
 
     resp = flask_client.get(f"/api/pba/download/{file_name}")