Agent: Modify exploiter wrapper to accept IEventQueue

This commit is contained in:
Ilija Lazoroski 2022-08-15 17:55:25 +02:00
parent c6cb477474
commit f171e548f3
1 changed files with 10 additions and 2 deletions

View File

@ -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
)