Island: Add MongoEventRepository

This commit is contained in:
Kekoa Kaaikala 2022-09-14 15:39:09 +00:00
parent 0b5dd3832b
commit f5e398c175
4 changed files with 213 additions and 0 deletions

View File

@ -27,3 +27,4 @@ from .mongo_machine_repository import MongoMachineRepository
from .mongo_agent_repository import MongoAgentRepository
from .mongo_node_repository import MongoNodeRepository
from .stubbed_event_repository import StubbedEventRepository
from .mongo_event_repository import MongoEventRepository

View File

@ -0,0 +1,52 @@
from typing import Sequence, Type
from pymongo import MongoClient
from common.events import AbstractAgentEvent
from common.types import AgentID
from monkey_island.cc.repository import IEventRepository
from . import RemovalError, RetrievalError, StorageError
class MongoEventRepository(IEventRepository):
def __init__(self, mongo_client: MongoClient):
self._events_collection = mongo_client.monkey_island.events
def save_event(self, event: AbstractAgentEvent):
try:
self._events_collection.insert_one(event.dict(simplify=True))
except Exception as err:
raise StorageError(err)
def get_events(self) -> Sequence[AbstractAgentEvent]:
try:
return list(self._events_collection.find())
except Exception as err:
raise RetrievalError(f"Error retrieving events: {err}")
def get_events_by_type(
self, event_type: Type[AbstractAgentEvent]
) -> Sequence[AbstractAgentEvent]:
try:
return []
except Exception as err:
raise RetrievalError(f"Error retrieving events for type {event_type}: {err}")
def get_events_by_tag(self, tag: str) -> Sequence[AbstractAgentEvent]:
try:
return list(self._events_collection.find({"tags": {"$in": [tag]}}))
except Exception as err:
raise RetrievalError(f"Error retrieving events for tag {tag}: {err}")
def get_events_by_source(self, source: AgentID) -> Sequence[AbstractAgentEvent]:
try:
return list(self._events_collection.find({"source": source}))
except Exception as err:
raise RetrievalError(f"Error retrieving events for source {source}: {err}")
def reset(self):
try:
self._events_collection.drop()
except Exception as err:
raise RemovalError(f"Error resetting the repository: {err}")

View File

@ -0,0 +1,159 @@
import uuid
from typing import List
from unittest.mock import MagicMock
import mongomock
import pytest
from monkey.common.events.abstract_agent_event import AbstractAgentEvent
from monkey.monkey_island.cc.repository import (
IEventRepository,
MongoEventRepository,
RemovalError,
RetrievalError,
StorageError,
)
class FakeAgentEvent(AbstractAgentEvent):
item: str
EVENTS: List[AbstractAgentEvent] = [
AbstractAgentEvent(source=uuid.uuid4(), tags={"foo"}),
AbstractAgentEvent(source=uuid.uuid4(), tags={"foo", "bar"}),
AbstractAgentEvent(source=uuid.uuid4(), tags={"bar", "baz"}),
FakeAgentEvent(source=uuid.uuid4(), tags={"baz"}, item="blah"),
]
@pytest.fixture
def mongo_client():
client = mongomock.MongoClient()
client.monkey_island.events.insert_many((e.dict(simplify=True) for e in EVENTS))
return client
@pytest.fixture
def mongo_repository(mongo_client) -> IEventRepository:
return MongoEventRepository(mongo_client)
@pytest.fixture
def error_raising_mongo_client(mongo_client) -> mongomock.MongoClient:
mongo_client = MagicMock(spec=mongomock.MongoClient)
mongo_client.monkey_island = MagicMock(spec=mongomock.Database)
mongo_client.monkey_island.events = MagicMock(spec=mongomock.Collection)
# The first call to find() must succeed
mongo_client.monkey_island.events.find = MagicMock(
# side_effect=chain([MagicMock()], repeat(Exception("some exception")))
side_effect=Exception("some exception")
)
mongo_client.monkey_island.events.find_one = MagicMock(side_effect=Exception("some exception"))
mongo_client.monkey_island.events.insert_one = MagicMock(
side_effect=Exception("some exception")
)
mongo_client.monkey_island.events.replace_one = MagicMock(
side_effect=Exception("some exception")
)
mongo_client.monkey_island.events.drop = MagicMock(side_effect=Exception("some exception"))
return mongo_client
@pytest.fixture
def error_raising_mongo_repository(error_raising_mongo_client) -> IEventRepository:
return MongoEventRepository(error_raising_mongo_client)
def assert_same_contents(a, b):
assert len(a) == len(b)
difference = set(a) ^ set(b)
assert not difference
def test_mongo_event_repository__save_event(mongo_repository: IEventRepository):
event = AbstractAgentEvent(source=uuid.uuid4())
mongo_repository.save_event(event)
events = mongo_repository.get_events()
assert event in events
def test_mongo_event_repository__save_event_raises(
error_raising_mongo_repository: IEventRepository,
):
event = AbstractAgentEvent(source=uuid.uuid4())
with pytest.raises(StorageError):
error_raising_mongo_repository.save_event(event)
def test_mongo_event_repository__get_events(mongo_repository: IEventRepository):
events = mongo_repository.get_events()
assert_same_contents(events, EVENTS)
def test_mongo_event_repository__get_events_raises(
error_raising_mongo_repository: IEventRepository,
):
with pytest.raises(RetrievalError):
error_raising_mongo_repository.get_events()
def test_mongo_event_repository__get_events_by_type(mongo_repository: IEventRepository):
events = mongo_repository.get_events_by_type(FakeAgentEvent)
expected_events = [EVENTS[3]]
assert_same_contents(events, expected_events)
def test_mongo_event_repository__get_events_by_type_raises(
error_raising_mongo_repository: IEventRepository,
):
with pytest.raises(RetrievalError):
error_raising_mongo_repository.get_events_by_type(FakeAgentEvent)
def test_mongo_event_repository__get_events_by_tag(mongo_repository: IEventRepository):
events = mongo_repository.get_events_by_tag("bar")
expected_events = [EVENTS[1], EVENTS[2]]
assert_same_contents(events, expected_events)
def test_mongo_event_repository__get_events_by_tag_raises(
error_raising_mongo_repository: IEventRepository,
):
with pytest.raises(RetrievalError):
error_raising_mongo_repository.get_events_by_tag("bar")
def test_mongo_event_repository__get_events_by_source(mongo_repository: IEventRepository):
source_event = EVENTS[2]
events = mongo_repository.get_events_by_source(source_event.source)
expected_events = [source_event]
assert_same_contents(events, expected_events)
def test_mongo_event_repository__get_events_by_source_raises(
error_raising_mongo_repository: IEventRepository,
):
with pytest.raises(RetrievalError):
source_event = EVENTS[2]
error_raising_mongo_repository.get_events_by_source(source_event.source)
def test_mongo_event_repository__reset(mongo_repository: IEventRepository):
mongo_repository.reset()
events = mongo_repository.get_events()
assert not events
def test_mongo_event_repository__reset_raises(error_raising_mongo_repository: IEventRepository):
with pytest.raises(RemovalError):
error_raising_mongo_repository.reset()

View File

@ -312,6 +312,7 @@ IEventRepository.save_event
IEventRepository.get_events_by_type
IEventRepository.get_events_by_tag
IEventRepository.get_events_by_source
MongoEventRepository
# pydantic base models