forked from p15670423/monkey
Island: Add repository.FileNotFoundError
This commit is contained in:
parent
c424262f12
commit
cd34cd5eae
|
@ -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
|
||||
|
|
|
@ -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}"
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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):
|
||||
|
|
|
@ -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):
|
||||
|
|
|
@ -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")
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
Loading…
Reference in New Issue