diff --git a/monkey/infection_monkey/exploit/exploiter_wrapper.py b/monkey/infection_monkey/exploit/exploiter_wrapper.py index 540e0b4a4..8e9d4f22f 100644 --- a/monkey/infection_monkey/exploit/exploiter_wrapper.py +++ b/monkey/infection_monkey/exploit/exploiter_wrapper.py @@ -1,6 +1,7 @@ import threading from typing import Dict, Type +from common.event_queue import IEventQueue from infection_monkey.model import VictimHost from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger @@ -21,10 +22,12 @@ class ExploiterWrapper: self, exploit_class: Type[HostExploiter], telemetry_messenger: ITelemetryMessenger, + event_queue: IEventQueue, agent_repository: IAgentRepository, ): self._exploit_class = exploit_class self._telemetry_messenger = telemetry_messenger + self._event_queue = event_queue self._agent_repository = agent_repository def exploit_host( @@ -35,18 +38,23 @@ class ExploiterWrapper: host, current_depth, self._telemetry_messenger, + self._event_queue, self._agent_repository, options, interrupt, ) def __init__( - self, telemetry_messenger: ITelemetryMessenger, agent_repository: IAgentRepository + self, + telemetry_messenger: ITelemetryMessenger, + event_queue: IEventQueue, + agent_repository: IAgentRepository, ): self._telemetry_messenger = telemetry_messenger + self._event_queue = event_queue self._agent_repository = agent_repository def wrap(self, exploit_class: Type[HostExploiter]): return ExploiterWrapper.Inner( - exploit_class, self._telemetry_messenger, self._agent_repository + exploit_class, self._telemetry_messenger, self._event_queue, self._agent_repository )