p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Island: Raise StorageError in LocalStorageFileRepository.save_file()

This commit is contained in:
Mike Salvatore 2022-06-21 13:19:56 -04:00
parent 3446dbf0aa
commit 63404c7bed
2 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
monkey/monkey_island/cc/repository/file_storage/i_file_repository.py
View File

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

9
monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py
View File

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