Agent: Pass servers to exploit_host in ExploiterWrapper.Inner

This commit is contained in:
Ilija Lazoroski 2022-09-05 13:24:33 +02:00 committed by Mike Salvatore
parent 8fc0d52b8b
commit c51217507a
1 changed files with 8 additions and 8 deletions

View File

@ -1,5 +1,5 @@
import threading
from typing import Dict, List, Type
from typing import Dict, Sequence, Type
from common.event_queue import IAgentEventQueue
from infection_monkey.model import VictimHost
@ -24,21 +24,24 @@ class ExploiterWrapper:
telemetry_messenger: ITelemetryMessenger,
event_queue: IAgentEventQueue,
agent_binary_repository: IAgentBinaryRepository,
servers: List[str],
):
self._exploit_class = exploit_class
self._telemetry_messenger = telemetry_messenger
self._event_queue = event_queue
self._agent_binary_repository = agent_binary_repository
self._servers = servers
def exploit_host(
self, host: VictimHost, current_depth: int, options: Dict, interrupt: threading.Event
self,
host: VictimHost,
servers: Sequence[str],
current_depth: int,
options: Dict,
interrupt: threading.Event,
):
exploiter = self._exploit_class()
return exploiter.exploit_host(
host,
self._servers,
servers,
current_depth,
self._telemetry_messenger,
self._event_queue,
@ -52,12 +55,10 @@ class ExploiterWrapper:
telemetry_messenger: ITelemetryMessenger,
event_queue: IAgentEventQueue,
agent_binary_repository: IAgentBinaryRepository,
servers: List[str],
):
self._telemetry_messenger = telemetry_messenger
self._event_queue = event_queue
self._agent_binary_repository = agent_binary_repository
self._servers = servers
def wrap(self, exploit_class: Type[HostExploiter]):
return ExploiterWrapper.Inner(
@ -65,5 +66,4 @@ class ExploiterWrapper:
self._telemetry_messenger,
self._event_queue,
self._agent_binary_repository,
self._servers,
)