Agent: Add servers list to ExploiterWrapper

This commit is contained in:
Ilija Lazoroski 2022-09-02 16:10:37 +02:00 committed by Mike Salvatore
parent f16f111543
commit e4d49f5a12
1 changed files with 7 additions and 1 deletions

View File

@ -1,5 +1,5 @@
import threading
from typing import Dict, Type
from typing import Dict, List, Type
from common.event_queue import IAgentEventQueue
from infection_monkey.model import VictimHost
@ -24,11 +24,13 @@ 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
@ -36,6 +38,7 @@ class ExploiterWrapper:
exploiter = self._exploit_class()
return exploiter.exploit_host(
host,
self._servers,
current_depth,
self._telemetry_messenger,
self._event_queue,
@ -49,10 +52,12 @@ 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(
@ -60,4 +65,5 @@ class ExploiterWrapper:
self._telemetry_messenger,
self._event_queue,
self._agent_binary_repository,
self._servers,
)