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 import threading
from typing import Dict, Type from typing import Dict, Type
from common.event_queue import IEventQueue
from infection_monkey.model import VictimHost from infection_monkey.model import VictimHost
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
@ -21,10 +22,12 @@ class ExploiterWrapper:
self, self,
exploit_class: Type[HostExploiter], exploit_class: Type[HostExploiter],
telemetry_messenger: ITelemetryMessenger, telemetry_messenger: ITelemetryMessenger,
event_queue: IEventQueue,
agent_repository: IAgentRepository, agent_repository: IAgentRepository,
): ):
self._exploit_class = exploit_class self._exploit_class = exploit_class
self._telemetry_messenger = telemetry_messenger self._telemetry_messenger = telemetry_messenger
self._event_queue = event_queue
self._agent_repository = agent_repository self._agent_repository = agent_repository
def exploit_host( def exploit_host(
@ -35,18 +38,23 @@ class ExploiterWrapper:
host, host,
current_depth, current_depth,
self._telemetry_messenger, self._telemetry_messenger,
self._event_queue,
self._agent_repository, self._agent_repository,
options, options,
interrupt, interrupt,
) )
def __init__( 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._telemetry_messenger = telemetry_messenger
self._event_queue = event_queue
self._agent_repository = agent_repository self._agent_repository = agent_repository
def wrap(self, exploit_class: Type[HostExploiter]): def wrap(self, exploit_class: Type[HostExploiter]):
return ExploiterWrapper.Inner( return ExploiterWrapper.Inner(
exploit_class, self._telemetry_messenger, self._agent_repository exploit_class, self._telemetry_messenger, self._event_queue, self._agent_repository
) )