From 51a9bad0008fcb3e2c56897121aeb316a93ee16e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 20 Sep 2022 15:48:07 +0530 Subject: [PATCH 1/4] Island: Register instances of INodeRepository, IMachineRepository, IAgentRepository in DI container --- monkey/monkey_island/cc/services/initialize.py | 10 ++++++++++ 1 file changed, 10 insertions(+) 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( From f93565b2cfcf69d4e413f7f490131a84c19b2ed2 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 20 Sep 2022 15:50:43 +0530 Subject: [PATCH 2/4] Island: Subscribe repository reset functions to CLEAR_SIMULATION_DATA event --- .../cc/setup/island_event_handlers.py | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index ab3e856d6..ea32e8e65 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,17 @@ def _subscribe_clear_simulation_data_events( island_event_queue.subscribe( IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials ) + + node_repository = container.resolve(INodeRepository) + island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, node_repository.reset) + + agent_event_repository = container.resolve(IAgentEventRepository) + island_event_queue.subscribe( + IslandEventTopic.CLEAR_SIMULATION_DATA, agent_event_repository.reset + ) + + agent_repository = container.resolve(IAgentRepository) + island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, agent_repository.reset) + + machine_repository = container.resolve(IMachineRepository) + island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, machine_repository.reset) From 638e4cd7ed955c2717f286d7c679e70efceda73e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 20 Sep 2022 17:05:53 +0530 Subject: [PATCH 3/4] Island: Simplify logic in island_event_handlers._subscribe_clear_simulation_data_events --- .../cc/setup/island_event_handlers.py | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index ea32e8e65..859483ee9 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -39,16 +39,11 @@ def _subscribe_clear_simulation_data_events( IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials ) - node_repository = container.resolve(INodeRepository) - island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, node_repository.reset) - - agent_event_repository = container.resolve(IAgentEventRepository) - island_event_queue.subscribe( - IslandEventTopic.CLEAR_SIMULATION_DATA, agent_event_repository.reset - ) - - agent_repository = container.resolve(IAgentRepository) - island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, agent_repository.reset) - - machine_repository = container.resolve(IMachineRepository) - island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, machine_repository.reset) + for i_repository in [ + INodeRepository, + IAgentEventRepository, + IAgentRepository, + IMachineRepository, + ]: + repository = container.resolve(i_repository) + island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, repository.reset) From 5032588c23630369ee1d21f53b161e34ba1e9ab8 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 20 Sep 2022 17:07:28 +0530 Subject: [PATCH 4/4] Island: Initialize MongoDB connection before DI container so that repositories using mongo can be registered and subscribed --- monkey/monkey_island/cc/server_setup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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)