diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index 6614e165b..c9f1f33d7 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -2,7 +2,7 @@ from typing import Any, Dict, MutableMapping, Sequence, Type from pymongo import MongoClient -from common.event_serializers import EVENT_TYPE_FIELD, EventSerializerRegistry +from common.event_serializers import EVENT_TYPE_FIELD, AgentEventSerializerRegistry from common.events import AbstractAgentEvent from common.types import AgentID from monkey_island.cc.repository import IEventRepository @@ -14,7 +14,9 @@ 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): + def __init__( + self, mongo_client: MongoClient, serializer_registry: AgentEventSerializerRegistry + ): self._events_collection = mongo_client.monkey_island.events self._serializers = serializer_registry diff --git a/monkey/monkey_island/cc/resources/events.py b/monkey/monkey_island/cc/resources/events.py index d1dedeb90..97c42c252 100644 --- a/monkey/monkey_island/cc/resources/events.py +++ b/monkey/monkey_island/cc/resources/events.py @@ -4,7 +4,7 @@ from http import HTTPStatus from flask import request from common.event_queue import IAgentEventQueue -from common.event_serializers import EVENT_TYPE_FIELD, EventSerializerRegistry +from common.event_serializers import EVENT_TYPE_FIELD, AgentEventSerializerRegistry from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__name__) @@ -16,7 +16,7 @@ class Events(AbstractResource): def __init__( self, agent_event_queue: IAgentEventQueue, - event_serializer_registry: EventSerializerRegistry, + event_serializer_registry: AgentEventSerializerRegistry, ): self._agent_event_queue = agent_event_queue self._event_serializer_registry = event_serializer_registry diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 4f6989e12..27fc3bd77 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -24,7 +24,7 @@ if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path: from common import DIContainer # noqa: E402 from common.event_serializers import ( # noqa: E402 - EventSerializerRegistry, + AgentEventSerializerRegistry, register_common_agent_event_serializers, ) from common.version import get_version # noqa: E402 @@ -140,10 +140,10 @@ def _initialize_di_container( def _setup_agent_event_serializers(container: DIContainer): - agent_event_serializer_registry = EventSerializerRegistry() + agent_event_serializer_registry = AgentEventSerializerRegistry() register_common_agent_event_serializers(agent_event_serializer_registry) - container.register_instance(EventSerializerRegistry, agent_event_serializer_registry) + container.register_instance(AgentEventSerializerRegistry, agent_event_serializer_registry) def _initialize_mongodb_connection(start_mongodb: bool, data_dir: Path):