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 .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 .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError
|
||||||
from .agent_binary_repository import AgentBinaryRepository
|
from .agent_binary_repository import AgentBinaryRepository
|
||||||
from .i_agent_configuration_repository import IAgentConfigurationRepository
|
from .i_agent_configuration_repository import IAgentConfigurationRepository
|
||||||
|
|
|
@ -1,6 +1,8 @@
|
||||||
from typing import BinaryIO
|
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"
|
LINUX_AGENT_FILE_NAME = "monkey-linux-64"
|
||||||
WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe"
|
WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe"
|
||||||
|
@ -20,7 +22,8 @@ class AgentBinaryRepository(IAgentBinaryRepository):
|
||||||
try:
|
try:
|
||||||
agent_binary = self._file_repository.open_file(filename)
|
agent_binary = self._file_repository.open_file(filename)
|
||||||
return agent_binary
|
return agent_binary
|
||||||
except FileRetrievalError as err:
|
# TODO: Reevaluate this
|
||||||
|
except repository.FileNotFoundError as err:
|
||||||
raise AgentRetrievalError(
|
raise AgentRetrievalError(
|
||||||
f"An error occurred while retrieving the {filename}"
|
f"An error occurred while retrieving the {filename}"
|
||||||
f" agent binary from {self._file_repository}: {err}"
|
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
|
from .local_storage_file_repository import LocalStorageFileRepository
|
||||||
|
|
|
@ -4,7 +4,7 @@ from typing import BinaryIO
|
||||||
from monkey_island.cc.repository import RetrievalError
|
from monkey_island.cc.repository import RetrievalError
|
||||||
|
|
||||||
|
|
||||||
class FileRetrievalError(RetrievalError):
|
class FileNotFoundError(RetrievalError):
|
||||||
pass
|
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
|
: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
|
: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 FileNotFoundError: if the file cannot be found
|
||||||
|
:raises RetrievalError: if the file cannot be retrieved
|
||||||
"""
|
"""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
|
@ -1,13 +1,13 @@
|
||||||
import errno
|
|
||||||
import logging
|
import logging
|
||||||
import shutil
|
import shutil
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
|
||||||
from common.utils.file_utils import get_all_regular_files_in_directory
|
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 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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
@ -44,13 +44,15 @@ class LocalStorageFileRepository(IFileRepository):
|
||||||
try:
|
try:
|
||||||
logger.debug(f"Opening {safe_file_path}")
|
logger.debug(f"Opening {safe_file_path}")
|
||||||
return open(safe_file_path, "rb")
|
return open(safe_file_path, "rb")
|
||||||
except OSError as err:
|
except FileNotFoundError as err:
|
||||||
if err.errno == errno.ENOENT:
|
# Wrap Python's FileNotFound error, which is-an OSError, in repository.FileNotFoundError
|
||||||
raise FileNotFoundError(f"No file {safe_file_path} found: {err}") from err
|
raise i_file_repository.FileNotFoundError(
|
||||||
else:
|
f'The requested file "{unsafe_file_name}" does not exist: {err}'
|
||||||
raise FileRetrievalError(
|
)
|
||||||
f"Failed to retrieve file {safe_file_path}: {err}"
|
except Exception as err:
|
||||||
) from err
|
raise RetrievalError(
|
||||||
|
f'Error retrieving file "{unsafe_file_name}" from the repository: {err}'
|
||||||
|
)
|
||||||
|
|
||||||
def delete_file(self, unsafe_file_name: str):
|
def delete_file(self, unsafe_file_name: str):
|
||||||
safe_file_path = self._get_safe_file_path(unsafe_file_name)
|
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 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
|
from monkey_island.cc.resources.AbstractResource import AbstractResource
|
||||||
|
|
||||||
logger = logging.getLogger(__file__)
|
logger = logging.getLogger(__file__)
|
||||||
|
@ -24,7 +25,9 @@ class PBAFileDownload(AbstractResource):
|
||||||
|
|
||||||
# `send_file()` handles the closing of the open file.
|
# `send_file()` handles the closing of the open file.
|
||||||
return send_file(file, mimetype="application/octet-stream")
|
return send_file(file, mimetype="application/octet-stream")
|
||||||
except FileRetrievalError as err:
|
except repository.FileNotFoundError as err:
|
||||||
error_msg = f"Failed to open file {filename}: {err}"
|
# TODO: Do we need to log? Or will flask handle it when we `make_response()`?
|
||||||
logger.error(error_msg)
|
logger.error(str(err))
|
||||||
return make_response({"error": error_msg}, 404)
|
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 werkzeug.utils import secure_filename as sanitize_filename
|
||||||
|
|
||||||
from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH
|
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.AbstractResource import AbstractResource
|
||||||
from monkey_island.cc.resources.request_authentication import jwt_required
|
from monkey_island.cc.resources.request_authentication import jwt_required
|
||||||
from monkey_island.cc.services.config import ConfigService
|
from monkey_island.cc.services.config import ConfigService
|
||||||
|
@ -57,10 +58,12 @@ class FileUpload(AbstractResource):
|
||||||
|
|
||||||
# `send_file()` handles the closing of the open file.
|
# `send_file()` handles the closing of the open file.
|
||||||
return send_file(file, mimetype="application/octet-stream")
|
return send_file(file, mimetype="application/octet-stream")
|
||||||
except FileRetrievalError as err:
|
except repository.FileNotFoundError as err:
|
||||||
error_msg = f"Failed to open file {filename}: {err}"
|
# TODO: Do we need to log? Or will flask handle it when we `make_response()`?
|
||||||
logger.error(error_msg)
|
logger.error(str(err))
|
||||||
return make_response({"error": error_msg}, 404)
|
return make_response({"error": str(err)}, 404)
|
||||||
|
|
||||||
|
# TODO: Add unit tests that test 404 vs 500 errors
|
||||||
|
|
||||||
@jwt_required
|
@jwt_required
|
||||||
def post(self, target_os):
|
def post(self, target_os):
|
||||||
|
|
|
@ -1,7 +1,8 @@
|
||||||
import io
|
import io
|
||||||
from typing import BinaryIO
|
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):
|
class SingleFileRepository(IFileRepository):
|
||||||
|
@ -13,7 +14,7 @@ class SingleFileRepository(IFileRepository):
|
||||||
|
|
||||||
def open_file(self, unsafe_file_name: str) -> BinaryIO:
|
def open_file(self, unsafe_file_name: str) -> BinaryIO:
|
||||||
if self._file is None:
|
if self._file is None:
|
||||||
raise FileRetrievalError()
|
raise repository.FileNotFoundError()
|
||||||
return self._file
|
return self._file
|
||||||
|
|
||||||
def delete_file(self, unsafe_file_name: str):
|
def delete_file(self, unsafe_file_name: str):
|
||||||
|
|
|
@ -1,4 +1,3 @@
|
||||||
import errno
|
|
||||||
import io
|
import io
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import Mock, patch
|
from unittest.mock import Mock, patch
|
||||||
|
@ -6,7 +5,8 @@ from unittest.mock import Mock, patch
|
||||||
import pytest
|
import pytest
|
||||||
from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions
|
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
|
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):
|
def test_open_nonexistant_file(tmp_path):
|
||||||
fss = LocalStorageFileRepository(tmp_path)
|
fss = LocalStorageFileRepository(tmp_path)
|
||||||
|
|
||||||
with pytest.raises(FileNotFoundError):
|
with pytest.raises(repository.FileNotFoundError):
|
||||||
fss.open_file("nonexistant_file.txt")
|
fss.open_file("nonexistant_file.txt")
|
||||||
|
|
||||||
|
|
||||||
def test_open_locked_file(tmp_path, monkeypatch):
|
def test_open_locked_file(tmp_path, monkeypatch):
|
||||||
fss = LocalStorageFileRepository(tmp_path)
|
fss = LocalStorageFileRepository(tmp_path)
|
||||||
|
|
||||||
with patch(
|
with patch("builtins.open", Mock(side_effect=Exception())):
|
||||||
"builtins.open", Mock(side_effect=OSError(errno.EIO, "File could not be retrieved"))
|
with pytest.raises(repository.RetrievalError):
|
||||||
):
|
|
||||||
with pytest.raises(FileRetrievalError):
|
|
||||||
fss.open_file("locked_file.txt")
|
fss.open_file("locked_file.txt")
|
||||||
|
|
|
@ -5,7 +5,7 @@ import pytest
|
||||||
from tests.common import StubDIContainer
|
from tests.common import StubDIContainer
|
||||||
from tests.unit_tests.monkey_island.conftest import get_url_for_resource
|
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
|
from monkey_island.cc.resources.pba_file_download import PBAFileDownload
|
||||||
|
|
||||||
FILE_NAME = "test_file"
|
FILE_NAME = "test_file"
|
||||||
|
@ -21,7 +21,7 @@ class MockFileRepository(IFileRepository):
|
||||||
|
|
||||||
def open_file(self, unsafe_file_name: str) -> BinaryIO:
|
def open_file(self, unsafe_file_name: str) -> BinaryIO:
|
||||||
if unsafe_file_name != FILE_NAME:
|
if unsafe_file_name != FILE_NAME:
|
||||||
raise FileRetrievalError()
|
raise FileNotFoundError()
|
||||||
|
|
||||||
return self._file
|
return self._file
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue