diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py index f7033842d..9dd1996f9 100644 --- a/monkey/tests/monkey_island/__init__.py +++ b/monkey/tests/monkey_island/__init__.py @@ -4,3 +4,4 @@ from .open_error_file_repository import OpenErrorFileRepository from .in_memory_agent_configuration_repository import InMemoryAgentConfigurationRepository from .in_memory_simulation_configuration import InMemorySimulationRepository from .in_memory_credentials_repository import InMemoryCredentialsRepository +from .in_memory_file_repository import InMemoryFileRepository diff --git a/monkey/tests/monkey_island/in_memory_file_repository.py b/monkey/tests/monkey_island/in_memory_file_repository.py new file mode 100644 index 000000000..76778d226 --- /dev/null +++ b/monkey/tests/monkey_island/in_memory_file_repository.py @@ -0,0 +1,33 @@ +import io +import re +from typing import BinaryIO, Dict + +from common.utils.code_utils import del_key +from monkey_island.cc.repository import IFileRepository, UnknownRecordError + + +class InMemoryFileRepository(IFileRepository): + def __init__(self): + self._files: Dict[str, bytes] = {} + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + self._files[unsafe_file_name] = file_contents.read() + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + try: + return io.BytesIO(self._files[unsafe_file_name]) + except KeyError: + raise UnknownRecordError(f"Unknown file {unsafe_file_name}") + + def delete_file(self, unsafe_file_name: str): + del_key(self._files, "unsafe_file_name") + + def delete_files_by_regex(self, file_name_regex: re.Pattern): + self._files = { + name: contents + for name, contents in self._files.items() + if not re.match(file_name_regex, name) + } + + def delete_all_files(self): + self._files = {}