diff --git a/monkey/monkey_island/cc/setup/agent_event_handlers.py b/monkey/monkey_island/cc/setup/agent_event_handlers.py index b01392361..4d61a5610 100644 --- a/monkey/monkey_island/cc/setup/agent_event_handlers.py +++ b/monkey/monkey_island/cc/setup/agent_event_handlers.py @@ -1,11 +1,18 @@ from common import DIContainer -from common.agent_events import CredentialsStolenEvent +from common.agent_events import CredentialsStolenEvent, PingScanEvent from common.event_queue import IAgentEventQueue from monkey_island.cc.agent_event_handlers import ( + handle_scan_data, save_event_to_event_repository, save_stolen_credentials_to_repository, ) -from monkey_island.cc.repository import IAgentEventRepository, ICredentialsRepository +from monkey_island.cc.repository import ( + IAgentEventRepository, + IAgentRepository, + ICredentialsRepository, + IMachineRepository, + INodeRepository, +) def setup_agent_event_handlers(container: DIContainer): @@ -22,3 +29,16 @@ def _subscribe_and_store_to_event_repository(container: DIContainer): container.resolve(ICredentialsRepository) ) agent_event_queue.subscribe_type(CredentialsStolenEvent, save_stolen_credentials_subscriber) + + +def _subscribe_ping_scan_event(container: DIContainer): + # Mypy don't handle cases where abstract class is passed as Type[...] + # https://github.com/python/mypy/issues/4717 + agent_event_queue = container.resolve(IAgentEventQueue) # type: ignore + agent_repository = container.resolve(IAgentRepository) # type: ignore + machine_repository = container.resolve(IMachineRepository) # type: ignore + node_repository = container.resolve(INodeRepository) # type: ignore + + handler = handle_scan_data(agent_repository, machine_repository, node_repository) + + agent_event_queue.subscribe_type(PingScanEvent, handler)