Merge pull request #1628 from guardicore/1593-i-master

Add IMaster and MockMaster

monkey/infection_monkey/i_master.py

@@ -0,0 +1,22 @@
+import abc
+
+
+class IMaster(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def start(self) -> None:
+        """
+        Run the control logic that will instruct the Puppet to perform various actions like scanning
+        or exploiting a specific host.
+        """
+
+    @abc.abstractmethod
+    def terminate(self) -> None:
+        """
+        Stop the master and interrupt any actions that are currently being executed.
+        """
+
+    @abc.abstractmethod
+    def cleanup(self) -> None:
+        """
+        Revert any changes that the master has directly or indirectly caused to the system.
+        """

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

@@ -0,0 +1,124 @@
+import logging
+
+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
+from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
+
+logger = logging.getLogger()
+
+
+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()
+        self._run_pbas()
+        self._scan_victims()
+        self._fingerprint()
+        self._exploit()
+        self._run_payload()
+
+    def _run_sys_info_collectors(self):
+        logging.info("Running system info collectors")
+        system_info_telemetry = {}
+        system_info_telemetry["ProcessListCollector"] = self._puppet.run_sys_info_collector(
+            "ProcessListCollector"
+        )
+        self._telemetry_messenger.send_telemetry(
+            SystemInfoTelem({"collectors": system_info_telemetry})
+        )
+        system_info = self._puppet.run_sys_info_collector("LinuxInfoCollector")
+        self._telemetry_messenger.send_telemetry(SystemInfoTelem(system_info))
+        logging.info("Finished running system info collectors")
+
+    def _run_pbas(self):
+        logging.info("Running post breach actions")
+        name = "AccountDiscovery"
+        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))
+        logging.info("Finished running post breach actions")
+
+    def _scan_victims(self):
+        logging.info("Scanning network for potential victims")
+        ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3"]
+        ports = [22, 445, 3389, 8008]
+        for ip in ips:
+            h = self._hosts[ip]
+
+            (response_received, os) = self._puppet.ping(ip)
+            h.icmp = response_received
+            if os is not None:
+                h.os["type"] = os
+
+            for p in ports:
+                port_scan_data = self._puppet.scan_tcp_port(ip, p)
+                if port_scan_data.status == PortStatus.OPEN:
+                    h.services[port_scan_data.service] = {}
+                    h.services[port_scan_data.service]["display_name"] = "unknown(TCP)"
+                    h.services[port_scan_data.service]["port"] = port_scan_data.port
+                    if port_scan_data.banner is not None:
+                        h.services[port_scan_data.service]["banner"] = port_scan_data.banner
+
+            self._telemetry_messenger.send_telemetry(ScanTelem(h))
+        logging.info("Finished scanning network for potential victims")
+
+    def _fingerprint(self):
+        logging.info("Running fingerprinters on potential victims")
+        machine_1 = self._hosts["10.0.0.1"]
+        machine_3 = self._hosts["10.0.0.3"]
+
+        self._puppet.fingerprint("SMBFinger", machine_1)
+        self._telemetry_messenger.send_telemetry(ScanTelem(machine_1))
+
+        self._puppet.fingerprint("SMBFinger", machine_3)
+        self._telemetry_messenger.send_telemetry(ScanTelem(machine_3))
+
+        self._puppet.fingerprint("HTTPFinger", machine_3)
+        self._telemetry_messenger.send_telemetry(ScanTelem(machine_3))
+        logging.info("Finished running fingerprinters on potential victims")
+
+    def _exploit(self):
+        logging.info("Exploiting victims")
+        result, info, attempts = self._puppet.exploit_host(
+            "PowerShellExploiter", "10.0.0.1", {}, None
+        )
+        self._telemetry_messenger.send_telemetry(
+            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)
+        )
+        logging.info("Finished exploiting victims")
+
+    def _run_payload(self):
+        logging.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))
+        logging.info("Finished running payloads")
+
+    def terminate(self) -> None:
+        logger.info("Terminating MockMaster")
+
+    def cleanup(self) -> None:
+        pass

monkey/infection_monkey/monkey.py

@@ -13,18 +13,23 @@ from common.version import get_version
 from infection_monkey.config import WormConfiguration
 from infection_monkey.control import ControlClient
 from infection_monkey.exploit.HostExploiter import HostExploiter
+from infection_monkey.master.mock_master import MockMaster
 from infection_monkey.model import DELAY_DELETE_CMD
 from infection_monkey.network.firewall import app as firewall
 from infection_monkey.network.HostFinger import HostFinger
 from infection_monkey.network.network_scanner import NetworkScanner
 from infection_monkey.network.tools import get_interface_to_target, is_running_on_island
 from infection_monkey.post_breach.post_breach_handler import PostBreach
+from infection_monkey.puppet.mock_puppet import MockPuppet
 from infection_monkey.ransomware.ransomware_payload_builder import build_ransomware_payload
 from infection_monkey.system_info import SystemInfoCollector
 from infection_monkey.system_singleton import SystemSingleton
 from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
 from infection_monkey.telemetry.attack.t1107_telem import T1107Telem
 from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem
+from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
+    LegacyTelemetryMessengerAdapter,
+)
 from infection_monkey.telemetry.scan_telem import ScanTelem
 from infection_monkey.telemetry.state_telem import StateTelem
 from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
@@ -38,6 +43,7 @@ from infection_monkey.utils.monkey_dir import (
     remove_monkey_dir,
 )
 from infection_monkey.utils.monkey_log_path import get_monkey_log_path
+from infection_monkey.utils.signal_handler import register_signal_handlers
 from infection_monkey.windows_upgrader import WindowsUpgrader
 
 MAX_DEPTH_REACHED_MESSAGE = "Reached max depth, skipping propagation phase."
@@ -107,6 +113,9 @@ class InfectionMonkey(object):
             logger.info("Monkey is starting...")
 
             logger.debug("Starting the setup phase.")
+            mock_master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
+            register_signal_handlers(mock_master)
+
             # Sets island's IP and port for monkey to communicate to
             self.set_default_server()
             self.set_default_port()

monkey/infection_monkey/puppet/mock_puppet.py

@@ -2,7 +2,13 @@ import logging
 import threading
 from typing import Dict, 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,13 @@ class MockPuppet(IPuppet):
 
         return {}
 
-    def run_pba(self, name: str, options: Dict) -> None:
+    def run_pba(self, name: str, options: Dict) -> PostBreachData:
         logger.debug(f"run_pba({name}, {options})")
-        return None
+
+        if name == "AccountDiscovery":
+            return PostBreachData("pba command 1", ["pba result 1", True])
+        else:
+            return PostBreachData("pba command 2", ["pba result 2", False])
 
     def ping(self, host: str) -> Tuple[bool, Optional[str]]:
         logger.debug(f"run_ping({host})")
@@ -154,7 +164,7 @@ class MockPuppet(IPuppet):
             return (False, None)
 
         if host == DOT_3:
-            return (True, "Linux")
+            return (True, "linux")
 
         if host == DOT_4:
             return (False, None)
@@ -206,15 +216,30 @@ class MockPuppet(IPuppet):
 
         return {}
 
-    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:
         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]
 
-    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 result: The result from the 'exploit_host' method.
+        :param name: The name of exploiter used
+        :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__,
-            "exploiter": self.exploiter.__class__.__name__,
-            "info": self.exploiter.exploit_info,
-            "attempts": self.exploiter.exploit_attempts,
+            "machine": self.host,
+            "exploiter": self.name,
+            "info": self.info,
+            "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"

monkey/infection_monkey/utils/signal_handler.py

@@ -0,0 +1,33 @@
+import logging
+import signal
+
+from infection_monkey.i_master import IMaster
+from infection_monkey.utils.environment import is_windows_os
+from infection_monkey.utils.exceptions.planned_shutdown_exception import PlannedShutdownException
+
+logger = logging.getLogger(__name__)
+
+
+class StopSignalHandler:
+    def __init__(self, master: IMaster):
+        self._master = master
+
+    def __call__(self, signum, _=None):
+        logger.info(f"The Monkey Agent received signal {signum}")
+        self._master.terminate()
+        raise PlannedShutdownException("Monkey Agent got an interrupt signal")
+
+
+def register_signal_handlers(master: IMaster):
+    stop_signal_handler = StopSignalHandler(master)
+    signal.signal(signal.SIGINT, stop_signal_handler)
+    signal.signal(signal.SIGTERM, stop_signal_handler)
+
+    if is_windows_os():
+        import win32api
+
+        signal.signal(signal.SIGBREAK, stop_signal_handler)
+
+        # CTRL_CLOSE_EVENT signal has a timeout of 5000ms,
+        # after that OS will forcefully kill the process
+        win32api.SetConsoleCtrlHandler(stop_signal_handler, True)

monkey/tests/unit_tests/infection_monkey/telemetry/test_exploit_telem.py

@@ -19,7 +19,6 @@ HOST_AS_DICT = {
     "default_tunnel": None,
     "default_server": None,
 }
-EXPLOITER = SSHExploiter(HOST)
 EXPLOITER_NAME = "SSHExploiter"
 EXPLOITER_INFO = {
     "display_name": SSHExploiter._EXPLOITED_SERVICE,
@@ -35,7 +34,7 @@ RESULT = False
 
 @pytest.fixture
 def exploit_telem_test_instance():
-    return ExploitTelem(EXPLOITER, RESULT)
+    return ExploitTelem(EXPLOITER_NAME, HOST, RESULT, EXPLOITER_INFO, EXPLOITER_ATTEMPTS)
 
 
 def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry):

monkey/tests/unit_tests/infection_monkey/telemetry/test_post_breach_telem.py

@@ -20,10 +20,9 @@ class StubSomePBA:
 
 @pytest.fixture
 def post_breach_telem_test_instance(monkeypatch):
-    PBA = StubSomePBA()
     monkeypatch.setattr(PostBreachTelem, "_get_hostname_and_ip", lambda: (HOSTNAME, IP))
     monkeypatch.setattr(PostBreachTelem, "_get_os", lambda: OS)
-    return PostBreachTelem(PBA, RESULT)
+    return PostBreachTelem(PBA_NAME, PBA_COMMAND, RESULT)
 
 
 def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry):

vulture_allowlist.py

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