Agent: Modify IPuppet interface to take VictimHost instead of object

This commit is contained in:
Mike Salvatore 2022-02-22 12:09:40 -05:00
parent b17c85cd01
commit 32d618ac92
3 changed files with 8 additions and 6 deletions

View File

@ -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

View File

@ -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 = [

View File

@ -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)