Agent: Finish implementing MockMaster

Also modified ExploitTelem and PostBreachTelem internals, and
MockPuppet.

monkey/infection_monkey/i_puppet.py

@@ -10,7 +10,9 @@ class PortStatus(Enum):
     CLOSED = 2
 
 
+ExploiterResultData = namedtuple("ExploiterResultData", ["result", "info", "attempts"])
 PortScanData = namedtuple("PortScanData", ["port", "status", "banner", "service"])
+PostBreachData = namedtuple("PostBreachData", ["command", "result"])
 
 
 class IPuppet(metaclass=abc.ABCMeta):
@@ -24,11 +26,12 @@ class IPuppet(metaclass=abc.ABCMeta):
         """
 
     @abc.abstractmethod
-    def run_pba(self, name: str, options: Dict) -> None:
+    def run_pba(self, name: str, options: Dict) -> PostBreachData:
         """
         Runs a post-breach action (PBA)
         :param str name: The name of the post-breach action to run
         :param Dict options: A dictionary containing options that modify the behavior of the PBA
+        :rtype: PostBreachData
         """
 
     @abc.abstractmethod
@@ -63,7 +66,9 @@ class IPuppet(metaclass=abc.ABCMeta):
         """
 
     @abc.abstractmethod
-    def exploit_host(self, name: str, host: str, options: Dict, interrupt: threading.Event) -> bool:
+    def exploit_host(
+        self, name: str, host: str, options: Dict, interrupt: threading.Event
+    ) -> ExploiterResultData:
         """
         Runs an exploiter against a remote host
         :param str name: The name of the exploiter to run
@@ -71,7 +76,7 @@ class IPuppet(metaclass=abc.ABCMeta):
         :param Dict options: A dictionary containing options that modify the behavior of the
                              exploiter
         :return: True if exploitation was successful, False otherwise
-        :rtype: bool
+        :rtype: ExploiterResultData
         """
 
     @abc.abstractmethod

monkey/infection_monkey/master/mock_master.py

@@ -6,6 +6,7 @@ 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
 from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
 
@@ -16,6 +17,12 @@ class MockMaster(IMaster):
     def __init__(self, puppet: IPuppet, telemetry_messenger: ITelemetryMessenger):
         self._puppet = puppet
         self._telemetry_messenger = telemetry_messenger
+        self._hosts = {
+            "10.0.0.1": VictimHost("10.0.0.1"),
+            "10.0.0.2": VictimHost("10.0.0.2"),
+            "10.0.0.3": VictimHost("10.0.0.3"),
+            "10.0.0.4": VictimHost("10.0.0.4"),
+        }
 
     def start(self) -> None:
         self._run_sys_info_collectors()
@@ -37,8 +44,13 @@ class MockMaster(IMaster):
         self._telemetry_messenger.send_telemetry(SystemInfoTelem(system_info))
 
     def _run_pbas(self):
-        self._puppet.run_pba("AccountDiscovery", {})
+        name = "AccountDiscovery"
-        self._puppet.run_pba("CommunicateAsBackdoorUser", {})
+        command, result = self._puppet.run_pba(name, {})
+        self._telemetry_messenger.send_telemetry(PostBreachTelem(name, command, result))
+
+        name = "CommunicateAsBackdoorUser"
+        command, result = self._puppet.run_pba(name, {})
+        self._telemetry_messenger.send_telemetry(PostBreachTelem(name, command, result))
 
     def _scan_victims(self):
         # TODO: The telemetry must be malformed somehow, or something else is wrong. This causes the
@@ -46,7 +58,7 @@ class MockMaster(IMaster):
         ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3"]
         ports = [22, 445, 3389, 8008]
         for ip in ips:
-            h = VictimHost(ip)
+            h = self._hosts[ip]
 
             (response_received, os) = self._puppet.ping(ip)
             h.icmp = response_received
@@ -65,8 +77,8 @@ class MockMaster(IMaster):
             self._telemetry_messenger.send_telemetry(ScanTelem(h))
 
     def _fingerprint(self):
-        machine_1 = VictimHost("10.0.0.1")
+        machine_1 = self._hosts["10.0.0.1"]
-        machine_3 = VictimHost("10.0.0.3")
+        machine_3 = self._hosts["10.0.0.3"]
 
         self._puppet.fingerprint("SMBFinger", machine_1)
         self._telemetry_messenger.send_telemetry(ScanTelem(machine_1))
@@ -78,19 +90,22 @@ class MockMaster(IMaster):
         self._telemetry_messenger.send_telemetry(ScanTelem(machine_3))
 
     def _exploit(self):
-        # TODO: modify what ExploitTelem gets
+        result, info, attempts = self._puppet.exploit_host(
-        self._telemetry_messenger.send_telemetry(
+            "PowerShellExploiter", "10.0.0.1", {}, None
-            ExploitTelem(self._puppet.exploit_host("PowerShellExploiter", "10.0.0.1", {}, None))
         )
         self._telemetry_messenger.send_telemetry(
-            ExploitTelem(self._puppet.exploit_host("SSHExploiter", "10.0.0.3", {}, None))
+            ExploitTelem("PowerShellExploiter", self._hosts["10.0.0.1"], result, info, attempts)
+        )
+
+        result, info, attempts = self._puppet.exploit_host("SSHExploiter", "10.0.0.3", {}, None)
+        self._telemetry_messenger.send_telemetry(
+            ExploitTelem("SSHExploiter", self._hosts["10.0.0.3"], result, info, attempts)
         )
 
     def _run_payload(self):
         # TODO: modify what FileEncryptionTelem gets
-        self._telemetry_messenger.send_telemetry(
+        path, success, error = self._puppet.run_payload("RansomwarePayload", {}, None)
-            FileEncryptionTelem(self._run_payload("RansomwarePayload", {}, None))
+        self._telemetry_messenger.send_telemetry(FileEncryptionTelem(path, success, error))
-        )
 
     def terminate(self) -> None:
         logger.info("Terminating MockMaster")

monkey/infection_monkey/puppet/mock_puppet.py

@@ -1,8 +1,14 @@
 import logging
 import threading
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
-from infection_monkey.i_puppet import IPuppet, PortScanData, PortStatus
+from infection_monkey.i_puppet import (
+    ExploiterResultData,
+    IPuppet,
+    PortScanData,
+    PortStatus,
+    PostBreachData,
+)
 
 DOT_1 = "10.0.0.1"
 DOT_2 = "10.0.0.2"
@@ -141,9 +147,15 @@ class MockPuppet(IPuppet):
 
         return {}
 
-    def run_pba(self, name: str, options: Dict) -> None:
+    def run_pba(self, name: str, options: Dict) -> List[Tuple[str, bool]]:
         logger.debug(f"run_pba({name}, {options})")
-        return None
+        result_1 = PostBreachData("pba command 1", "pba result 1")
+        result_2 = PostBreachData("pba command 2", "pba result 2")
+
+        return [
+            (result_1.command, result_1.result, True),
+            (result_2.command, result_2.result, False),
+        ]
 
     def ping(self, host: str) -> Tuple[bool, Optional[str]]:
         logger.debug(f"run_ping({host})")
@@ -208,13 +220,30 @@ class MockPuppet(IPuppet):
 
     def exploit_host(self, name: str, host: str, options: Dict, interrupt: threading.Event) -> bool:
         logger.debug(f"exploit_hosts({name}, {host}, {options})")
-        successful_exploiters = {DOT_1: {"PowerShellExploiter"}, DOT_3: {"SSHExploiter"}}
+        successful_exploiters = {
+            DOT_1: {
+                "PowerShellExploiter": ExploiterResultData(
+                    True, {"info": "important success stuff"}, ["attempt 1"]
+                )
+            },
+            DOT_3: {
+                "SSHExploiter": ExploiterResultData(
+                    False, {"info": "important failure stuff"}, ["attempt 2"]
+                )
+            },
+        }
 
-        return name in successful_exploiters.get(host, {})
+        return (
+            successful_exploiters[host][name].result,
+            successful_exploiters[host][name].info,
+            successful_exploiters[host][name].attempts,
+        )
 
-    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]:
         logger.debug(f"run_payload({name}, {options})")
-        return None
+        return (None, True, "")
 
     def cleanup(self) -> None:
         print("Cleanup called!")

monkey/infection_monkey/telemetry/exploit_telem.py

@@ -1,25 +1,35 @@
+from typing import Dict, List
+
 from common.common_consts.telem_categories import TelemCategoryEnum
+from infection_monkey.model.host import VictimHost
 from infection_monkey.telemetry.base_telem import BaseTelem
 
 
 class ExploitTelem(BaseTelem):
-    def __init__(self, exploiter, result):
+    def __init__(self, name: str, host: VictimHost, result: bool, info: Dict, attempts: List):
         """
         Default exploit telemetry constructor
-        :param exploiter: The instance of exploiter used
+        :param name: The name of exploiter used
-        :param result: The result from the 'exploit_host' method.
+        :param host: The host machine
+        :param result: The result from the 'exploit_host' method
+        :param info: Information about the exploiter
+        :param attempts: Information about the exploiter's attempts
         """
         super(ExploitTelem, self).__init__()
-        self.exploiter = exploiter
+
+        self.name = name
+        self.host = host.__dict__
         self.result = result
+        self.info = info
+        self.attempts = attempts
 
     telem_category = TelemCategoryEnum.EXPLOIT
 
-    def get_data(self):
+    def get_data(self) -> Dict:
         return {
             "result": self.result,
-            "machine": self.exploiter.host.__dict__,
+            "machine": self.host,
-            "exploiter": self.exploiter.__class__.__name__,
+            "exploiter": self.name,
-            "info": self.exploiter.exploit_info,
+            "info": self.info,
-            "attempts": self.exploiter.exploit_attempts,
+            "attempts": self.attempts,
         }

monkey/infection_monkey/telemetry/post_breach_telem.py

@@ -1,4 +1,5 @@
 import socket
+from typing import Dict, Tuple
 
 from common.common_consts.telem_categories import TelemCategoryEnum
 from infection_monkey.telemetry.base_telem import BaseTelem
@@ -6,31 +7,33 @@ from infection_monkey.utils.environment import is_windows_os
 
 
 class PostBreachTelem(BaseTelem):
-    def __init__(self, pba, result):
+    def __init__(self, name: str, command: str, result: str) -> None:
         """
         Default post breach telemetry constructor
-        :param pba: Post breach action which was used
+        :param name: Name of post breach action
+        :param command: Command used as PBA
         :param result: Result of PBA
         """
         super(PostBreachTelem, self).__init__()
-        self.pba = pba
+        self.name = name
+        self.command = command
         self.result = result
         self.hostname, self.ip = PostBreachTelem._get_hostname_and_ip()
 
     telem_category = TelemCategoryEnum.POST_BREACH
 
-    def get_data(self):
+    def get_data(self) -> Dict:
         return {
-            "command": self.pba.command,
+            "command": self.command,
             "result": self.result,
-            "name": self.pba.name,
+            "name": self.name,
             "hostname": self.hostname,
             "ip": self.ip,
             "os": PostBreachTelem._get_os(),
         }
 
     @staticmethod
-    def _get_hostname_and_ip():
+    def _get_hostname_and_ip() -> Tuple[str, str]:
         try:
             hostname = socket.gethostname()
             ip = socket.gethostbyname(hostname)
@@ -40,5 +43,5 @@ class PostBreachTelem(BaseTelem):
         return hostname, ip
 
     @staticmethod
-    def _get_os():
+    def _get_os() -> str:
         return "Windows" if is_windows_os() else "Linux"

vulture_allowlist.py

@@ -212,3 +212,4 @@ MockPuppet
 ControlChannel
 should_agent_stop
 get_credentials_for_propagation
+MockMaster