Island: Add MongoEventRepository to DI container

monkey/monkey_island/cc/server_setup.py

@@ -23,11 +23,6 @@ if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path:
     sys.path.insert(0, MONKEY_ISLAND_DIR_BASE_PATH)
 
 from common import DIContainer  # noqa: E402
-from common.agent_event_serializers import (  # noqa: E402
-    AgentEventSerializerRegistry,
-    register_common_agent_event_serializers,
-)
-from common.network.network_utils import get_my_ip_addresses  # noqa: E402
 from common.version import get_version  # noqa: E402
 from monkey_island.cc.app import init_app  # noqa: E402
 from monkey_island.cc.arg_parser import IslandCmdArgs  # noqa: E402
@@ -45,6 +40,7 @@ from monkey_island.cc.setup import (  # noqa: E402
     setup_agent_event_handlers,
     setup_island_event_handlers,
 )
+from common.network.network_utils import get_my_ip_addresses  # noqa: E402
 from monkey_island.cc.setup.data_dir import IncompatibleDataDirectory, setup_data_dir  # noqa: E402
 from monkey_island.cc.setup.gevent_hub_error_handler import GeventHubErrorHandler  # noqa: E402
 from monkey_island.cc.setup.island_config_options import IslandConfigOptions  # noqa: E402
@@ -71,7 +67,6 @@ def run_monkey_island():
     container = _initialize_di_container(ip_addresses, version, config_options.data_dir)
     setup_island_event_handlers(container)
     setup_agent_event_handlers(container)
-    _setup_agent_event_serializers(container)
 
     _start_island_server(ip_addresses, island_args.setup_only, config_options, container)
 
@@ -141,13 +136,6 @@ def _initialize_di_container(
     return container
 
 
-def _setup_agent_event_serializers(container: DIContainer):
-    agent_event_serializer_registry = AgentEventSerializerRegistry()
-    register_common_agent_event_serializers(agent_event_serializer_registry)
-
-    container.register_instance(AgentEventSerializerRegistry, agent_event_serializer_registry)
-
-
 def _initialize_mongodb_connection(start_mongodb: bool, data_dir: Path):
     mongo_db_process = None
     if start_mongodb:

monkey/monkey_island/cc/services/initialize.py

@@ -12,6 +12,10 @@ from common.agent_configuration import (
 )
 from common.aws import AWSInstance
 from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue
+from common.agent_event_serializers import (
+    AgentEventSerializerRegistry,
+    register_common_agent_event_serializers,
+)
 from common.utils.file_utils import get_binary_io_sha256_hash
 from monkey_island.cc.event_queue import IIslandEventQueue, PyPubSubIslandEventQueue
 from monkey_island.cc.repository import (
@@ -35,10 +39,10 @@ from monkey_island.cc.repository import (
     LocalStorageFileRepository,
     MongoAgentRepository,
     MongoCredentialsRepository,
+    MongoAgentEventRepository,
     MongoMachineRepository,
     MongoNodeRepository,
     RetrievalError,
-    StubbedEventRepository,
 )
 from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
 from monkey_island.cc.server_utils.encryption import ILockableEncryptor, RepositoryEncryptor
@@ -68,6 +72,7 @@ def initialize_services(container: DIContainer, data_dir: Path):
     container.register_instance(IAgentEventQueue, container.resolve(PyPubSubAgentEventQueue))
     container.register_instance(IIslandEventQueue, container.resolve(PyPubSubIslandEventQueue))
 
+    _setup_agent_event_serializers(container)
     _register_repositories(container, data_dir)
     _register_services(container)
 
@@ -107,8 +112,8 @@ def _register_repositories(container: DIContainer, data_dir: Path):
     )
     container.register_instance(IUserRepository, container.resolve(JSONFileUserRepository))
 
-    # TODO: Replace with MongoEventRepository
+    # TODO: Figure out how to manage encryptor locking for MongoEventRepository
-    container.register_instance(IAgentEventRepository, StubbedEventRepository())
+    container.register_instance(IAgentEventRepository, container.resolve(MongoAgentEventRepository))
 
     container.register_instance(INodeRepository, container.resolve(MongoNodeRepository))
     container.register_instance(IMachineRepository, container.resolve(MongoMachineRepository))
@@ -130,9 +135,16 @@ def _build_agent_binary_repository():
     return agent_binary_repository
 
 
+def _setup_agent_event_serializers(container: DIContainer):
+    agent_event_serializer_registry = AgentEventSerializerRegistry()
+    register_common_agent_event_serializers(agent_event_serializer_registry)
+
+    container.register_instance(AgentEventSerializerRegistry, agent_event_serializer_registry)
+
+
 def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository):
     """
-    Logs all the hashes of the agent executables for debbuging ease
+    Logs all the hashes of the agent executables for debugging ease
 
     :param agent_binary_repository: Used to retrieve the agent binaries
     """

vulture_allowlist.py

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