From 58b1a04bd702e3a3bd3dbab4d82b9aabb8c6a763 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Tue, 22 Feb 2022 19:30:53 +0100
Subject: [PATCH] Agent: Modify exploit_host() to accept object instead of
 string

---
 monkey/infection_monkey/i_puppet/i_puppet.py  | 7 +++++--
 monkey/infection_monkey/master/exploiter.py   | 2 +-
 monkey/infection_monkey/puppet/mock_puppet.py | 5 +++--
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py
index 79bd3b4fe..78c0c9659 100644
--- a/monkey/infection_monkey/i_puppet/i_puppet.py
+++ b/monkey/infection_monkey/i_puppet/i_puppet.py
@@ -102,16 +102,19 @@ class IPuppet(metaclass=abc.ABCMeta):
         :rtype: FingerprintData
         """
 
+    # TODO: host should be VictimHost, at the moment it can't because of circular dependency
     @abc.abstractmethod
     def exploit_host(
-        self, name: str, host: str, options: Dict, interrupt: threading.Event
+        self, name: str, host: object, options: Dict, interrupt: threading.Event
     ) -> ExploiterResultData:
         """
         Runs an exploiter against a remote host
         :param str name: The name of the exploiter to run
-        :param str host: The domain name or IP address of a host
+        :param object host: The domain name or IP address of a host
         :param Dict options: A dictionary containing options that modify the behavior of the
                              exploiter
+        :param threading.Event interrupt: A threading.Event object that signals the exploit to stop
+                                          executing and clean itself up.
         :return: True if exploitation was successful, False otherwise
         :rtype: ExploiterResultData
         """
diff --git a/monkey/infection_monkey/master/exploiter.py b/monkey/infection_monkey/master/exploiter.py
index 092bc78d8..5a76b20a8 100644
--- a/monkey/infection_monkey/master/exploiter.py
+++ b/monkey/infection_monkey/master/exploiter.py
@@ -115,7 +115,7 @@ class Exploiter:
         credentials = self._get_credentials_for_propagation()
         options = {"credentials": credentials, **options}
 
-        return self._puppet.exploit_host(exploiter_name, victim_host.ip_addr, options, stop)
+        return self._puppet.exploit_host(exploiter_name, victim_host, options, stop)
 
     def _get_credentials_for_propagation(self) -> Mapping:
         try:
diff --git a/monkey/infection_monkey/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py
index 453265f55..8a7f5935d 100644
--- a/monkey/infection_monkey/puppet/mock_puppet.py
+++ b/monkey/infection_monkey/puppet/mock_puppet.py
@@ -134,8 +134,9 @@ class MockPuppet(IPuppet):
 
         return empty_fingerprint_data
 
+    # TODO: host should be VictimHost, at the moment it can't because of circular dependency
     def exploit_host(
-        self, name: str, host: str, options: Dict, interrupt: threading.Event
+        self, name: str, host: object, options: Dict, interrupt: threading.Event
     ) -> ExploiterResultData:
         logger.debug(f"exploit_hosts({name}, {host}, {options})")
         attempts = [
@@ -209,7 +210,7 @@ class MockPuppet(IPuppet):
         }
 
         try:
-            return successful_exploiters[host][name]
+            return successful_exploiters[host.ip_addr][name]
         except KeyError:
             return ExploiterResultData(
                 False, False, os_linux, {}, [], f"{name} failed for host {host}"