Agent: Modify IPuppet interface to take VictimHost instead of object

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -5,6 +5,8 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Dict, Iterable, List, Mapping, Sequence
 
+from infection_monkey.model import VictimHost
+
 from . import PluginType
 from .credential_collection import Credentials
 
@@ -110,15 +112,14 @@ 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: object, options: Dict, interrupt: threading.Event
+        self, name: str, host: VictimHost, options: Dict, interrupt: threading.Event
     ) -> ExploiterResultData:
         """
         Runs an exploiter against a remote host
         :param str name: The name of the exploiter to run
-        :param object host: The domain name or IP address of a host
+        :param VictimHost host: A VictimHost object representing the target to exploit
         :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

monkey/infection_monkey/puppet/mock_puppet.py

@@ -14,6 +14,7 @@ from infection_monkey.i_puppet import (
     PortStatus,
     PostBreachData,
 )
+from infection_monkey.model import VictimHost
 
 DOT_1 = "10.0.0.1"
 DOT_2 = "10.0.0.2"
@@ -136,7 +137,7 @@ class MockPuppet(IPuppet):
 
     # TODO: host should be VictimHost, at the moment it can't because of circular dependency
     def exploit_host(
-        self, name: str, host: object, options: Dict, interrupt: threading.Event
+        self, name: str, host: VictimHost, options: Dict, interrupt: threading.Event
     ) -> ExploiterResultData:
         logger.debug(f"exploit_hosts({name}, {host}, {options})")
         attempts = [

monkey/infection_monkey/puppet/puppet.py

@@ -13,6 +13,7 @@ from infection_monkey.i_puppet import (
     PortScanData,
     PostBreachData,
 )
+from infection_monkey.model import VictimHost
 
 from ..telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from .mock_puppet import MockPuppet
@@ -58,9 +59,8 @@ class Puppet(IPuppet):
         fingerprinter = self._plugin_registry.get_plugin(name, PluginType.FINGERPRINTER)
         return fingerprinter.get_host_fingerprint(host, ping_scan_data, port_scan_data, options)
 
-    # TODO: host should be VictimHost, at the moment it can't because of circular dependency
     def exploit_host(
-        self, name: str, host: object, options: Dict, interrupt: threading.Event
+        self, name: str, host: VictimHost, options: Dict, interrupt: threading.Event
     ) -> ExploiterResultData:
         exploiter = self._plugin_registry.get_plugin(name, PluginType.EXPLOITER)
         return exploiter.exploit_host(host, self._telemetry_messenger, options)