From dc1eeefbc19c2fdec08b3cb5075833fb20bc156d Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 14:36:40 +0000 Subject: [PATCH] Island: Add MongoEventRepository to DI container --- monkey/monkey_island/cc/server_setup.py | 14 +------------ .../monkey_island/cc/services/initialize.py | 20 +++++++++++++++---- vulture_allowlist.py | 1 - 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 710d45465..aca537411 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/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: diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 695db13c7..70556d811 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/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 - container.register_instance(IAgentEventRepository, StubbedEventRepository()) + # TODO: Figure out how to manage encryptor locking for MongoEventRepository + 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 """ diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 53001be59..2843b4efb 100644 --- a/vulture_allowlist.py +++ b/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