Agent: Mock out unimplemented functions in Puppet

This commit is contained in:
Mike Salvatore 2021-12-17 10:22:42 -05:00
parent 973c88678e
commit 7b8b485b57
1 changed files with 9 additions and 6 deletions

View File

@ -13,27 +13,30 @@ from infection_monkey.i_puppet import (
) )
from infection_monkey.puppet.plugin_registry import PluginRegistry from infection_monkey.puppet.plugin_registry import PluginRegistry
from .mock_puppet import MockPuppet
logger = logging.getLogger() logger = logging.getLogger()
class Puppet(IPuppet): class Puppet(IPuppet):
def __init__(self) -> None: def __init__(self) -> None:
self._mock_puppet = MockPuppet()
self._plugin_registry = PluginRegistry() self._plugin_registry = PluginRegistry()
def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None:
self._plugin_registry.load_plugin(plugin_name, plugin, plugin_type) self._plugin_registry.load_plugin(plugin_name, plugin, plugin_type)
def run_sys_info_collector(self, name: str) -> Dict: def run_sys_info_collector(self, name: str) -> Dict:
pass return self._mock_puppet.run_sys_info_collector(name)
def run_pba(self, name: str, options: Dict) -> PostBreachData: def run_pba(self, name: str, options: Dict) -> PostBreachData:
pass return self._mock_puppet.run_pba(name, options)
def ping(self, host: str, timeout: float = 1) -> PingScanData: def ping(self, host: str, timeout: float = 1) -> PingScanData:
pass return self._mock_puppet.ping(host, timeout)
def scan_tcp_port(self, host: str, port: int, timeout: float = 3) -> PortScanData: def scan_tcp_port(self, host: str, port: int, timeout: float = 3) -> PortScanData:
pass return self._mock_puppet.scan_tcp_port(host, port, timeout)
def fingerprint( def fingerprint(
self, self,
@ -42,12 +45,12 @@ class Puppet(IPuppet):
ping_scan_data: PingScanData, ping_scan_data: PingScanData,
port_scan_data: Dict[int, PortScanData], port_scan_data: Dict[int, PortScanData],
) -> FingerprintData: ) -> FingerprintData:
pass return self._mock_puppet.fingerprint(name, host, ping_scan_data, port_scan_data)
def exploit_host( def exploit_host(
self, name: str, host: str, options: Dict, interrupt: threading.Event self, name: str, host: str, options: Dict, interrupt: threading.Event
) -> ExploiterResultData: ) -> ExploiterResultData:
pass return self._mock_puppet.exploit_host(name, host, options, interrupt)
def run_payload(self, name: str, options: Dict, interrupt: threading.Event): def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
payload = self._plugin_registry.get_plugin(name, PluginType.PAYLOAD) payload = self._plugin_registry.get_plugin(name, PluginType.PAYLOAD)