From cd34cd5eae4d2fb22958ff20a1172ba6d55b5e87 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 08:49:36 -0400 Subject: [PATCH] Island: Add repository.FileNotFoundError --- .../monkey_island/cc/repository/__init__.py | 2 +- .../cc/repository/agent_binary_repository.py | 7 +++++-- .../cc/repository/file_storage/__init__.py | 2 +- .../file_storage/i_file_repository.py | 4 ++-- .../local_storage_file_repository.py | 20 ++++++++++--------- .../cc/resources/pba_file_download.py | 13 +++++++----- .../cc/resources/pba_file_upload.py | 13 +++++++----- .../monkey_island/single_file_repository.py | 5 +++-- .../test_local_storage_file_repository.py | 12 +++++------ .../cc/resources/test_pba_file_download.py | 4 ++-- 10 files changed, 46 insertions(+), 36 deletions(-) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index cc89c5c57..a44e7b1ad 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,5 +1,5 @@ from .errors import RetrievalError -from .file_storage import FileRetrievalError, IFileRepository, LocalStorageFileRepository +from .file_storage import FileNotFoundError, IFileRepository, LocalStorageFileRepository from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError from .agent_binary_repository import AgentBinaryRepository from .i_agent_configuration_repository import IAgentConfigurationRepository diff --git a/monkey/monkey_island/cc/repository/agent_binary_repository.py b/monkey/monkey_island/cc/repository/agent_binary_repository.py index 4c1334f71..a09d88868 100644 --- a/monkey/monkey_island/cc/repository/agent_binary_repository.py +++ b/monkey/monkey_island/cc/repository/agent_binary_repository.py @@ -1,6 +1,8 @@ from typing import BinaryIO -from . import AgentRetrievalError, FileRetrievalError, IAgentBinaryRepository, IFileRepository +from monkey_island.cc import repository + +from . import AgentRetrievalError, IAgentBinaryRepository, IFileRepository LINUX_AGENT_FILE_NAME = "monkey-linux-64" WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe" @@ -20,7 +22,8 @@ class AgentBinaryRepository(IAgentBinaryRepository): try: agent_binary = self._file_repository.open_file(filename) return agent_binary - except FileRetrievalError as err: + # TODO: Reevaluate this + except repository.FileNotFoundError as err: raise AgentRetrievalError( f"An error occurred while retrieving the {filename}" f" agent binary from {self._file_repository}: {err}" diff --git a/monkey/monkey_island/cc/repository/file_storage/__init__.py b/monkey/monkey_island/cc/repository/file_storage/__init__.py index 30b59b5d9..14227bc66 100644 --- a/monkey/monkey_island/cc/repository/file_storage/__init__.py +++ b/monkey/monkey_island/cc/repository/file_storage/__init__.py @@ -1,2 +1,2 @@ -from .i_file_repository import IFileRepository, FileRetrievalError +from .i_file_repository import IFileRepository, FileNotFoundError from .local_storage_file_repository import LocalStorageFileRepository 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 5069b36d7..681b3d7cf 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 @@ -4,7 +4,7 @@ from typing import BinaryIO from monkey_island.cc.repository import RetrievalError -class FileRetrievalError(RetrievalError): +class FileNotFoundError(RetrievalError): pass @@ -30,8 +30,8 @@ class IFileRepository(metaclass=abc.ABCMeta): :param unsafe_file_name: An unsanitized file name that identifies the file to be opened :return: A file-like object providing access to the file's contents - :raises FileRetrievalError: if the file cannot be retrieved :raises FileNotFoundError: if the file cannot be found + :raises RetrievalError: if the file cannot be retrieved """ 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 a23b25228..cb18c6ba8 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 @@ -1,13 +1,13 @@ -import errno import logging import shutil 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.server_utils.file_utils import create_secure_directory -from . import FileRetrievalError, IFileRepository +from . import IFileRepository, i_file_repository logger = logging.getLogger(__name__) @@ -44,13 +44,15 @@ class LocalStorageFileRepository(IFileRepository): try: logger.debug(f"Opening {safe_file_path}") return open(safe_file_path, "rb") - except OSError as err: - if err.errno == errno.ENOENT: - raise FileNotFoundError(f"No file {safe_file_path} found: {err}") from err - else: - raise FileRetrievalError( - f"Failed to retrieve file {safe_file_path}: {err}" - ) from err + except FileNotFoundError as err: + # Wrap Python's FileNotFound error, which is-an OSError, in repository.FileNotFoundError + raise i_file_repository.FileNotFoundError( + f'The requested file "{unsafe_file_name}" does not exist: {err}' + ) + except Exception as err: + raise RetrievalError( + f'Error retrieving file "{unsafe_file_name}" from the repository: {err}' + ) def delete_file(self, unsafe_file_name: str): safe_file_path = self._get_safe_file_path(unsafe_file_name) diff --git a/monkey/monkey_island/cc/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index 85395d5e3..3b7e92047 100644 --- a/monkey/monkey_island/cc/resources/pba_file_download.py +++ b/monkey/monkey_island/cc/resources/pba_file_download.py @@ -2,7 +2,8 @@ import logging from flask import make_response, send_file -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__file__) @@ -24,7 +25,9 @@ class PBAFileDownload(AbstractResource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") - except FileRetrievalError as err: - error_msg = f"Failed to open file {filename}: {err}" - logger.error(error_msg) - return make_response({"error": error_msg}, 404) + except repository.FileNotFoundError as err: + # TODO: Do we need to log? Or will flask handle it when we `make_response()`? + logger.error(str(err)) + return make_response({"error": str(err)}, 404) + + # TODO: Add unit tests that test 404 vs 500 errors diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 9d5cce8b1..833086998 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -5,7 +5,8 @@ from flask import Response, make_response, request, send_file from werkzeug.utils import secure_filename as sanitize_filename from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.config import ConfigService @@ -57,10 +58,12 @@ class FileUpload(AbstractResource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") - except FileRetrievalError as err: - error_msg = f"Failed to open file {filename}: {err}" - logger.error(error_msg) - return make_response({"error": error_msg}, 404) + except repository.FileNotFoundError as err: + # TODO: Do we need to log? Or will flask handle it when we `make_response()`? + logger.error(str(err)) + return make_response({"error": str(err)}, 404) + + # TODO: Add unit tests that test 404 vs 500 errors @jwt_required def post(self, target_os): diff --git a/monkey/tests/monkey_island/single_file_repository.py b/monkey/tests/monkey_island/single_file_repository.py index d00ee9ae9..462969acb 100644 --- a/monkey/tests/monkey_island/single_file_repository.py +++ b/monkey/tests/monkey_island/single_file_repository.py @@ -1,7 +1,8 @@ import io from typing import BinaryIO -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository class SingleFileRepository(IFileRepository): @@ -13,7 +14,7 @@ class SingleFileRepository(IFileRepository): def open_file(self, unsafe_file_name: str) -> BinaryIO: if self._file is None: - raise FileRetrievalError() + raise repository.FileNotFoundError() return self._file def delete_file(self, unsafe_file_name: str): diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py index 804b03c79..64541bcbe 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py @@ -1,4 +1,3 @@ -import errno import io from pathlib import Path from unittest.mock import Mock, patch @@ -6,7 +5,8 @@ from unittest.mock import Mock, patch import pytest from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions -from monkey_island.cc.repository import FileRetrievalError, LocalStorageFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import LocalStorageFileRepository from monkey_island.cc.server_utils.file_utils import is_windows_os @@ -133,15 +133,13 @@ def test_remove_nonexistant_file(tmp_path): def test_open_nonexistant_file(tmp_path): fss = LocalStorageFileRepository(tmp_path) - with pytest.raises(FileNotFoundError): + with pytest.raises(repository.FileNotFoundError): fss.open_file("nonexistant_file.txt") def test_open_locked_file(tmp_path, monkeypatch): fss = LocalStorageFileRepository(tmp_path) - with patch( - "builtins.open", Mock(side_effect=OSError(errno.EIO, "File could not be retrieved")) - ): - with pytest.raises(FileRetrievalError): + with patch("builtins.open", Mock(side_effect=Exception())): + with pytest.raises(repository.RetrievalError): fss.open_file("locked_file.txt") 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 cf2a109b9..caf86edab 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 @@ -5,7 +5,7 @@ import pytest from tests.common import StubDIContainer from tests.unit_tests.monkey_island.conftest import get_url_for_resource -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc.repository import FileNotFoundError, IFileRepository from monkey_island.cc.resources.pba_file_download import PBAFileDownload FILE_NAME = "test_file" @@ -21,7 +21,7 @@ class MockFileRepository(IFileRepository): def open_file(self, unsafe_file_name: str) -> BinaryIO: if unsafe_file_name != FILE_NAME: - raise FileRetrievalError() + raise FileNotFoundError() return self._file