Merge pull request #2061 from guardicore/file-repository-decorators

File repository decorators

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"

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": {

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

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()

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()

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()

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

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.

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)

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)