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

FileAgentLogRepository
This commit is contained in:
Mike Salvatore 2022-09-28 14:11:31 -04:00 committed by GitHub
commit 078574998a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 222 additions and 8 deletions

View File

@ -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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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:

View File

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

View File

@ -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

View File

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

View File

@ -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

View File

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

View File

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

View File

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