Merge pull request #2363 from guardicore/2274-file-agent-log-repo

FileAgentLogRepository

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

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.
     """

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"

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

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

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

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

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)

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:

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)

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

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 = {}

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

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

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)

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"}