Merge branch '2952-smb-preconditions' into 2952-smb-exploiter-plugin

Issue #2952
PR #3143

monkey/agent_plugins/exploiters/smb/src/plugin.py

@@ -22,8 +22,15 @@
 
 from .smb_command_builder import build_smb_command
 from .smb_options import SMBOptions
+from .smb_remote_access_client import SMB_PORTS
 from .smb_remote_access_client_factory import SMBRemoteAccessClientFactory
 
+
+def should_attempt_exploit(host: TargetHost) -> bool:
+    closed_tcp_ports = host.ports_status.tcp_ports.closed
+    return not all([p in closed_tcp_ports for p in SMB_PORTS])
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -67,6 +74,13 @@ def run(
             logger.exception(msg)
             return ExploiterResultData(error_message=msg)
 
+        if not should_attempt_exploit(host):
+            msg = f"Host {host.ip} has no open SMB ports"
+            logger.debug(msg)
+            return ExploiterResultData(
+                exploitation_success=False, propagation_success=False, error_message=msg
+            )
+
         command_builder = partial(
             build_smb_command,
             servers,

monkey/infection_monkey/i_puppet/__init__.py

@@ -10,4 +10,4 @@
 )
 from .i_fingerprinter import IFingerprinter
 from .i_credential_collector import ICredentialCollector
-from .target_host import TargetHost, TargetHostPorts
+from .target_host import TargetHost, TargetHostPorts, PortScanDataDict

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -4,9 +4,9 @@
 from common.agent_plugins import AgentPluginType
 from common.credentials import Credentials
 from common.types import Event, NetworkPort
-from infection_monkey.i_puppet.target_host import TargetHost
 
-from . import ExploiterResultData, FingerprintData, PingScanData, PortScanData
+from . import ExploiterResultData, FingerprintData, PingScanData
+from .target_host import PortScanDataDict, TargetHost
 
 
 class UnknownPluginError(Exception):
@@ -58,7 +58,7 @@ def ping(self, host: str, timeout: float) -> PingScanData:
     @abc.abstractmethod
     def scan_tcp_ports(
         self, host: str, ports: Sequence[NetworkPort], timeout: float = 3
-    ) -> Dict[NetworkPort, PortScanData]:
+    ) -> PortScanDataDict:
         """
         Scans a list of TCP ports on a remote host
 
@@ -74,7 +74,7 @@ def fingerprint(
         name: str,
         host: str,
         ping_scan_data: PingScanData,
-        port_scan_data: Dict[NetworkPort, PortScanData],
+        port_scan_data: PortScanDataDict,
         options: Dict,
     ) -> FingerprintData:
         """

monkey/infection_monkey/i_puppet/target_host.py

@@ -1,22 +1,43 @@
 import pprint
+from collections import UserDict
 from ipaddress import IPv4Address
-from typing import Dict, Optional
+from typing import Optional, Set
 
-from pydantic import Field
+from pydantic import Field, validate_arguments
 
 from common import OperatingSystem
-from common.base_models import MutableInfectionMonkeyBaseModel
-from common.types import NetworkPort
+from common.base_models import MutableInfectionMonkeyBaseModel, MutableInfectionMonkeyModelConfig
+from common.types import NetworkPort, PortStatus
 
 from . import PortScanData
 
 
+class PortScanDataDict(UserDict[NetworkPort, PortScanData]):
+    @validate_arguments
+    def __setitem__(self, key: NetworkPort, value: PortScanData):
+        super().__setitem__(key, value)
+
+    @property
+    def closed(self) -> Set[NetworkPort]:
+        return {
+            port
+            for port, port_scan_data in self.data.items()
+            if port_scan_data.status == PortStatus.CLOSED
+        }
+
+
 class TargetHostPorts(MutableInfectionMonkeyBaseModel):
-    tcp_ports: Dict[NetworkPort, PortScanData] = Field(default={})
-    udp_ports: Dict[NetworkPort, PortScanData] = Field(default={})
+    class Config(MutableInfectionMonkeyModelConfig):
+        arbitrary_types_allowed = True
+
+    tcp_ports: PortScanDataDict = Field(default_factory=PortScanDataDict)
+    udp_ports: PortScanDataDict = Field(default_factory=PortScanDataDict)
 
 
 class TargetHost(MutableInfectionMonkeyBaseModel):
+    class Config(MutableInfectionMonkeyModelConfig):
+        json_encoders = {PortScanDataDict: dict}
+
     ip: IPv4Address
     operating_system: Optional[OperatingSystem] = Field(default=None)
     icmp: bool = Field(default=False)

monkey/infection_monkey/master/ip_scan_results.py

@@ -1,14 +1,13 @@
 from dataclasses import dataclass
 from typing import Dict
 
-from common.types import NetworkPort
-from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData
+from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanDataDict
 
 FingerprinterName = str
 
 
 @dataclass
 class IPScanResults:
     ping_scan_data: PingScanData
-    port_scan_data: Dict[NetworkPort, PortScanData]
+    port_scan_data: PortScanDataDict
     fingerprint_data: Dict[FingerprinterName, FingerprintData]

monkey/infection_monkey/master/ip_scanner.py

@@ -8,8 +8,8 @@
     NetworkScanConfiguration,
     PluginConfiguration,
 )
-from common.types import Event, NetworkPort, PortStatus
-from infection_monkey.i_puppet import FingerprintData, IPuppet, PingScanData, PortScanData
+from common.types import Event, PortStatus
+from infection_monkey.i_puppet import FingerprintData, IPuppet, PingScanData, PortScanDataDict
 from infection_monkey.network import NetworkAddress
 from infection_monkey.utils.threading import interruptible_iter, run_worker_threads
 
@@ -86,15 +86,15 @@ def _scan_addresses(
             )
 
     @staticmethod
-    def port_scan_found_open_port(port_scan_data: Dict[NetworkPort, PortScanData]):
+    def port_scan_found_open_port(port_scan_data: PortScanDataDict):
         return any(psd.status == PortStatus.OPEN for psd in port_scan_data.values())
 
     def _run_fingerprinters(
         self,
         ip: str,
         fingerprinters: Sequence[PluginConfiguration],
         ping_scan_data: PingScanData,
-        port_scan_data: Dict[NetworkPort, PortScanData],
+        port_scan_data: PortScanDataDict,
         stop: Event,
     ) -> Dict[str, FingerprintData]:
         fingerprint_data = {}

monkey/infection_monkey/network_scanning/tcp_scanner.py

@@ -11,22 +11,22 @@
 from common.agent_events import TCPScanEvent
 from common.event_queue import IAgentEventQueue
 from common.types import NetworkPort, PortStatus
-from infection_monkey.i_puppet import PortScanData
+from infection_monkey.i_puppet import PortScanData, PortScanDataDict
 from infection_monkey.network.tools import BANNER_READ, DEFAULT_TIMEOUT
 from infection_monkey.utils.ids import get_agent_id
 
 logger = logging.getLogger(__name__)
 
 POLL_INTERVAL = 0.5
-EMPTY_PORT_SCAN: Dict[NetworkPort, PortScanData] = {}
+EMPTY_PORT_SCAN = PortScanDataDict()
 
 
 def scan_tcp_ports(
     host: str,
     ports_to_scan: Collection[NetworkPort],
     timeout: float,
     agent_event_queue: IAgentEventQueue,
-) -> Dict[NetworkPort, PortScanData]:
+) -> PortScanDataDict:
     try:
         return _scan_tcp_ports(host, ports_to_scan, timeout, agent_event_queue)
     except Exception:
@@ -39,7 +39,7 @@ def _scan_tcp_ports(
     ports_to_scan: Collection[NetworkPort],
     timeout: float,
     agent_event_queue: IAgentEventQueue,
-) -> Dict[NetworkPort, PortScanData]:
+) -> PortScanDataDict:
     event_timestamp, open_ports = _check_tcp_ports(host, ports_to_scan, timeout)
 
     port_scan_data = _build_port_scan_data(ports_to_scan, open_ports)
@@ -51,9 +51,9 @@ def _scan_tcp_ports(
 
 
 def _generate_tcp_scan_event(
-    host: str, port_scan_data: Dict[NetworkPort, PortScanData], event_timestamp: float
+    host: str, port_scan_data_dict: PortScanDataDict, event_timestamp: float
 ):
-    port_statuses = {port: psd.status for port, psd in port_scan_data.items()}
+    port_statuses = {port: psd.status for port, psd in port_scan_data_dict.items()}
 
     return TCPScanEvent(
         source=get_agent_id(),
@@ -65,8 +65,8 @@ def _generate_tcp_scan_event(
 
 def _build_port_scan_data(
     ports_to_scan: Iterable[NetworkPort], open_ports: Mapping[NetworkPort, str]
-) -> Dict[NetworkPort, PortScanData]:
-    port_scan_data = {}
+) -> PortScanDataDict:
+    port_scan_data = PortScanDataDict()
     for port in ports_to_scan:
         if port in open_ports:
             banner = open_ports[port]

monkey/infection_monkey/puppet/puppet.py

@@ -13,7 +13,7 @@
     IncompatibleOperatingSystemError,
     IPuppet,
     PingScanData,
-    PortScanData,
+    PortScanDataDict,
     TargetHost,
 )
 from infection_monkey.puppet import PluginCompatabilityVerifier
@@ -50,15 +50,15 @@ def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData:
 
     def scan_tcp_ports(
         self, host: str, ports: Sequence[NetworkPort], timeout: float = CONNECTION_TIMEOUT
-    ) -> Dict[NetworkPort, PortScanData]:
+    ) -> PortScanDataDict:
         return network_scanning.scan_tcp_ports(host, ports, timeout, self._agent_event_queue)
 
     def fingerprint(
         self,
         name: str,
         host: str,
         ping_scan_data: PingScanData,
-        port_scan_data: Dict[NetworkPort, PortScanData],
+        port_scan_data: PortScanDataDict,
         options: Dict,
     ) -> FingerprintData:
         try:

monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_hadoop_exploiter.py

@@ -11,7 +11,13 @@
 from common import OperatingSystem
 from common.types import NetworkPort, NetworkProtocol, NetworkService, PortStatus
 from infection_monkey.exploit.tools import HTTPBytesServer
-from infection_monkey.i_puppet import ExploiterResultData, PortScanData, TargetHost, TargetHostPorts
+from infection_monkey.i_puppet import (
+    ExploiterResultData,
+    PortScanData,
+    PortScanDataDict,
+    TargetHost,
+    TargetHostPorts,
+)
 
 TARGET_IP = IPv4Address("1.1.1.1")
 SERVERS = ["10.10.10.10"]
@@ -187,11 +193,13 @@ def test_exploit_attempt_on_all_discovered_open_http_ports(
         ip=IPv4Address("1.1.1.1"),
         operating_system=OperatingSystem.WINDOWS,
         ports_status=TargetHostPorts(
-            tcp_ports={
-                HTTP_PORT: HTTP_PORT_DATA,
-                CLOSED_PORT: CLOSED_PORT_DATA,
-                HTTPS_PORT: HTTPS_PORT_DATA,
-            }
+            tcp_ports=PortScanDataDict(
+                {
+                    HTTP_PORT: HTTP_PORT_DATA,
+                    CLOSED_PORT: CLOSED_PORT_DATA,
+                    HTTPS_PORT: HTTPS_PORT_DATA,
+                }
+            )
         ),
     )
     hadoop_exploiter.exploit_host(
@@ -219,7 +227,9 @@ def test_exploit_attempt_skips_configured_ports_if_closed(
     target_host = TargetHost(
         ip=IPv4Address("1.1.1.1"),
         operating_system=OperatingSystem.WINDOWS,
-        ports_status=TargetHostPorts(tcp_ports={CLOSED_PORT_80: CLOSED_PORT_80_DATA}),
+        ports_status=TargetHostPorts(
+            tcp_ports=PortScanDataDict({CLOSED_PORT_80: CLOSED_PORT_80_DATA})
+        ),
     )
     hadoop_exploiter.exploit_host(
         target_host=target_host,