From cc9cfc5e3b6f865623732a3300f20f58ec7e4d97 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 28 Feb 2022 15:43:20 -0500
Subject: [PATCH] Agent: Inject IAgentRepository into exploiters

---
 .../infection_monkey/exploit/HostExploiter.py | 11 +++++++++-
 .../exploit/exploiter_wrapper.py              | 20 +++++++++++++++----
 monkey/infection_monkey/monkey.py             |  7 +++++--
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/monkey/infection_monkey/exploit/HostExploiter.py b/monkey/infection_monkey/exploit/HostExploiter.py
index b74dc3871..69924b61a 100644
--- a/monkey/infection_monkey/exploit/HostExploiter.py
+++ b/monkey/infection_monkey/exploit/HostExploiter.py
@@ -9,6 +9,8 @@ from infection_monkey.config import WormConfiguration
 from infection_monkey.i_puppet import ExploiterResultData
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
+from . import IAgentRepository
+
 logger = logging.getLogger(__name__)
 
 
@@ -67,9 +69,16 @@ class HostExploiter:
         )
 
     # TODO: host should be VictimHost, at the moment it can't because of circular dependency
-    def exploit_host(self, host, telemetry_messenger: ITelemetryMessenger, options: Dict):
+    def exploit_host(
+        self,
+        host,
+        telemetry_messenger: ITelemetryMessenger,
+        agent_repository: IAgentRepository,
+        options: Dict,
+    ):
         self.host = host
         self.telemetry_messenger = telemetry_messenger
+        self.agent_repository = agent_repository
         self.options = options
 
         self.pre_exploit()
diff --git a/monkey/infection_monkey/exploit/exploiter_wrapper.py b/monkey/infection_monkey/exploit/exploiter_wrapper.py
index 444c89b31..c621ecaea 100644
--- a/monkey/infection_monkey/exploit/exploiter_wrapper.py
+++ b/monkey/infection_monkey/exploit/exploiter_wrapper.py
@@ -3,6 +3,7 @@ from typing import Dict, Type
 from infection_monkey.model import VictimHost
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
+from . import IAgentRepository
 from .HostExploiter import HostExploiter
 
 
@@ -16,17 +17,28 @@ class ExploiterWrapper:
 
     class Inner:
         def __init__(
-            self, exploit_class: Type[HostExploiter], telemetry_messenger: ITelemetryMessenger
+            self,
+            exploit_class: Type[HostExploiter],
+            telemetry_messenger: ITelemetryMessenger,
+            agent_repository: IAgentRepository,
         ):
             self._exploit_class = exploit_class
             self._telemetry_messenger = telemetry_messenger
+            self._agent_repository = agent_repository
 
         def exploit_host(self, host: VictimHost, options: Dict):
             exploiter = self._exploit_class()
-            return exploiter.exploit_host(host, self._telemetry_messenger, options)
+            return exploiter.exploit_host(
+                host, self._telemetry_messenger, self._agent_repository, options
+            )
 
-    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+    def __init__(
+        self, telemetry_messenger: ITelemetryMessenger, agent_repository: IAgentRepository
+    ):
         self._telemetry_messenger = telemetry_messenger
+        self._agent_repository = agent_repository
 
     def wrap(self, exploit_class: Type[HostExploiter]):
-        return ExploiterWrapper.Inner(exploit_class, self._telemetry_messenger)
+        return ExploiterWrapper.Inner(
+            exploit_class, self._telemetry_messenger, self._agent_repository
+        )
diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 3fb26f348..eaa6e0d90 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -16,7 +16,7 @@ from infection_monkey.credential_collectors import (
     MimikatzCredentialCollector,
     SSHCredentialCollector,
 )
-from infection_monkey.exploit import ExploiterWrapper
+from infection_monkey.exploit import CachingAgentRepository, ExploiterWrapper
 from infection_monkey.exploit.hadoop import HadoopExploiter
 from infection_monkey.exploit.sshexec import SSHExploiter
 from infection_monkey.i_puppet import IPuppet, PluginType
@@ -200,7 +200,10 @@ class InfectionMonkey:
         puppet.load_plugin("smb", SMBFingerprinter(), PluginType.FINGERPRINTER)
         puppet.load_plugin("ssh", SSHFingerprinter(), PluginType.FINGERPRINTER)
 
-        exploit_wrapper = ExploiterWrapper(self.telemetry_messenger)
+        agent_repoitory = CachingAgentRepository(
+            f"https://{self._default_server}", ControlClient.proxies
+        )
+        exploit_wrapper = ExploiterWrapper(self.telemetry_messenger, agent_repoitory)
 
         puppet.load_plugin(
             "SSHExploiter",