forked from p15670423/monkey
Island: Rename and move file storage classes
This commit is contained in:
parent
b3710e4772
commit
154d59b837
|
@ -0,0 +1,2 @@
|
|||
from .file_storage.filesystem_storage import FilesystemStorage
|
||||
from .file_storage.i_file_repository import IFileRepository, FileRetrievalError
|
|
@ -6,12 +6,12 @@ from typing import BinaryIO
|
|||
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 . import FileRetrievalError, IFileStorageService
|
||||
from .i_file_repository import FileRetrievalError, IFileRepository
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DirectoryFileStorageService(IFileStorageService):
|
||||
class FilesystemStorage(IFileRepository):
|
||||
"""
|
||||
A implementation of IFileStorageService that reads and writes files from/to the local
|
||||
filesystem.
|
|
@ -6,7 +6,7 @@ class FileRetrievalError(ValueError):
|
|||
pass
|
||||
|
||||
|
||||
class IFileStorageService(metaclass=abc.ABCMeta):
|
||||
class IFileRepository(metaclass=abc.ABCMeta):
|
||||
"""
|
||||
A service that allows the storage and retrieval of individual files.
|
||||
"""
|
|
@ -2,8 +2,8 @@ import logging
|
|||
|
||||
from flask import make_response, send_file
|
||||
|
||||
from monkey_island.cc.repository import FileRetrievalError, IFileRepository
|
||||
from monkey_island.cc.resources.AbstractResource import AbstractResource
|
||||
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||
|
||||
logger = logging.getLogger(__file__)
|
||||
|
||||
|
@ -14,7 +14,7 @@ class PBAFileDownload(AbstractResource):
|
|||
File download endpoint used by monkey to download user's PBA file
|
||||
"""
|
||||
|
||||
def __init__(self, file_storage_service: IFileStorageService):
|
||||
def __init__(self, file_storage_service: IFileRepository):
|
||||
self._file_storage_service = file_storage_service
|
||||
|
||||
# Used by monkey. can't secure.
|
||||
|
|
|
@ -5,9 +5,9 @@ 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.resources.AbstractResource import AbstractResource
|
||||
from monkey_island.cc.resources.request_authentication import jwt_required
|
||||
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||
from monkey_island.cc.services.config import ConfigService
|
||||
|
||||
logger = logging.getLogger(__file__)
|
||||
|
@ -29,8 +29,8 @@ class FileUpload(AbstractResource):
|
|||
"/api/file-upload/<string:target_os>?restore=<string:filename>",
|
||||
]
|
||||
|
||||
def __init__(self, file_storage_service: IFileStorageService):
|
||||
self._file_storage_service = file_storage_service
|
||||
def __init__(self, file_storage_repository: IFileRepository):
|
||||
self._file_storage_service = file_storage_repository
|
||||
|
||||
# This endpoint is basically a duplicate of PBAFileDownload.get(). They serve slightly different
|
||||
# purposes. This endpoint is authenticated, whereas the one in PBAFileDownload can not be (at
|
||||
|
|
|
@ -1,6 +1,3 @@
|
|||
from .i_file_storage_service import IFileStorageService, FileRetrievalError
|
||||
from .directory_file_storage_service import DirectoryFileStorageService
|
||||
|
||||
from .authentication.authentication_service import AuthenticationService
|
||||
from .authentication.json_file_user_datastore import JsonFileUserDatastore
|
||||
|
||||
|
|
|
@ -2,7 +2,8 @@ from pathlib import Path
|
|||
|
||||
from common import DIContainer
|
||||
from common.aws import AWSInstance
|
||||
from monkey_island.cc.services import AWSService, DirectoryFileStorageService, IFileStorageService
|
||||
from monkey_island.cc.repository import FilesystemStorage, IFileRepository
|
||||
from monkey_island.cc.services import AWSService
|
||||
from monkey_island.cc.services.post_breach_files import PostBreachFilesService
|
||||
from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
|
||||
|
||||
|
@ -14,13 +15,11 @@ def initialize_services(data_dir: Path) -> DIContainer:
|
|||
container = DIContainer()
|
||||
container.register_instance(AWSInstance, AWSInstance())
|
||||
|
||||
container.register_instance(
|
||||
IFileStorageService, DirectoryFileStorageService(data_dir / "custom_pbas")
|
||||
)
|
||||
container.register_instance(IFileRepository, FilesystemStorage(data_dir / "custom_pbas"))
|
||||
container.register_instance(AWSService, container.resolve(AWSService))
|
||||
|
||||
# This is temporary until we get DI all worked out.
|
||||
PostBreachFilesService.initialize(container.resolve(IFileStorageService))
|
||||
PostBreachFilesService.initialize(container.resolve(IFileRepository))
|
||||
LocalMonkeyRunService.initialize(data_dir)
|
||||
AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir))
|
||||
ReportService.initialize(container.resolve(AWSService))
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
import logging
|
||||
|
||||
from monkey_island.cc.services import IFileStorageService
|
||||
from monkey_island.cc.repository import IFileRepository
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
@ -16,7 +16,7 @@ class PostBreachFilesService:
|
|||
# static/singleton hybrids. At the moment, this requires invasive refactoring that's
|
||||
# not a priority.
|
||||
@classmethod
|
||||
def initialize(cls, file_storage_service: IFileStorageService):
|
||||
def initialize(cls, file_storage_service: IFileRepository):
|
||||
cls._file_storage_service = file_storage_service
|
||||
|
||||
@classmethod
|
||||
|
|
|
@ -5,14 +5,14 @@ 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.resources.pba_file_download import PBAFileDownload
|
||||
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||
|
||||
FILE_NAME = "test_file"
|
||||
FILE_CONTENTS = b"HelloWorld!"
|
||||
|
||||
|
||||
class MockFileStorageService(IFileStorageService):
|
||||
class MockFileStorageService(IFileRepository):
|
||||
def __init__(self):
|
||||
self._file = io.BytesIO(FILE_CONTENTS)
|
||||
|
||||
|
@ -35,7 +35,7 @@ class MockFileStorageService(IFileStorageService):
|
|||
@pytest.fixture
|
||||
def flask_client(build_flask_client):
|
||||
container = StubDIContainer()
|
||||
container.register(IFileStorageService, MockFileStorageService)
|
||||
container.register(IFileRepository, MockFileStorageService)
|
||||
|
||||
with build_flask_client(container) as flask_client:
|
||||
yield flask_client
|
||||
|
|
|
@ -6,8 +6,8 @@ from tests.common import StubDIContainer
|
|||
from tests.unit_tests.monkey_island.conftest import get_url_for_resource
|
||||
from tests.utils import raise_
|
||||
|
||||
from monkey_island.cc.repository import FileRetrievalError, IFileRepository
|
||||
from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload
|
||||
from monkey_island.cc.services import FileRetrievalError, IFileStorageService
|
||||
|
||||
TEST_FILE_CONTENTS = b"m0nk3y"
|
||||
TEST_FILE = (
|
||||
|
@ -40,7 +40,7 @@ def mock_get_config_value(monkeypatch):
|
|||
)
|
||||
|
||||
|
||||
class MockFileStorageService(IFileStorageService):
|
||||
class MockFileStorageService(IFileRepository):
|
||||
def __init__(self):
|
||||
self._file = None
|
||||
|
||||
|
@ -67,7 +67,7 @@ def file_storage_service():
|
|||
@pytest.fixture
|
||||
def flask_client(build_flask_client, file_storage_service):
|
||||
container = StubDIContainer()
|
||||
container.register_instance(IFileStorageService, file_storage_service)
|
||||
container.register_instance(IFileRepository, file_storage_service)
|
||||
|
||||
with build_flask_client(container) as flask_client:
|
||||
yield flask_client
|
||||
|
|
|
@ -4,8 +4,8 @@ from pathlib import Path
|
|||
import pytest
|
||||
from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions
|
||||
|
||||
from monkey_island.cc.repository import FileRetrievalError, FilesystemStorage
|
||||
from monkey_island.cc.server_utils.file_utils import is_windows_os
|
||||
from monkey_island.cc.services import DirectoryFileStorageService, FileRetrievalError
|
||||
|
||||
|
||||
def test_error_if_storage_directory_is_file(tmp_path):
|
||||
|
@ -13,13 +13,13 @@ def test_error_if_storage_directory_is_file(tmp_path):
|
|||
new_file.write_text("HelloWorld!")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
DirectoryFileStorageService(new_file)
|
||||
FilesystemStorage(new_file)
|
||||
|
||||
|
||||
def test_directory_created(tmp_path):
|
||||
new_dir = tmp_path / "new_dir"
|
||||
|
||||
DirectoryFileStorageService(new_dir)
|
||||
FilesystemStorage(new_dir)
|
||||
|
||||
assert new_dir.exists() and new_dir.is_dir()
|
||||
|
||||
|
@ -28,7 +28,7 @@ def test_directory_created(tmp_path):
|
|||
def test_directory_permissions__linux(tmp_path):
|
||||
new_dir = tmp_path / "new_dir"
|
||||
|
||||
DirectoryFileStorageService(new_dir)
|
||||
FilesystemStorage(new_dir)
|
||||
|
||||
assert_linux_permissions(new_dir)
|
||||
|
||||
|
@ -37,7 +37,7 @@ def test_directory_permissions__linux(tmp_path):
|
|||
def test_directory_permissions__windows(tmp_path):
|
||||
new_dir = tmp_path / "new_dir"
|
||||
|
||||
DirectoryFileStorageService(new_dir)
|
||||
FilesystemStorage(new_dir)
|
||||
|
||||
assert_windows_permissions(new_dir)
|
||||
|
||||
|
@ -47,7 +47,7 @@ def save_file(tmp_path, file_path_prefix=""):
|
|||
file_contents = "Hello World!"
|
||||
expected_file_path = tmp_path / file_name
|
||||
|
||||
fss = DirectoryFileStorageService(tmp_path)
|
||||
fss = FilesystemStorage(tmp_path)
|
||||
fss.save_file(Path(file_path_prefix) / file_name, io.BytesIO(file_contents.encode()))
|
||||
|
||||
assert expected_file_path.is_file()
|
||||
|
@ -60,7 +60,7 @@ def delete_file(tmp_path, file_path_prefix=""):
|
|||
file.touch()
|
||||
assert file.is_file()
|
||||
|
||||
fss = DirectoryFileStorageService(tmp_path)
|
||||
fss = FilesystemStorage(tmp_path)
|
||||
fss.delete_file(Path(file_path_prefix) / file_name)
|
||||
|
||||
assert not file.exists()
|
||||
|
@ -72,7 +72,7 @@ def open_file(tmp_path, file_path_prefix=""):
|
|||
expected_file_path = tmp_path / file_name
|
||||
expected_file_path.write_text(expected_file_contents)
|
||||
|
||||
fss = DirectoryFileStorageService(tmp_path)
|
||||
fss = FilesystemStorage(tmp_path)
|
||||
with fss.open_file(Path(file_path_prefix) / file_name) as f:
|
||||
actual_file_contents = f.read()
|
||||
|
||||
|
@ -101,7 +101,7 @@ def test_remove_all_files(tmp_path):
|
|||
for filename in ["1.txt", "2.txt", "3.txt"]:
|
||||
(tmp_path / filename).touch()
|
||||
|
||||
fss = DirectoryFileStorageService(tmp_path)
|
||||
fss = FilesystemStorage(tmp_path)
|
||||
fss.delete_all_files()
|
||||
|
||||
for file in tmp_path.iterdir():
|
||||
|
@ -114,7 +114,7 @@ def test_remove_all_files__skip_directories(tmp_path):
|
|||
for filename in ["1.txt", "2.txt", "3.txt"]:
|
||||
(tmp_path / filename).touch()
|
||||
|
||||
fss = DirectoryFileStorageService(tmp_path)
|
||||
fss = FilesystemStorage(tmp_path)
|
||||
fss.delete_all_files()
|
||||
|
||||
for file in tmp_path.iterdir():
|
||||
|
@ -122,14 +122,14 @@ def test_remove_all_files__skip_directories(tmp_path):
|
|||
|
||||
|
||||
def test_remove_nonexistant_file(tmp_path):
|
||||
fss = DirectoryFileStorageService(tmp_path)
|
||||
fss = FilesystemStorage(tmp_path)
|
||||
|
||||
# This test will fail if this call raises an exception.
|
||||
fss.delete_file("nonexistant_file.txt")
|
||||
|
||||
|
||||
def test_open_nonexistant_file(tmp_path):
|
||||
fss = DirectoryFileStorageService(tmp_path)
|
||||
fss = FilesystemStorage(tmp_path)
|
||||
|
||||
with pytest.raises(FileRetrievalError):
|
||||
fss.open_file("nonexistant_file.txt")
|
||||
|
|
|
@ -4,13 +4,13 @@ import os
|
|||
import pytest
|
||||
from tests.utils import raise_
|
||||
|
||||
from monkey_island.cc.services import DirectoryFileStorageService
|
||||
from monkey_island.cc.repository import FilesystemStorage
|
||||
from monkey_island.cc.services.post_breach_files import PostBreachFilesService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def file_storage_service(tmp_path):
|
||||
return DirectoryFileStorageService(tmp_path)
|
||||
return FilesystemStorage(tmp_path)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
Loading…
Reference in New Issue