diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 197cae607..9012afb8b 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,4 +1,4 @@ -from .errors import RemovalError, RetrievalError, StorageError, UnknownRecordError +from .errors import RemovalError, RepositoryError, RetrievalError, StorageError, UnknownRecordError from .i_file_repository import FileNotFoundError, IFileRepository @@ -28,5 +28,6 @@ from .mongo_machine_repository import MongoMachineRepository from .mongo_agent_repository import MongoAgentRepository from .mongo_node_repository import MongoNodeRepository from .mongo_agent_event_repository import MongoAgentEventRepository +from .file_agent_log_repository import FileAgentLogRepository from .utils import initialize_machine_repository diff --git a/monkey/monkey_island/cc/repository/errors.py b/monkey/monkey_island/cc/repository/errors.py index a5c26fe15..f01114cc0 100644 --- a/monkey/monkey_island/cc/repository/errors.py +++ b/monkey/monkey_island/cc/repository/errors.py @@ -1,22 +1,28 @@ -class RemovalError(RuntimeError): +class RepositoryError(RuntimeError): + """ + Raised when a repository encounters an error while attempting any operation. + """ + + +class RemovalError(RepositoryError): """ Raised when a repository encounters an error while attempting to remove data. """ -class RetrievalError(RuntimeError): +class RetrievalError(RepositoryError): """ Raised when a repository encounters an error while attempting to retrieve data. """ -class StorageError(RuntimeError): +class StorageError(RepositoryError): """ Raised when a repository encounters an error while attempting to store data. """ -class UnknownRecordError(RuntimeError): +class UnknownRecordError(RepositoryError): """ Raised when the repository does not contain any data matching the request. """ diff --git a/monkey/monkey_island/cc/repository/file_agent_log_repository.py b/monkey/monkey_island/cc/repository/file_agent_log_repository.py new file mode 100644 index 000000000..3ded768dd --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_agent_log_repository.py @@ -0,0 +1,40 @@ +import io +import re + +from monkey_island.cc.models import AgentID +from monkey_island.cc.repository import ( + IAgentLogRepository, + IFileRepository, + RepositoryError, + RetrievalError, +) + +AGENT_LOG_FILE_NAME_REGEX = re.compile(r"^agent-[\w-]+.log$") + + +class FileAgentLogRepository(IAgentLogRepository): + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + + def upsert_agent_log(self, agent_id: AgentID, log_contents: str): + log_file_name = FileAgentLogRepository._get_agent_log_file_name(agent_id) + self._file_repository.save_file(log_file_name, io.BytesIO(log_contents.encode())) + + def get_agent_log(self, agent_id: AgentID) -> str: + log_file_name = FileAgentLogRepository._get_agent_log_file_name(agent_id) + + try: + with self._file_repository.open_file(log_file_name) as f: + log_contents = f.read().decode() + return log_contents + except RepositoryError as err: + raise err + except Exception as err: + raise RetrievalError(f"Error retrieving the agent logs: {err}") + + def reset(self): + self._file_repository.delete_files_by_regex(AGENT_LOG_FILE_NAME_REGEX) + + @staticmethod + def _get_agent_log_file_name(agent_id: AgentID) -> str: + return f"agent-{agent_id}.log" diff --git a/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py index 0e026d035..324ac0f89 100644 --- a/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py @@ -1,4 +1,5 @@ import io +import re import shutil from functools import lru_cache from typing import BinaryIO @@ -36,6 +37,10 @@ class FileRepositoryCachingDecorator(IFileRepository): self._open_file.cache_clear() return self._file_repository.delete_file(unsafe_file_name) + def delete_files_by_regex(self, file_name_regex: re.Pattern): + self._open_file.cache_clear() + return self._file_repository.delete_files_by_regex(file_name_regex) + 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 index 1bd76b4dd..72af1206a 100644 --- a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py @@ -1,3 +1,4 @@ +import re from typing import BinaryIO from readerwriterlock import rwlock @@ -26,6 +27,10 @@ class FileRepositoryLockingDecorator(IFileRepository): with self._rwlock.gen_wlock(): return self._file_repository.delete_file(unsafe_file_name) + def delete_files_by_regex(self, file_name_regex: re.Pattern): + with self._rwlock.gen_wlock(): + return self._file_repository.delete_files_by_regex(file_name_regex) + 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 index 2bf8bfe6e..2ddb31fe1 100644 --- a/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py @@ -1,4 +1,5 @@ import logging +import re from typing import BinaryIO from . import IFileRepository @@ -26,6 +27,10 @@ class FileRepositoryLoggingDecorator(IFileRepository): logger.debug(f"Deleting file {unsafe_file_name}") return self._file_repository.delete_file(unsafe_file_name) + def delete_files_by_regex(self, file_name_regex: re.Pattern): + logger.debug(f'Deleting files whose names match the regex "{file_name_regex}"') + return self._file_repository.delete_files_by_regex(file_name_regex) + 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 bedc9cfdb..f28fc2bbd 100644 --- a/monkey/monkey_island/cc/repository/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/i_file_repository.py @@ -1,10 +1,12 @@ import abc +import re from typing import BinaryIO -from monkey_island.cc.repository import RetrievalError +from monkey_island.cc.repository import UnknownRecordError -class FileNotFoundError(RetrievalError): +# TODO: Remove this and use UnknownRecordError directly wherever needed. +class FileNotFoundError(UnknownRecordError): pass @@ -49,6 +51,17 @@ class IFileRepository(metaclass=abc.ABCMeta): """ pass + @abc.abstractmethod + def delete_files_by_regex(self, file_name_regex: re.Pattern): + """ + Delete files whose names match a particular regex + + This method matches relevant files and deletes them using `delete_file()`. + + :param file_name_regex: A regex with which a file's name should match before deleting it + """ + pass + @abc.abstractmethod def delete_all_files(self): """ 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 58fd20bdb..7f7be398c 100644 --- a/monkey/monkey_island/cc/repository/local_storage_file_repository.py +++ b/monkey/monkey_island/cc/repository/local_storage_file_repository.py @@ -1,4 +1,6 @@ import logging +import os +import re import shutil from pathlib import Path from typing import BinaryIO @@ -55,6 +57,11 @@ class LocalStorageFileRepository(IFileRepository): f'Error retrieving file "{unsafe_file_name}" from the repository: {err}' ) + def delete_files_by_regex(self, file_name_regex: re.Pattern): + for file_name in os.listdir(self._storage_directory): + if re.match(file_name_regex, file_name): + self.delete_file(file_name) + def delete_file(self, unsafe_file_name: str): try: safe_file_path = self._get_safe_file_path(unsafe_file_name) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index c0dc600c3..1cd020fd8 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -21,6 +21,7 @@ from monkey_island.cc.event_queue import IIslandEventQueue, PyPubSubIslandEventQ from monkey_island.cc.repository import ( AgentBinaryRepository, FileAgentConfigurationRepository, + FileAgentLogRepository, FileRepositoryCachingDecorator, FileRepositoryLockingDecorator, FileRepositoryLoggingDecorator, @@ -28,6 +29,7 @@ from monkey_island.cc.repository import ( IAgentBinaryRepository, IAgentConfigurationRepository, IAgentEventRepository, + IAgentLogRepository, IAgentRepository, ICredentialsRepository, IFileRepository, @@ -117,6 +119,7 @@ def _register_repositories(container: DIContainer, data_dir: Path): container.register_instance(INodeRepository, container.resolve(MongoNodeRepository)) container.register_instance(IMachineRepository, _build_machine_repository(container)) container.register_instance(IAgentRepository, container.resolve(MongoAgentRepository)) + container.register_instance(IAgentLogRepository, container.resolve(FileAgentLogRepository)) def _decorate_file_repository(file_repository: IFileRepository) -> IFileRepository: diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 96d2e1941..bf07be1f0 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -10,6 +10,7 @@ from monkey_island.cc.island_event_handlers import ( ) from monkey_island.cc.repository import ( IAgentEventRepository, + IAgentLogRepository, IAgentRepository, ICredentialsRepository, INodeRepository, @@ -59,9 +60,10 @@ def _subscribe_clear_simulation_data_events( island_event_queue.subscribe(topic, container.resolve(reset_machine_repository)) for i_repository in [ - INodeRepository, IAgentEventRepository, + IAgentLogRepository, IAgentRepository, + INodeRepository, ]: repository = container.resolve(i_repository) island_event_queue.subscribe(topic, repository.reset) 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 = {} diff --git a/monkey/tests/monkey_island/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py index 782c9838b..0d6e61e9b 100644 --- a/monkey/tests/monkey_island/mock_file_repository.py +++ b/monkey/tests/monkey_island/mock_file_repository.py @@ -1,4 +1,5 @@ import io +import re from typing import BinaryIO from monkey_island.cc import repository @@ -24,5 +25,8 @@ class MockFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): pass + def delete_files_by_regex(self, file_name_regex: re.Pattern): + pass + def delete_all_files(self): pass diff --git a/monkey/tests/monkey_island/single_file_repository.py b/monkey/tests/monkey_island/single_file_repository.py index 462969acb..be9f896ba 100644 --- a/monkey/tests/monkey_island/single_file_repository.py +++ b/monkey/tests/monkey_island/single_file_repository.py @@ -1,4 +1,5 @@ import io +import re from typing import BinaryIO from monkey_island.cc import repository @@ -8,9 +9,11 @@ from monkey_island.cc.repository import IFileRepository class SingleFileRepository(IFileRepository): def __init__(self): self._file = None + self._file_name = "" def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): self._file = io.BytesIO(file_contents.read()) + self._file_name = unsafe_file_name def open_file(self, unsafe_file_name: str) -> BinaryIO: if self._file is None: @@ -19,6 +22,11 @@ class SingleFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): self._file = None + self._file_name = "" + + def delete_files_by_regex(self, file_name_regex: re.Pattern): + if re.match(file_name_regex, self._file_name): + self.delete_file("") def delete_all_files(self): self.delete_file("") diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py new file mode 100644 index 000000000..71a0f878e --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -0,0 +1,69 @@ +import io +from unittest.mock import MagicMock +from uuid import UUID + +import pytest +from tests.monkey_island import InMemoryFileRepository, OpenErrorFileRepository + +from monkey_island.cc.repository import ( + FileAgentLogRepository, + IAgentLogRepository, + IFileRepository, + RetrievalError, + UnknownRecordError, +) + +LOG_CONTENTS = "lots of useful information" +AGENT_ID_1 = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") +AGENT_ID_2 = UUID("789abcd4-20d7-abcd-ef7a-0123acaabcde") + + +@pytest.fixture +def repository() -> IAgentLogRepository: + return FileAgentLogRepository(InMemoryFileRepository()) + + +def test_store_agent_log(repository: IAgentLogRepository): + repository.upsert_agent_log(AGENT_ID_1, LOG_CONTENTS) + retrieved_log_contents = repository.get_agent_log(AGENT_ID_1) + + assert retrieved_log_contents == LOG_CONTENTS + + +def test_get_agent_log__unknown_record_error(repository: IAgentLogRepository): + with pytest.raises(UnknownRecordError): + repository.get_agent_log(AGENT_ID_1) + + +def test_get_agent_log__retrieval_error(): + repository = FileAgentLogRepository(OpenErrorFileRepository()) + with pytest.raises(RetrievalError): + repository.get_agent_log(AGENT_ID_1) + + +def test_get_agent_log__corrupt_data(): + file_repository = MagicMock(spec=IFileRepository) + # Return invalid unicode + file_repository.open_file = MagicMock(return_value=io.BytesIO(b"\xff\xfe")) + repository = FileAgentLogRepository(file_repository) + + with pytest.raises(RetrievalError): + repository.get_agent_log(AGENT_ID_1) + + +def test_multiple_logs(repository: IAgentLogRepository): + log_contents_1 = "hello" + log_contents_2 = "world" + + repository.upsert_agent_log(AGENT_ID_1, log_contents_1) + repository.upsert_agent_log(AGENT_ID_2, log_contents_2) + + assert repository.get_agent_log(AGENT_ID_1) == log_contents_1 + assert repository.get_agent_log(AGENT_ID_2) == log_contents_2 + + +def test_reset_agent_logs(repository: IAgentLogRepository): + repository.upsert_agent_log(AGENT_ID_1, LOG_CONTENTS) + repository.reset() + with pytest.raises(UnknownRecordError): + repository.get_agent_log(AGENT_ID_1) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py index 64541bcbe..f088ee16e 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py @@ -1,4 +1,5 @@ import io +import re from pathlib import Path from unittest.mock import Mock, patch @@ -143,3 +144,14 @@ def test_open_locked_file(tmp_path, monkeypatch): with patch("builtins.open", Mock(side_effect=Exception())): with pytest.raises(repository.RetrievalError): fss.open_file("locked_file.txt") + + +def test_delete_files_by_regex(tmp_path): + for filename in {"xyz-1.txt", "abc-2.txt", "pqr-3.txt", "abc-4.txt", "abc-5.pdf"}: + (tmp_path / filename).touch() + + fss = LocalStorageFileRepository(tmp_path) + fss.delete_files_by_regex(re.compile(r"^abc-[\w-]+.txt$")) + + files = {f.name for f in tmp_path.iterdir()} + assert files == {"xyz-1.txt", "pqr-3.txt", "abc-5.pdf"}