diff --git a/monkey/monkey_island/Pipfile b/monkey/monkey_island/Pipfile index fbc8b16ae..052a46039 100644 --- a/monkey/monkey_island/Pipfile +++ b/monkey/monkey_island/Pipfile @@ -30,6 +30,7 @@ pywin32 = {version = "*", sys_platform = "== 'win32'"} # Lock file is not create pefile = {version = "*", sys_platform = "== 'win32'"} # Pyinstaller requirement on windows marshmallow = "*" marshmallow-enum = "*" +readerwriterlock = "*" [dev-packages] virtualenv = ">=20.0.26" diff --git a/monkey/monkey_island/Pipfile.lock b/monkey/monkey_island/Pipfile.lock index c3e59c1d5..67647710a 100644 --- a/monkey/monkey_island/Pipfile.lock +++ b/monkey/monkey_island/Pipfile.lock @@ -1,7 +1,7 @@ { "_meta": { "hash": { - "sha256": "04efa1f593acdfcdc48b7089108921a46421acbacec80d8a664ec674b221dd4b" + "sha256": "91b8cfcf1408b3709300f47d420c550fe355df76ad396e455049fef1cceca3ad" }, "pipfile-spec": 6, "requires": { @@ -338,7 +338,7 @@ "sha256:84d9dd047ffa80596e0f246e2eab0b391788b0503584e8945f2368256d2735ff", "sha256:9d643ff0a55b762d5cdb124b8eaa99c66322e2157b69160bc32796e824360e6d" ], - "markers": "python_full_version >= '3.5.0'", + "markers": "python_version >= '3.5'", "version": "==3.3" }, "importlib-metadata": { @@ -792,6 +792,14 @@ "markers": "sys_platform == 'win32'", "version": "==0.2.0" }, + "readerwriterlock": { + "hashes": [ + "sha256:8c4b704e60d15991462081a27ef46762fea49b478aa4426644f2146754759ca7", + "sha256:b7c4cc003435d7a8ff15b312b0a62a88d9800ba6164af88991f87f8b748f9bea" + ], + "index": "pypi", + "version": "==1.0.9" + }, "requests": { "hashes": [ "sha256:bc7861137fbce630f17b03d3ad02ad0bf978c844f3536d0edda6499dafce2b6f", @@ -1117,7 +1125,7 @@ "sha256:84d9dd047ffa80596e0f246e2eab0b391788b0503584e8945f2368256d2735ff", "sha256:9d643ff0a55b762d5cdb124b8eaa99c66322e2157b69160bc32796e824360e6d" ], - "markers": "python_full_version >= '3.5.0'", + "markers": "python_version >= '3.5'", "version": "==3.3" }, "imagesize": { diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 153d57e27..d90f6aac6 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,6 +1,9 @@ from .errors import RemovalError, RetrievalError, StorageError from .i_file_repository import FileNotFoundError, IFileRepository from .local_storage_file_repository import LocalStorageFileRepository +from .file_repository_caching_decorator import FileRepositoryCachingDecorator +from .file_repository_locking_decorator import FileRepositoryLockingDecorator +from .file_repository_logging_decorator import FileRepositoryLoggingDecorator from .i_agent_binary_repository import IAgentBinaryRepository from .agent_binary_repository import AgentBinaryRepository from .i_agent_configuration_repository import IAgentConfigurationRepository diff --git a/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py new file mode 100644 index 000000000..0e026d035 --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py @@ -0,0 +1,41 @@ +import io +import shutil +from functools import lru_cache +from typing import BinaryIO + +from . import IFileRepository + + +class FileRepositoryCachingDecorator(IFileRepository): + """ + An IFileRepository decorator that provides caching for other IFileRepositories. + """ + + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + self._open_file.cache_clear() + return self._file_repository.save_file(unsafe_file_name, file_contents) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + original_file = self._open_file(unsafe_file_name) + file_copy = io.BytesIO() + + shutil.copyfileobj(original_file, file_copy) + original_file.seek(0) + file_copy.seek(0) + + return file_copy + + @lru_cache(maxsize=16) + def _open_file(self, unsafe_file_name: str) -> BinaryIO: + return self._file_repository.open_file(unsafe_file_name) + + def delete_file(self, unsafe_file_name: str): + self._open_file.cache_clear() + return self._file_repository.delete_file(unsafe_file_name) + + def delete_all_files(self): + self._open_file.cache_clear() + return self._file_repository.delete_all_files() diff --git a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py new file mode 100644 index 000000000..1bd76b4dd --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py @@ -0,0 +1,31 @@ +from typing import BinaryIO + +from readerwriterlock import rwlock + +from . import IFileRepository + + +class FileRepositoryLockingDecorator(IFileRepository): + """ + An IFileRepository decorator that makes other IFileRepositories thread-safe. + """ + + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + self._rwlock = rwlock.RWLockFair() + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + with self._rwlock.gen_wlock(): + return self._file_repository.save_file(unsafe_file_name, file_contents) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + with self._rwlock.gen_rlock(): + return self._file_repository.open_file(unsafe_file_name) + + def delete_file(self, unsafe_file_name: str): + with self._rwlock.gen_wlock(): + return self._file_repository.delete_file(unsafe_file_name) + + def delete_all_files(self): + with self._rwlock.gen_wlock(): + return self._file_repository.delete_all_files() diff --git a/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py b/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py new file mode 100644 index 000000000..2bf8bfe6e --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py @@ -0,0 +1,31 @@ +import logging +from typing import BinaryIO + +from . import IFileRepository + +logger = logging.getLogger(__name__) + + +class FileRepositoryLoggingDecorator(IFileRepository): + """ + An IFileRepository decorator that provides debug logging for other IFileRepositories. + """ + + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + logger.debug(f"Saving file {unsafe_file_name}") + return self._file_repository.save_file(unsafe_file_name, file_contents) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + logger.debug(f"Opening file {unsafe_file_name}") + return self._file_repository.open_file(unsafe_file_name) + + def delete_file(self, unsafe_file_name: str): + logger.debug(f"Deleting file {unsafe_file_name}") + return self._file_repository.delete_file(unsafe_file_name) + + def delete_all_files(self): + logger.debug("Deleting all files in the repository") + return self._file_repository.delete_all_files() diff --git a/monkey/monkey_island/cc/repository/i_file_repository.py b/monkey/monkey_island/cc/repository/i_file_repository.py index baa52fdda..bedc9cfdb 100644 --- a/monkey/monkey_island/cc/repository/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/i_file_repository.py @@ -27,7 +27,7 @@ class IFileRepository(metaclass=abc.ABCMeta): @abc.abstractmethod def open_file(self, unsafe_file_name: str) -> BinaryIO: """ - Open a file and return a file-like object + Open a file and return a read-only file-like object :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 diff --git a/monkey/monkey_island/cc/repository/local_storage_file_repository.py b/monkey/monkey_island/cc/repository/local_storage_file_repository.py index 4c7662b0c..58fd20bdb 100644 --- a/monkey/monkey_island/cc/repository/local_storage_file_repository.py +++ b/monkey/monkey_island/cc/repository/local_storage_file_repository.py @@ -4,10 +4,11 @@ from pathlib import Path from typing import BinaryIO from common.utils.file_utils import get_all_regular_files_in_directory +from monkey_island.cc import repository from monkey_island.cc.repository import RemovalError, RetrievalError, StorageError from monkey_island.cc.server_utils.file_utils import create_secure_directory -from . import IFileRepository, i_file_repository +from . import IFileRepository logger = logging.getLogger(__name__) @@ -35,7 +36,6 @@ class LocalStorageFileRepository(IFileRepository): try: safe_file_path = self._get_safe_file_path(unsafe_file_name) - logger.debug(f"Saving file contents to {safe_file_path}") with open(safe_file_path, "wb") as dest: shutil.copyfileobj(file_contents, dest) except Exception as err: @@ -44,12 +44,10 @@ class LocalStorageFileRepository(IFileRepository): def open_file(self, unsafe_file_name: str) -> BinaryIO: try: safe_file_path = self._get_safe_file_path(unsafe_file_name) - - logger.debug(f"Opening {safe_file_path}") return open(safe_file_path, "rb") except FileNotFoundError as err: # Wrap Python's FileNotFound error, which is-an OSError, in repository.FileNotFoundError - raise i_file_repository.FileNotFoundError( + raise repository.FileNotFoundError( f'The requested file "{unsafe_file_name}" does not exist: {err}' ) except Exception as err: @@ -60,8 +58,6 @@ class LocalStorageFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): try: safe_file_path = self._get_safe_file_path(unsafe_file_name) - - logger.debug(f"Deleting {safe_file_path}") safe_file_path.unlink() except FileNotFoundError: # This method is idempotent. diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 28a88558c..b5a16ec9f 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -12,6 +12,9 @@ from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, FileAgentConfigurationRepository, + FileRepositoryCachingDecorator, + FileRepositoryLockingDecorator, + FileRepositoryLoggingDecorator, FileSimulationRepository, IAgentBinaryRepository, IAgentConfigurationRepository, @@ -62,7 +65,8 @@ def _register_conventions(container: DIContainer, data_dir: Path): def _register_repositories(container: DIContainer, data_dir: Path): container.register_instance( - IFileRepository, LocalStorageFileRepository(data_dir / "runtime_data") + IFileRepository, + _decorate_file_repository(LocalStorageFileRepository(data_dir / "runtime_data")), ) container.register_instance(IAgentBinaryRepository, _build_agent_binary_repository()) container.register_instance( @@ -71,8 +75,14 @@ def _register_repositories(container: DIContainer, data_dir: Path): container.register_instance(ISimulationRepository, container.resolve(FileSimulationRepository)) +def _decorate_file_repository(file_repository: IFileRepository) -> IFileRepository: + return FileRepositoryLockingDecorator( + FileRepositoryLoggingDecorator(FileRepositoryCachingDecorator(file_repository)) + ) + + def _build_agent_binary_repository(): - file_repository = LocalStorageFileRepository(AGENT_BINARIES_PATH) + file_repository = _decorate_file_repository(LocalStorageFileRepository(AGENT_BINARIES_PATH)) agent_binary_repository = AgentBinaryRepository(file_repository) _log_agent_binary_hashes(agent_binary_repository) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_repository_caching_decorator.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_repository_caching_decorator.py new file mode 100644 index 000000000..14fdf8c27 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_repository_caching_decorator.py @@ -0,0 +1,55 @@ +import io + +import pytest +from tests.monkey_island import SingleFileRepository + +from monkey_island.cc import repository +from monkey_island.cc.repository import FileRepositoryCachingDecorator + + +@pytest.fixture +def file_repository(): + return FileRepositoryCachingDecorator(SingleFileRepository()) + + +def test_open_cache_file(file_repository): + file_name = "test.txt" + file_contents = b"Hello World!" + + file_repository.save_file(file_name, io.BytesIO(file_contents)) + assert file_repository.open_file(file_name).read() == file_contents + assert file_repository.open_file(file_name).read() == file_contents + + +def test_overwrite_file(file_repository): + file_name = "test.txt" + file_contents_1 = b"Hello World!" + file_contents_2 = b"Goodbye World!" + + file_repository.save_file(file_name, io.BytesIO(file_contents_1)) + assert file_repository.open_file(file_name).read() == file_contents_1 + + file_repository.save_file(file_name, io.BytesIO(file_contents_2)) + assert file_repository.open_file(file_name).read() == file_contents_2 + + +def test_delete_file(file_repository): + file_name = "test.txt" + file_contents = b"Hello World!" + + file_repository.save_file(file_name, io.BytesIO(file_contents)) + file_repository.delete_file(file_name) + + with pytest.raises(repository.FileNotFoundError): + file_repository.open_file(file_name) + + +def test_delete_all_files(file_repository): + file_name = "test.txt" + file_contents = b"Hello World!" + + file_repository.save_file(file_name, io.BytesIO(file_contents)) + file_repository.delete_all_files() + + with pytest.raises(repository.FileNotFoundError): + file_repository.open_file(file_name)