diff --git a/monkey/infection_monkey/i_puppet.py b/monkey/infection_monkey/i_puppet.py index 285e32bca..11da6a260 100644 --- a/monkey/infection_monkey/i_puppet.py +++ b/monkey/infection_monkey/i_puppet.py @@ -2,7 +2,9 @@ import abc import threading from collections import namedtuple from enum import Enum -from typing import Dict +from typing import Dict, Tuple + +from infection_monkey.puppet.plugin_type import PluginType class PortStatus(Enum): @@ -10,6 +12,10 @@ class PortStatus(Enum): CLOSED = 2 +class UnknownPluginError(Exception): + pass + + ExploiterResultData = namedtuple("ExploiterResultData", ["result", "info", "attempts"]) PingScanData = namedtuple("PingScanData", ["response_received", "os"]) PortScanData = namedtuple("PortScanData", ["port", "status", "banner", "service"]) @@ -18,6 +24,14 @@ PostBreachData = namedtuple("PostBreachData", ["command", "result"]) class IPuppet(metaclass=abc.ABCMeta): + @abc.abstractmethod + def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: + """ + Loads a plugin into the puppet. + :param object plugin: The plugin object to load + :param PluginType plugin_type: The type of plugin being loaded + """ + @abc.abstractmethod def run_sys_info_collector(self, name: str) -> Dict: """ @@ -91,7 +105,9 @@ class IPuppet(metaclass=abc.ABCMeta): """ @abc.abstractmethod - def run_payload(self, name: str, options: Dict, interrupt: threading.Event) -> None: + def run_payload( + self, name: str, options: Dict, interrupt: threading.Event + ) -> Tuple[None, bool, str]: """ Runs a payload :param str name: The name of the payload to run diff --git a/monkey/infection_monkey/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index d5c8fa2f8..5f0389752 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/monkey/infection_monkey/puppet/mock_puppet.py @@ -11,6 +11,7 @@ from infection_monkey.i_puppet import ( PortStatus, PostBreachData, ) +from infection_monkey.puppet.plugin_type import PluginType DOT_1 = "10.0.0.1" DOT_2 = "10.0.0.2" @@ -21,6 +22,9 @@ logger = logging.getLogger() class MockPuppet(IPuppet): + def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: + logger.debug(f"load_plugin({plugin}, {plugin_type})") + def run_sys_info_collector(self, name: str) -> Dict: logger.debug(f"run_sys_info_collector({name})") # TODO: More collectors diff --git a/monkey/infection_monkey/puppet/plugin_type.py b/monkey/infection_monkey/puppet/plugin_type.py new file mode 100644 index 000000000..4e20d7360 --- /dev/null +++ b/monkey/infection_monkey/puppet/plugin_type.py @@ -0,0 +1,9 @@ +from enum import Enum + + +class PluginType(Enum): + EXPLOITER = "Exploiter" + FINGERPRINTER = "Fingerprinter" + PAYLOAD = "Payload" + POST_BREACH_ACTION = "PBA" + SYSTEM_INFO_COLLECTOR = "SystemInfoCollector" diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py new file mode 100644 index 000000000..f932d84a4 --- /dev/null +++ b/monkey/infection_monkey/puppet/puppet.py @@ -0,0 +1,54 @@ +import logging +import threading +from typing import Dict, Tuple + +from infection_monkey.i_puppet import ( + ExploiterResultData, + FingerprintData, + IPuppet, + PingScanData, + PortScanData, + PostBreachData, +) +from infection_monkey.puppet.plugin_type import PluginType + +logger = logging.getLogger() + + +class Puppet(IPuppet): + def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: + pass + + def run_sys_info_collector(self, name: str) -> Dict: + pass + + def run_pba(self, name: str, options: Dict) -> PostBreachData: + pass + + def ping(self, host: str, timeout: float = 1) -> PingScanData: + pass + + def scan_tcp_port(self, host: str, port: int, timeout: float = 3) -> PortScanData: + pass + + def fingerprint( + self, + name: str, + host: str, + ping_scan_data: PingScanData, + port_scan_data: Dict[int, PortScanData], + ) -> FingerprintData: + pass + + def exploit_host( + self, name: str, host: str, options: Dict, interrupt: threading.Event + ) -> ExploiterResultData: + pass + + def run_payload( + self, name: str, options: Dict, interrupt: threading.Event + ) -> Tuple[None, bool, str]: + pass + + def cleanup(self) -> None: + pass