Island: Add repository.FileNotFoundError

This commit is contained in:
Mike Salvatore 2022-06-21 08:49:36 -04:00
parent c424262f12
commit cd34cd5eae
10 changed files with 46 additions and 36 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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