forked from p34709852/monkey
Island: Raise FileRetrievalError in DirectoryFileStorageService
This commit is contained in:
parent
a0b4dc1bcb
commit
4ddcd5e9a8
monkey
monkey_island/cc
resources
services
tests/unit_tests/monkey_island/cc
|
@ -3,7 +3,7 @@ import logging
|
||||||
import flask_restful
|
import flask_restful
|
||||||
from flask import make_response, send_file
|
from flask import make_response, send_file
|
||||||
|
|
||||||
from monkey_island.cc.services import IFileStorageService
|
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||||
|
|
||||||
logger = logging.getLogger(__file__)
|
logger = logging.getLogger(__file__)
|
||||||
|
|
||||||
|
@ -23,7 +23,7 @@ class PBAFileDownload(flask_restful.Resource):
|
||||||
|
|
||||||
# `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 OSError as ex:
|
except FileRetrievalError as err:
|
||||||
error_msg = f"Failed to open file {filename}: {ex}"
|
error_msg = f"Failed to open file {filename}: {err}"
|
||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
return make_response({"error": error_msg}, 404)
|
return make_response({"error": error_msg}, 404)
|
||||||
|
|
|
@ -9,7 +9,7 @@ from werkzeug.utils import secure_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.resources.auth.auth import jwt_required
|
from monkey_island.cc.resources.auth.auth import jwt_required
|
||||||
from monkey_island.cc.services import IFileStorageService
|
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||||
from monkey_island.cc.services.config import ConfigService
|
from monkey_island.cc.services.config import ConfigService
|
||||||
|
|
||||||
logger = logging.getLogger(__file__)
|
logger = logging.getLogger(__file__)
|
||||||
|
@ -53,8 +53,8 @@ class FileUpload(flask_restful.Resource):
|
||||||
|
|
||||||
# `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 OSError as ex:
|
except FileRetrievalError as err:
|
||||||
error_msg = f"Failed to open file {filename}: {ex}"
|
error_msg = f"Failed to open file {filename}: {err}"
|
||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
return make_response({"error": error_msg}, 404)
|
return make_response({"error": error_msg}, 404)
|
||||||
|
|
||||||
|
|
|
@ -1,4 +1,4 @@
|
||||||
from .i_file_storage_service import IFileStorageService
|
from .i_file_storage_service import IFileStorageService, FileRetrievalError
|
||||||
from .directory_file_storage_service import DirectoryFileStorageService
|
from .directory_file_storage_service import DirectoryFileStorageService
|
||||||
|
|
||||||
from .authentication.authentication_service import AuthenticationService
|
from .authentication.authentication_service import AuthenticationService
|
||||||
|
|
|
@ -5,7 +5,7 @@ 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.server_utils.file_utils import create_secure_directory
|
from monkey_island.cc.server_utils.file_utils import create_secure_directory
|
||||||
|
|
||||||
from . import IFileStorageService
|
from . import FileRetrievalError, IFileStorageService
|
||||||
|
|
||||||
|
|
||||||
class DirectoryFileStorageService(IFileStorageService):
|
class DirectoryFileStorageService(IFileStorageService):
|
||||||
|
@ -35,7 +35,11 @@ class DirectoryFileStorageService(IFileStorageService):
|
||||||
|
|
||||||
def open_file(self, unsafe_file_name: str) -> BinaryIO:
|
def open_file(self, unsafe_file_name: str) -> BinaryIO:
|
||||||
safe_file_path = self._get_safe_file_path(unsafe_file_name)
|
safe_file_path = self._get_safe_file_path(unsafe_file_name)
|
||||||
return open(safe_file_path, "rb")
|
|
||||||
|
try:
|
||||||
|
return open(safe_file_path, "rb")
|
||||||
|
except OSError as err:
|
||||||
|
raise FileRetrievalError(f"Failed to retrieve file {safe_file_path}: {err}") from 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,6 +2,10 @@ import abc
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
|
||||||
|
|
||||||
|
class FileRetrievalError(ValueError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
class IFileStorageService(metaclass=abc.ABCMeta):
|
class IFileStorageService(metaclass=abc.ABCMeta):
|
||||||
"""
|
"""
|
||||||
A service that allows the storage and retrieval of individual files.
|
A service that allows the storage and retrieval of individual files.
|
||||||
|
@ -25,6 +29,7 @@ class IFileStorageService(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
|
||||||
:rtype: io.BinaryIO
|
:rtype: io.BinaryIO
|
||||||
|
:raises FileRetrievalError: if the file cannot be opened
|
||||||
"""
|
"""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
|
@ -4,7 +4,7 @@ from typing import BinaryIO
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from common import DIContainer
|
from common import DIContainer
|
||||||
from monkey_island.cc.services import IFileStorageService
|
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||||
|
|
||||||
FILE_NAME = "test_file"
|
FILE_NAME = "test_file"
|
||||||
FILE_CONTENTS = b"HelloWorld!"
|
FILE_CONTENTS = b"HelloWorld!"
|
||||||
|
@ -19,7 +19,7 @@ class MockFileStorageService(IFileStorageService):
|
||||||
|
|
||||||
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 OSError()
|
raise FileRetrievalError()
|
||||||
|
|
||||||
return self._file
|
return self._file
|
||||||
|
|
||||||
|
|
|
@ -6,7 +6,7 @@ from tests.utils import raise_
|
||||||
|
|
||||||
from common import DIContainer
|
from common import DIContainer
|
||||||
from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE
|
from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE
|
||||||
from monkey_island.cc.services import IFileStorageService
|
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||||
|
|
||||||
TEST_FILE_CONTENTS = b"m0nk3y"
|
TEST_FILE_CONTENTS = b"m0nk3y"
|
||||||
TEST_FILE = (
|
TEST_FILE = (
|
||||||
|
@ -48,8 +48,7 @@ class MockFileStorageService(IFileStorageService):
|
||||||
|
|
||||||
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:
|
||||||
# TODO: Add FileRetrievalError
|
raise FileRetrievalError()
|
||||||
raise OSError()
|
|
||||||
return self._file
|
return self._file
|
||||||
|
|
||||||
def delete_file(self, unsafe_file_name: str):
|
def delete_file(self, unsafe_file_name: str):
|
||||||
|
|
|
@ -5,10 +5,10 @@ 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.server_utils.file_utils import is_windows_os
|
from monkey_island.cc.server_utils.file_utils import is_windows_os
|
||||||
from monkey_island.cc.services import DirectoryFileStorageService
|
from monkey_island.cc.services import DirectoryFileStorageService, FileRetrievalError
|
||||||
|
|
||||||
|
|
||||||
def test_error_if_file(tmp_path):
|
def test_error_if_storage_directory_is_file(tmp_path):
|
||||||
new_file = tmp_path / "new_file.txt"
|
new_file = tmp_path / "new_file.txt"
|
||||||
new_file.write_text("HelloWorld!")
|
new_file.write_text("HelloWorld!")
|
||||||
|
|
||||||
|
@ -126,3 +126,10 @@ def test_remove_nonexistant_file(tmp_path):
|
||||||
|
|
||||||
# This test will fail if this call raises an exception.
|
# This test will fail if this call raises an exception.
|
||||||
fss.delete_file("nonexistant_file.txt")
|
fss.delete_file("nonexistant_file.txt")
|
||||||
|
|
||||||
|
|
||||||
|
def test_open_nonexistant_file(tmp_path):
|
||||||
|
fss = DirectoryFileStorageService(tmp_path)
|
||||||
|
|
||||||
|
with pytest.raises(FileRetrievalError):
|
||||||
|
fss.open_file("nonexistant_file.txt")
|
||||||
|
|
Loading…
Reference in New Issue