From f5e398c1759922469af5c0d3c04b184a94269fc2 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 15:39:09 +0000 Subject: [PATCH 01/10] Island: Add MongoEventRepository --- .../monkey_island/cc/repository/__init__.py | 1 + .../cc/repository/mongo_event_repository.py | 52 ++++++ .../repository/test_mongo_event_repository.py | 159 ++++++++++++++++++ vulture_allowlist.py | 1 + 4 files changed, 213 insertions(+) create mode 100644 monkey/monkey_island/cc/repository/mongo_event_repository.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index a414b5eb0..63e9ac214 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -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 diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py new file mode 100644 index 000000000..6708e3341 --- /dev/null +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -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}") diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py new file mode 100644 index 000000000..cedd43005 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py @@ -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() diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 6aa14abad..23ec0f252 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -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 From ad5bba9e2fb4994f42dadf9d6e2d45528bbae6a2 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 17:20:38 +0000 Subject: [PATCH 02/10] Island: Use EventSerializerRegistry --- .../cc/repository/mongo_event_repository.py | 29 +++++++--- .../repository/test_mongo_event_repository.py | 54 ++++++++++++------- 2 files changed, 58 insertions(+), 25 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index 6708e3341..772c0a3fb 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -1,27 +1,33 @@ -from typing import Sequence, Type +from typing import Any, MutableMapping, Sequence, Type from pymongo import MongoClient +from common.event_serializers import EVENT_TYPE_FIELD, EventSerializerRegistry from common.events import AbstractAgentEvent from common.types import AgentID from monkey_island.cc.repository import IEventRepository from . import RemovalError, RetrievalError, StorageError +from .consts import MONGO_OBJECT_ID_KEY class MongoEventRepository(IEventRepository): - def __init__(self, mongo_client: MongoClient): + def __init__(self, mongo_client: MongoClient, serializer_registry: EventSerializerRegistry): self._events_collection = mongo_client.monkey_island.events + self._serializers = serializer_registry def save_event(self, event: AbstractAgentEvent): try: - self._events_collection.insert_one(event.dict(simplify=True)) + serializer = self._serializers[type(event)] + serialized_event = serializer.serialize(event) + self._events_collection.insert_one(serialized_event) except Exception as err: raise StorageError(err) def get_events(self) -> Sequence[AbstractAgentEvent]: try: - return list(self._events_collection.find()) + serialized_events = list(self._events_collection.find()) + return list(map(self._deserialize, serialized_events)) except Exception as err: raise RetrievalError(f"Error retrieving events: {err}") @@ -29,19 +35,22 @@ class MongoEventRepository(IEventRepository): self, event_type: Type[AbstractAgentEvent] ) -> Sequence[AbstractAgentEvent]: try: - return [] + serialized_events = list(self._events_collection.find({EVENT_TYPE_FIELD: event_type})) + return list(map(self._deserialize, serialized_events)) 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]}})) + serialized_events = list(self._events_collection.find({"tags": {"$in": [tag]}})) + return list(map(self._deserialize, serialized_events)) 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})) + serialized_events = list(self._events_collection.find({"source": source})) + return list(map(self._deserialize, serialized_events)) except Exception as err: raise RetrievalError(f"Error retrieving events for source {source}: {err}") @@ -50,3 +59,9 @@ class MongoEventRepository(IEventRepository): self._events_collection.drop() except Exception as err: raise RemovalError(f"Error resetting the repository: {err}") + + def _deserialize(self, mongo_record: MutableMapping[str, Any]) -> AbstractAgentEvent: + del mongo_record[MONGO_OBJECT_ID_KEY] + event_type = mongo_record[EVENT_TYPE_FIELD] + serializer = self._serializers[event_type] + return serializer.deserialize(mongo_record) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py index cedd43005..ac5a02f34 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py @@ -4,9 +4,11 @@ from unittest.mock import MagicMock import mongomock import pytest +from pydantic import Field -from monkey.common.events.abstract_agent_event import AbstractAgentEvent -from monkey.monkey_island.cc.repository import ( +from common.event_serializers import EventSerializerRegistry, PydanticEventSerializer +from common.events import AbstractAgentEvent +from monkey_island.cc.repository import ( IEventRepository, MongoEventRepository, RemovalError, @@ -16,27 +18,41 @@ from monkey.monkey_island.cc.repository import ( class FakeAgentEvent(AbstractAgentEvent): + data = Field(default=435) + + +class FakeAgentItemEvent(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"), + FakeAgentEvent(source=uuid.uuid4(), tags={"foo"}), + FakeAgentEvent(source=uuid.uuid4(), tags={"foo", "bar"}), + FakeAgentEvent(source=uuid.uuid4(), tags={"bar", "baz"}), + FakeAgentItemEvent(source=uuid.uuid4(), tags={"baz"}, item="blah"), ] @pytest.fixture -def mongo_client(): +def event_serializer_registry() -> EventSerializerRegistry: + registry = EventSerializerRegistry() + registry[FakeAgentEvent] = PydanticEventSerializer(FakeAgentEvent) + registry[FakeAgentItemEvent] = PydanticEventSerializer(FakeAgentItemEvent) + return registry + + +@pytest.fixture +def mongo_client(event_serializer_registry): client = mongomock.MongoClient() - client.monkey_island.events.insert_many((e.dict(simplify=True) for e in EVENTS)) + client.monkey_island.events.insert_many( + (event_serializer_registry[type(e)].serialize(e) for e in EVENTS) + ) return client @pytest.fixture -def mongo_repository(mongo_client) -> IEventRepository: - return MongoEventRepository(mongo_client) +def mongo_repository(mongo_client, event_serializer_registry) -> IEventRepository: + return MongoEventRepository(mongo_client, event_serializer_registry) @pytest.fixture @@ -63,18 +79,20 @@ def error_raising_mongo_client(mongo_client) -> mongomock.MongoClient: @pytest.fixture -def error_raising_mongo_repository(error_raising_mongo_client) -> IEventRepository: - return MongoEventRepository(error_raising_mongo_client) +def error_raising_mongo_repository( + error_raising_mongo_client, event_serializer_registry +) -> IEventRepository: + return MongoEventRepository(error_raising_mongo_client, event_serializer_registry) def assert_same_contents(a, b): assert len(a) == len(b) - difference = set(a) ^ set(b) - assert not difference + for item in a: + assert item in b def test_mongo_event_repository__save_event(mongo_repository: IEventRepository): - event = AbstractAgentEvent(source=uuid.uuid4()) + event = FakeAgentEvent(source=uuid.uuid4()) mongo_repository.save_event(event) events = mongo_repository.get_events() @@ -84,7 +102,7 @@ def test_mongo_event_repository__save_event(mongo_repository: IEventRepository): def test_mongo_event_repository__save_event_raises( error_raising_mongo_repository: IEventRepository, ): - event = AbstractAgentEvent(source=uuid.uuid4()) + event = FakeAgentEvent(source=uuid.uuid4()) with pytest.raises(StorageError): error_raising_mongo_repository.save_event(event) @@ -104,7 +122,7 @@ def test_mongo_event_repository__get_events_raises( def test_mongo_event_repository__get_events_by_type(mongo_repository: IEventRepository): - events = mongo_repository.get_events_by_type(FakeAgentEvent) + events = mongo_repository.get_events_by_type(FakeAgentItemEvent) expected_events = [EVENTS[3]] assert_same_contents(events, expected_events) @@ -114,7 +132,7 @@ 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) + error_raising_mongo_repository.get_events_by_type(FakeAgentItemEvent) def test_mongo_event_repository__get_events_by_tag(mongo_repository: IEventRepository): From 89954223e8c0260082ad8e7f378dd8c3d969aefc Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 17:29:22 +0000 Subject: [PATCH 03/10] Island: Match type based on __name__ --- monkey/monkey_island/cc/repository/mongo_event_repository.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index 772c0a3fb..2a2d1451e 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -35,7 +35,9 @@ class MongoEventRepository(IEventRepository): self, event_type: Type[AbstractAgentEvent] ) -> Sequence[AbstractAgentEvent]: try: - serialized_events = list(self._events_collection.find({EVENT_TYPE_FIELD: event_type})) + serialized_events = list( + self._events_collection.find({EVENT_TYPE_FIELD: event_type.__name__}) + ) return list(map(self._deserialize, serialized_events)) except Exception as err: raise RetrievalError(f"Error retrieving events for type {event_type}: {err}") From 5c762930bf8fad3a3a4872c5abedd3984e27cfe3 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 17:34:31 +0000 Subject: [PATCH 04/10] Island: Search UUID by string value --- monkey/monkey_island/cc/repository/mongo_event_repository.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index 2a2d1451e..2bf793c68 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -51,7 +51,7 @@ class MongoEventRepository(IEventRepository): def get_events_by_source(self, source: AgentID) -> Sequence[AbstractAgentEvent]: try: - serialized_events = list(self._events_collection.find({"source": source})) + serialized_events = list(self._events_collection.find({"source": str(source)})) return list(map(self._deserialize, serialized_events)) except Exception as err: raise RetrievalError(f"Error retrieving events for source {source}: {err}") From 7e96c11204bdc87b837d6f2e6a43fa2dd370ca55 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 17:39:17 +0000 Subject: [PATCH 05/10] Island: Add docstring to MongoEventRepository --- monkey/monkey_island/cc/repository/mongo_event_repository.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index 2bf793c68..f95bfbdd4 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -12,6 +12,8 @@ from .consts import MONGO_OBJECT_ID_KEY class MongoEventRepository(IEventRepository): + """A repository for storing and retrieving events in MongoDB""" + def __init__(self, mongo_client: MongoClient, serializer_registry: EventSerializerRegistry): self._events_collection = mongo_client.monkey_island.events self._serializers = serializer_registry From 5f44204583ebb9fc9c1d16cce192acf07a935216 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 18:38:56 +0000 Subject: [PATCH 06/10] UT: Clean up mongo fixture --- .../cc/repository/test_mongo_event_repository.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py index ac5a02f34..4e9012443 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py @@ -61,18 +61,10 @@ def error_raising_mongo_client(mongo_client) -> 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.find = 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 From 8311fc79a95a5082a533d9dc3105e23d1fc80b4c Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 18:43:52 +0000 Subject: [PATCH 07/10] Island: Add message to error on save --- monkey/monkey_island/cc/repository/mongo_event_repository.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index f95bfbdd4..dbf080a84 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -24,7 +24,7 @@ class MongoEventRepository(IEventRepository): serialized_event = serializer.serialize(event) self._events_collection.insert_one(serialized_event) except Exception as err: - raise StorageError(err) + raise StorageError(f"Error saving event: {err}") def get_events(self) -> Sequence[AbstractAgentEvent]: try: From cd84280e7924785dbf135f81f102ce52c7b1ad99 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 19:14:35 +0000 Subject: [PATCH 08/10] Island: Simplify event query and deserialization --- .../cc/repository/mongo_event_repository.py | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index dbf080a84..fe85b3216 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -1,4 +1,4 @@ -from typing import Any, MutableMapping, Sequence, Type +from typing import Any, Dict, MutableMapping, Sequence, Type from pymongo import MongoClient @@ -28,8 +28,7 @@ class MongoEventRepository(IEventRepository): def get_events(self) -> Sequence[AbstractAgentEvent]: try: - serialized_events = list(self._events_collection.find()) - return list(map(self._deserialize, serialized_events)) + return self._query_events({}) except Exception as err: raise RetrievalError(f"Error retrieving events: {err}") @@ -37,24 +36,19 @@ class MongoEventRepository(IEventRepository): self, event_type: Type[AbstractAgentEvent] ) -> Sequence[AbstractAgentEvent]: try: - serialized_events = list( - self._events_collection.find({EVENT_TYPE_FIELD: event_type.__name__}) - ) - return list(map(self._deserialize, serialized_events)) + return self._query_events({EVENT_TYPE_FIELD: event_type.__name__}) 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: - serialized_events = list(self._events_collection.find({"tags": {"$in": [tag]}})) - return list(map(self._deserialize, serialized_events)) + return self._query_events({"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: - serialized_events = list(self._events_collection.find({"source": str(source)})) - return list(map(self._deserialize, serialized_events)) + return self._query_events({"source": str(source)}) except Exception as err: raise RetrievalError(f"Error retrieving events for source {source}: {err}") @@ -65,7 +59,10 @@ class MongoEventRepository(IEventRepository): raise RemovalError(f"Error resetting the repository: {err}") def _deserialize(self, mongo_record: MutableMapping[str, Any]) -> AbstractAgentEvent: - del mongo_record[MONGO_OBJECT_ID_KEY] event_type = mongo_record[EVENT_TYPE_FIELD] serializer = self._serializers[event_type] return serializer.deserialize(mongo_record) + + def _query_events(self, query: Dict[Any, Any]) -> Sequence[AbstractAgentEvent]: + serialized_events = list(self._events_collection.find(query, {MONGO_OBJECT_ID_KEY: False})) + return list(map(self._deserialize, serialized_events)) From 9ec91ec97afc4fbf53e31ecc2afc4652c55012a0 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 19:35:05 +0000 Subject: [PATCH 09/10] Island: Remove unneccessary call to list() --- monkey/monkey_island/cc/repository/mongo_event_repository.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index fe85b3216..6614e165b 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -64,5 +64,5 @@ class MongoEventRepository(IEventRepository): return serializer.deserialize(mongo_record) def _query_events(self, query: Dict[Any, Any]) -> Sequence[AbstractAgentEvent]: - serialized_events = list(self._events_collection.find(query, {MONGO_OBJECT_ID_KEY: False})) + serialized_events = self._events_collection.find(query, {MONGO_OBJECT_ID_KEY: False}) return list(map(self._deserialize, serialized_events)) From 25c3552061eac400da3f3c9b801cfc9c0923ac37 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 14 Sep 2022 19:35:37 +0000 Subject: [PATCH 10/10] UT: Clarify mongo repository reset test --- .../monkey_island/cc/repository/test_mongo_event_repository.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py index 4e9012443..64de7edc8 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py @@ -158,6 +158,9 @@ def test_mongo_event_repository__get_events_by_source_raises( def test_mongo_event_repository__reset(mongo_repository: IEventRepository): + initial_events = mongo_repository.get_events() + assert initial_events + mongo_repository.reset() events = mongo_repository.get_events()