From 32d618ac9299dbe99610593ef5f5c4bfc8d344e4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 22 Feb 2022 12:09:40 -0500 Subject: [PATCH] Agent: Modify IPuppet interface to take VictimHost instead of object --- monkey/infection_monkey/i_puppet/i_puppet.py | 7 ++++--- monkey/infection_monkey/puppet/mock_puppet.py | 3 ++- monkey/infection_monkey/puppet/puppet.py | 4 ++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py index fb861c76f..5b27de4f6 100644 --- a/monkey/infection_monkey/i_puppet/i_puppet.py +++ b/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 diff --git a/monkey/infection_monkey/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index 8a7f5935d..d43d48983 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/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 = [ diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index c06f047bf..1df3df885 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/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)