From 63404c7bed2c0f3ffed701e92f747b3808059cc0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 13:19:56 -0400 Subject: [PATCH] Island: Raise StorageError in LocalStorageFileRepository.save_file() --- .../cc/repository/file_storage/i_file_repository.py | 1 + .../file_storage/local_storage_file_repository.py | 9 ++++++--- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py b/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py index 237cc2c0d..73d3b4759 100644 --- a/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py @@ -21,6 +21,7 @@ class IFileRepository(metaclass=abc.ABCMeta): :param unsafe_file_name: An unsanitized file name that will identify the file :param file_contents: The data to be stored in the file :raises ValueError: If the file name is an attempted directory traversal + :raises StorageError: If an error was encountered while attempting to store the file """ pass diff --git a/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py b/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py index cb18c6ba8..a766746ae 100644 --- a/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py +++ b/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py @@ -4,7 +4,7 @@ from pathlib import Path from typing import BinaryIO from common.utils.file_utils import get_all_regular_files_in_directory -from monkey_island.cc.repository import RetrievalError +from monkey_island.cc.repository import RetrievalError, StorageError from monkey_island.cc.server_utils.file_utils import create_secure_directory from . import IFileRepository, i_file_repository @@ -35,8 +35,11 @@ class LocalStorageFileRepository(IFileRepository): safe_file_path = self._get_safe_file_path(unsafe_file_name) logger.debug(f"Saving file contents to {safe_file_path}") - with open(safe_file_path, "wb") as dest: - shutil.copyfileobj(file_contents, dest) + try: + with open(safe_file_path, "wb") as dest: + shutil.copyfileobj(file_contents, dest) + except Exception as err: + raise StorageError(f"Error while attempting to store {unsafe_file_name}: {err}") def open_file(self, unsafe_file_name: str) -> BinaryIO: safe_file_path = self._get_safe_file_path(unsafe_file_name)