Agent: Change return type of IPuppet.run_payload() to None

At the moment, we don't expect payloads to return any values. This may
be reevaluated as development proceeds or when telemetry is refactored.
This commit is contained in:
Mike Salvatore 2021-12-17 08:14:53 -05:00
parent 3adb1d5b07
commit c18af3c3fb
4 changed files with 9 additions and 17 deletions

View File

@ -2,7 +2,7 @@ import abc
import threading
from collections import namedtuple
from enum import Enum
from typing import Dict, Tuple
from typing import Dict
from infection_monkey.puppet.plugin_type import PluginType
@ -107,13 +107,13 @@ class IPuppet(metaclass=abc.ABCMeta):
"""
@abc.abstractmethod
def run_payload(
self, name: str, options: Dict, interrupt: threading.Event
) -> Tuple[None, bool, str]:
def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
"""
Runs a payload
:param str name: The name of the payload to run
:param Dict options: A dictionary containing options that modify the behavior of the payload
:param threading.Event interrupt: A threading.Event object that signals the payload to stop
executing and clean itself up.
"""
@abc.abstractmethod

View File

@ -4,7 +4,6 @@ from infection_monkey.i_master import IMaster
from infection_monkey.i_puppet import IPuppet, PortStatus
from infection_monkey.model.host import VictimHost
from infection_monkey.telemetry.exploit_telem import ExploitTelem
from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
from infection_monkey.telemetry.scan_telem import ScanTelem
@ -119,9 +118,7 @@ class MockMaster(IMaster):
def _run_payload(self):
logger.info("Running payloads")
# TODO: modify what FileEncryptionTelem gets
path, success, error = self._puppet.run_payload("RansomwarePayload", {}, None)
self._telemetry_messenger.send_telemetry(FileEncryptionTelem(path, success, error))
self._puppet.run_payload("RansomwarePayload", {}, None)
logger.info("Finished running payloads")
def terminate(self, block: bool = False) -> None:

View File

@ -1,6 +1,6 @@
import logging
import threading
from typing import Dict, Tuple
from typing import Dict
from infection_monkey.i_puppet import (
ExploiterResultData,
@ -299,11 +299,8 @@ class MockPuppet(IPuppet):
except KeyError:
return ExploiterResultData(False, {}, [], f"{name} failed for host {host}")
def run_payload(
self, name: str, options: Dict, interrupt: threading.Event
) -> Tuple[None, bool, str]:
def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
logger.debug(f"run_payload({name}, {options})")
return (None, True, "")
def cleanup(self) -> None:
print("Cleanup called!")

View File

@ -1,6 +1,6 @@
import logging
import threading
from typing import Dict, Tuple
from typing import Dict
from infection_monkey.i_puppet import (
ExploiterResultData,
@ -45,9 +45,7 @@ class Puppet(IPuppet):
) -> ExploiterResultData:
pass
def run_payload(
self, name: str, options: Dict, interrupt: threading.Event
) -> Tuple[None, bool, str]:
def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
pass
def cleanup(self) -> None: