From e3b5e0c01aebe68493253a9988d8b0b97b5ee07f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Sat, 2 Jul 2022 09:53:18 -0400 Subject: [PATCH] Island: Add FileRepositoryLockingDecorator --- .../monkey_island/cc/repository/__init__.py | 1 + .../file_repository_locking_decorator.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+) create mode 100644 monkey/monkey_island/cc/repository/file_repository_locking_decorator.py diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 3b18883c2..6f683cc1a 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,6 +1,7 @@ from .errors import RemovalError, RetrievalError, StorageError from .i_file_repository import FileNotFoundError, IFileRepository from .local_storage_file_repository import LocalStorageFileRepository +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 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..7539d4fc1 --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py @@ -0,0 +1,26 @@ +from multiprocessing import Lock +from typing import BinaryIO + +from . import IFileRepository + + +class FileRepositoryLockingDecorator(IFileRepository): + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + self._lock = Lock() + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + with self._lock: + return self._file_repository.save_file(unsafe_file_name, file_contents) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + with self._lock: + return self._file_repository.open_file(unsafe_file_name) + + def delete_file(self, unsafe_file_name: str): + with self._lock: + return self._file_repository.delete_file(unsafe_file_name) + + def delete_all_files(self): + with self._lock: + return self._file_repository.delete_all_files()