Island: Use AgentEventSerializerRegistry in place of EventSerializerRegistry

This commit is contained in:
Shreya Malviya 2022-09-16 12:28:18 +05:30 committed by Mike Salvatore
parent 19a59bf3c2
commit 4a23f5a95b
3 changed files with 9 additions and 7 deletions

View File

@ -2,7 +2,7 @@ from typing import Any, Dict, MutableMapping, Sequence, Type
from pymongo import MongoClient 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.events import AbstractAgentEvent
from common.types import AgentID from common.types import AgentID
from monkey_island.cc.repository import IEventRepository from monkey_island.cc.repository import IEventRepository
@ -14,7 +14,9 @@ from .consts import MONGO_OBJECT_ID_KEY
class MongoEventRepository(IEventRepository): class MongoEventRepository(IEventRepository):
"""A repository for storing and retrieving events in MongoDB""" """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._events_collection = mongo_client.monkey_island.events
self._serializers = serializer_registry self._serializers = serializer_registry

View File

@ -4,7 +4,7 @@ from http import HTTPStatus
from flask import request from flask import request
from common.event_queue import IAgentEventQueue 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 from monkey_island.cc.resources.AbstractResource import AbstractResource
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -16,7 +16,7 @@ class Events(AbstractResource):
def __init__( def __init__(
self, self,
agent_event_queue: IAgentEventQueue, agent_event_queue: IAgentEventQueue,
event_serializer_registry: EventSerializerRegistry, event_serializer_registry: AgentEventSerializerRegistry,
): ):
self._agent_event_queue = agent_event_queue self._agent_event_queue = agent_event_queue
self._event_serializer_registry = event_serializer_registry self._event_serializer_registry = event_serializer_registry

View File

@ -24,7 +24,7 @@ if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path:
from common import DIContainer # noqa: E402 from common import DIContainer # noqa: E402
from common.event_serializers import ( # noqa: E402 from common.event_serializers import ( # noqa: E402
EventSerializerRegistry, AgentEventSerializerRegistry,
register_common_agent_event_serializers, register_common_agent_event_serializers,
) )
from common.version import get_version # noqa: E402 from common.version import get_version # noqa: E402
@ -140,10 +140,10 @@ def _initialize_di_container(
def _setup_agent_event_serializers(container: DIContainer): 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) 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): def _initialize_mongodb_connection(start_mongodb: bool, data_dir: Path):