diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 4b31de5ce..710d45465 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -65,12 +65,14 @@ def run_monkey_island(): ip_addresses, deployment, version = _collect_system_info() _send_analytics(deployment, version) + + _initialize_mongodb_connection(config_options.start_mongodb, config_options.data_dir) + 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) - _initialize_mongodb_connection(config_options.start_mongodb, config_options.data_dir) _start_island_server(ip_addresses, island_args.setup_only, config_options, container) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index df4f92d6b..695db13c7 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -24,13 +24,19 @@ from monkey_island.cc.repository import ( IAgentBinaryRepository, IAgentConfigurationRepository, IAgentEventRepository, + IAgentRepository, ICredentialsRepository, IFileRepository, + IMachineRepository, + INodeRepository, ISimulationRepository, IUserRepository, JSONFileUserRepository, LocalStorageFileRepository, + MongoAgentRepository, MongoCredentialsRepository, + MongoMachineRepository, + MongoNodeRepository, RetrievalError, StubbedEventRepository, ) @@ -104,6 +110,10 @@ def _register_repositories(container: DIContainer, data_dir: Path): # TODO: Replace with MongoEventRepository container.register_instance(IAgentEventRepository, StubbedEventRepository()) + container.register_instance(INodeRepository, container.resolve(MongoNodeRepository)) + container.register_instance(IMachineRepository, container.resolve(MongoMachineRepository)) + container.register_instance(IAgentRepository, container.resolve(MongoAgentRepository)) + def _decorate_file_repository(file_repository: IFileRepository) -> IFileRepository: return FileRepositoryLockingDecorator( diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index ab3e856d6..859483ee9 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -3,7 +3,13 @@ from functools import partial from common import DIContainer from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic from monkey_island.cc.island_event_handlers import reset_agent_configuration -from monkey_island.cc.repository import ICredentialsRepository +from monkey_island.cc.repository import ( + IAgentEventRepository, + IAgentRepository, + ICredentialsRepository, + IMachineRepository, + INodeRepository, +) from monkey_island.cc.services.database import Database @@ -32,3 +38,12 @@ def _subscribe_clear_simulation_data_events( island_event_queue.subscribe( IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials ) + + for i_repository in [ + INodeRepository, + IAgentEventRepository, + IAgentRepository, + IMachineRepository, + ]: + repository = container.resolve(i_repository) + island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, repository.reset)