diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 17e86fc23..beb9e1468 100644 --- a/monkey/monkey_island/cc/app.py +++ b/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/", - 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/", "/api/file-upload/?load=", "/api/file-upload/?restore=", + resource_class_kwargs={"file_storage_service": file_storage_service}, ) + api.add_resource(PropagationCredentials, "/api/propagation-credentials/") api.add_resource(RemoteRun, "/api/remote-monkey") api.add_resource(VersionUpdate, "/api/version-update") diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 68be20e9a..bb2608872 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/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: diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py index 333ff4cf1..5b6dcad3f 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py +++ b/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}")