diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 8fadb9ccf..52e6f6873 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,4 +1,4 @@ -from .errors import RetrievalError +from .errors import RemovalError, RetrievalError, StorageError from .file_storage import FileNotFoundError, IFileRepository, LocalStorageFileRepository from .i_agent_binary_repository import IAgentBinaryRepository from .agent_binary_repository import AgentBinaryRepository diff --git a/monkey/monkey_island/cc/repository/errors.py b/monkey/monkey_island/cc/repository/errors.py index 586f23cb2..aeb9fa23c 100644 --- a/monkey/monkey_island/cc/repository/errors.py +++ b/monkey/monkey_island/cc/repository/errors.py @@ -1,3 +1,11 @@ +class RemovalError(RuntimeError): + """ + Raised when a repository encounters an error while attempting to remove data. + """ + + pass + + class RetrievalError(RuntimeError): """ Raised when a repository encounters an error while attempting to retrieve data. 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 bb9e345dc..baa52fdda 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 @@ -20,6 +20,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 StorageError: If an error was encountered while attempting to store the file """ pass @@ -44,6 +45,7 @@ class IFileRepository(metaclass=abc.ABCMeta): idempotent and will succeed if the file to be deleted does not exist. :param unsafe_file_name: An unsanitized file name that identifies the file to be deleted + :raises RemovalError: If an error was encountered while attempting to remove a file """ pass @@ -51,5 +53,7 @@ class IFileRepository(metaclass=abc.ABCMeta): def delete_all_files(self): """ Delete all files that have been stored using this service. + + :raises RemovalError: If an error was encountered while attempting to remove a 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..4c7662b0c 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 RemovalError, RetrievalError, StorageError from monkey_island.cc.server_utils.file_utils import create_secure_directory from . import IFileRepository, i_file_repository @@ -32,16 +32,19 @@ class LocalStorageFileRepository(IFileRepository): self._storage_directory = storage_directory def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): - safe_file_path = self._get_safe_file_path(unsafe_file_name) + try: + 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) + logger.debug(f"Saving file contents to {safe_file_path}") + 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) - try: + safe_file_path = self._get_safe_file_path(unsafe_file_name) + logger.debug(f"Opening {safe_file_path}") return open(safe_file_path, "rb") except FileNotFoundError as err: @@ -55,14 +58,16 @@ class LocalStorageFileRepository(IFileRepository): ) def delete_file(self, unsafe_file_name: str): - safe_file_path = self._get_safe_file_path(unsafe_file_name) - try: + safe_file_path = self._get_safe_file_path(unsafe_file_name) + logger.debug(f"Deleting {safe_file_path}") safe_file_path.unlink() except FileNotFoundError: # This method is idempotent. pass + except Exception as err: + raise RemovalError(f"Error while attempting to remove {unsafe_file_name}: {err}") def _get_safe_file_path(self, unsafe_file_name: str): # Remove any path information from the file name. @@ -71,12 +76,17 @@ class LocalStorageFileRepository(IFileRepository): # This is a paranoid check to avoid directory traversal attacks. if self._storage_directory.resolve() not in safe_file_path.parents: - raise ValueError(f"The file named {unsafe_file_name} can not be safely retrieved") + raise ValueError( + f'The file named "{unsafe_file_name}" cannot be safely retrieved or written' + ) logger.debug(f"Untrusted file name {unsafe_file_name} sanitized: {safe_file_path}") return safe_file_path def delete_all_files(self): - for file in get_all_regular_files_in_directory(self._storage_directory): - logger.debug(f"Deleting {file}") - file.unlink() + try: + for file in get_all_regular_files_in_directory(self._storage_directory): + logger.debug(f"Deleting {file}") + file.unlink() + except Exception as err: + raise RemovalError(f"Error while attempting to clear the repository: {err}")