From aa5632fd8ddc0e8a02703d9965674eac045c5b30 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Thu, 23 Mar 2023 17:50:06 +0000 Subject: [PATCH 01/10] SMB: Check for SMB ports before running exploiter Issue #2952 PR #3143 --- .../exploiters/smb/src/plugin.py | 15 +++++- .../infection_monkey/i_puppet/target_host.py | 13 +++++- .../exploiters/smb/test_smb_plugin.py | 46 ++++++++++++++++++- 3 files changed, 69 insertions(+), 5 deletions(-) diff --git a/monkey/agent_plugins/exploiters/smb/src/plugin.py b/monkey/agent_plugins/exploiters/smb/src/plugin.py index 0f721452b18..af8167ca6be 100644 --- a/monkey/agent_plugins/exploiters/smb/src/plugin.py +++ b/monkey/agent_plugins/exploiters/smb/src/plugin.py @@ -6,7 +6,7 @@ # common imports from common.event_queue import IAgentEventPublisher -from common.types import Event +from common.types import Event, PortStatus from common.utils.code_utils import del_key # dependencies to get rid of or internalize @@ -17,13 +17,19 @@ generate_brute_force_credentials, ) from infection_monkey.exploit.tools.helpers import get_agent_dst_path -from infection_monkey.i_puppet import ExploiterResultData, TargetHost +from infection_monkey.i_puppet import ExploiterResultData, PortScanData, TargetHost from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository 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 is_open_port(psd: PortScanData) -> bool: + return psd.status != PortStatus.CLOSED + + logger = logging.getLogger(__name__) @@ -67,6 +73,11 @@ def run( logger.exception(msg) return ExploiterResultData(error_message=msg) + if len(host.filter_selected_tcp_ports(SMB_PORTS, is_open_port)) == 0: + msg = f"Host {host.ip} has no open SMB ports" + logger.debug(msg) + return ExploiterResultData(error_message=msg) + command_builder = partial( build_smb_command, servers, diff --git a/monkey/infection_monkey/i_puppet/target_host.py b/monkey/infection_monkey/i_puppet/target_host.py index 2a43aa655c7..375d3b2ab53 100644 --- a/monkey/infection_monkey/i_puppet/target_host.py +++ b/monkey/infection_monkey/i_puppet/target_host.py @@ -1,6 +1,6 @@ import pprint from ipaddress import IPv4Address -from typing import Dict, Optional +from typing import Callable, Dict, Optional, Sequence from pydantic import Field @@ -10,6 +10,8 @@ from . import PortScanData +FilterPortFunc = Callable[[PortScanData], bool] + class TargetHostPorts(MutableInfectionMonkeyBaseModel): tcp_ports: Dict[NetworkPort, PortScanData] = Field(default={}) @@ -22,6 +24,15 @@ class TargetHost(MutableInfectionMonkeyBaseModel): icmp: bool = Field(default=False) ports_status: TargetHostPorts = Field(default=TargetHostPorts()) + def filter_selected_tcp_ports( + self, ports: Sequence[NetworkPort], filter: FilterPortFunc + ) -> Sequence[PortScanData]: + return [ + p + for p in ports + if p in self.ports_status.tcp_ports and filter(self.ports_status.tcp_ports[p]) + ] + def __hash__(self): return hash(self.ip) diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py index b5097773cb2..5e274699361 100644 --- a/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py +++ b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py @@ -7,14 +7,24 @@ from agent_plugins.exploiters.smb.src.plugin import Plugin from common import OperatingSystem +from common.types import NetworkPort, PortStatus from infection_monkey.exploit.tools import BruteForceExploiter -from infection_monkey.i_puppet import ExploiterResultData, TargetHost +from infection_monkey.i_puppet import ExploiterResultData, PortScanData, TargetHost, TargetHostPorts from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository AGENT_ID = UUID("5c145d4e-ec61-44f7-998e-17477112f50f") BAD_SMB_OPTIONS_DICT = {"blah": "blah"} TARGET_IP = IPv4Address("1.1.1.1") -TARGET_HOST = TargetHost(ip=TARGET_IP, operating_system=OperatingSystem.WINDOWS) +SMB_PORT = NetworkPort(139) +OPEN_SMB_PORTS = TargetHostPorts( + tcp_ports={SMB_PORT: PortScanData(port=SMB_PORT, status=PortStatus.OPEN)} +) +CLOSED_SMB_PORTS = TargetHostPorts() +TARGET_HOST = TargetHost( + ip=TARGET_IP, + operating_system=OperatingSystem.WINDOWS, + ports_status=OPEN_SMB_PORTS, +) SERVERS = ["10.10.10.10"] EXPLOITER_RESULT_DATA = ExploiterResultData(True, False, error_message="Test error") @@ -70,6 +80,38 @@ def test_run__fails_on_bad_options(plugin: Plugin): assert not result.propagation_success +def test_run__fails_if_no_open_smb_ports(plugin: Plugin): + host = TARGET_HOST.copy(deep=True) + host.ports_status = CLOSED_SMB_PORTS + result = plugin.run( + host=host, + servers=SERVERS, + current_depth=1, + options={}, + interrupt=Event(), + ) + + assert not result.exploitation_success + assert not result.propagation_success + + +@pytest.mark.parametrize("port", [NetworkPort(445), NetworkPort(139)]) +def test_run__proceeds_if_open_smb_ports(plugin: Plugin, port: NetworkPort): + host = TARGET_HOST.copy(deep=True) + host.ports_status = TargetHostPorts( + tcp_ports={port: PortScanData(port=port, status=PortStatus.OPEN)} + ) + result = plugin.run( + host=host, + servers=SERVERS, + current_depth=1, + options={}, + interrupt=Event(), + ) + + assert result == EXPLOITER_RESULT_DATA + + def test_run__returns_exploiter_result_data(plugin: Plugin): result = plugin.run( host=TARGET_HOST, From 2882b21cccbd2f3c6665c753c827273b98a42a92 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 23 Mar 2023 19:51:02 -0400 Subject: [PATCH 02/10] SMB: Refactor port checking logic --- .../exploiters/smb/src/plugin.py | 19 ++- .../exploiters/smb/test_smb_plugin.py | 136 +++++++++++++----- 2 files changed, 115 insertions(+), 40 deletions(-) diff --git a/monkey/agent_plugins/exploiters/smb/src/plugin.py b/monkey/agent_plugins/exploiters/smb/src/plugin.py index af8167ca6be..1aa956d59b2 100644 --- a/monkey/agent_plugins/exploiters/smb/src/plugin.py +++ b/monkey/agent_plugins/exploiters/smb/src/plugin.py @@ -17,7 +17,7 @@ generate_brute_force_credentials, ) from infection_monkey.exploit.tools.helpers import get_agent_dst_path -from infection_monkey.i_puppet import ExploiterResultData, PortScanData, TargetHost +from infection_monkey.i_puppet import ExploiterResultData, TargetHost from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository from .smb_command_builder import build_smb_command @@ -26,8 +26,15 @@ from .smb_remote_access_client_factory import SMBRemoteAccessClientFactory -def is_open_port(psd: PortScanData) -> bool: - return psd.status != PortStatus.CLOSED +def should_attempt_exploit(host: TargetHost) -> bool: + for smb_port in SMB_PORTS: + if smb_port not in host.ports_status.tcp_ports: + return True + + if host.ports_status.tcp_ports[smb_port].status == PortStatus.OPEN: + return True + + return False logger = logging.getLogger(__name__) @@ -73,10 +80,12 @@ def run( logger.exception(msg) return ExploiterResultData(error_message=msg) - if len(host.filter_selected_tcp_ports(SMB_PORTS, is_open_port)) == 0: + if not should_attempt_exploit(host): msg = f"Host {host.ip} has no open SMB ports" logger.debug(msg) - return ExploiterResultData(error_message=msg) + return ExploiterResultData( + exploitation_success=False, propagation_success=False, error_message=msg + ) command_builder = partial( build_smb_command, diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py index 5e274699361..b5c73a174f2 100644 --- a/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py +++ b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py @@ -1,10 +1,11 @@ from ipaddress import IPv4Address from threading import Event +from typing import Dict from unittest.mock import MagicMock from uuid import UUID import pytest -from agent_plugins.exploiters.smb.src.plugin import Plugin +from agent_plugins.exploiters.smb.src.plugin import SMB_PORTS, Plugin from common import OperatingSystem from common.types import NetworkPort, PortStatus @@ -15,31 +16,26 @@ AGENT_ID = UUID("5c145d4e-ec61-44f7-998e-17477112f50f") BAD_SMB_OPTIONS_DICT = {"blah": "blah"} TARGET_IP = IPv4Address("1.1.1.1") -SMB_PORT = NetworkPort(139) OPEN_SMB_PORTS = TargetHostPorts( - tcp_ports={SMB_PORT: PortScanData(port=SMB_PORT, status=PortStatus.OPEN)} -) -CLOSED_SMB_PORTS = TargetHostPorts() -TARGET_HOST = TargetHost( - ip=TARGET_IP, - operating_system=OperatingSystem.WINDOWS, - ports_status=OPEN_SMB_PORTS, + tcp_ports={p: PortScanData(port=p, status=PortStatus.OPEN) for p in SMB_PORTS} ) +EMPTY_TARGET_HOST_PORTS = TargetHostPorts() SERVERS = ["10.10.10.10"] EXPLOITER_RESULT_DATA = ExploiterResultData(True, False, error_message="Test error") @pytest.fixture -def propagation_credentials_repository(): - return MagicMock(spec=IPropagationCredentialsRepository) - +def target_host() -> TargetHost: + return TargetHost( + ip=TARGET_IP, + operating_system=OperatingSystem.WINDOWS, + ports_status=OPEN_SMB_PORTS, + ) -class MockSMBExploiter(BruteForceExploiter): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - def exploit_host(self, *args, **kwargs) -> ExploiterResultData: - return EXPLOITER_RESULT_DATA +@pytest.fixture +def propagation_credentials_repository(): + return MagicMock(spec=IPropagationCredentialsRepository) class ErrorRaisingMockSMBExploiter(BruteForceExploiter): @@ -50,12 +46,22 @@ def exploit_host(self, *args, **kwargs) -> ExploiterResultData: raise Exception("Test error") +@pytest.fixture +def mock_smb_exploiter(): + exploiter = MagicMock(spec=BruteForceExploiter) + exploiter.exploit_host.return_value = EXPLOITER_RESULT_DATA + return exploiter + + @pytest.fixture def plugin( - monkeypatch, propagation_credentials_repository: IPropagationCredentialsRepository + monkeypatch, + propagation_credentials_repository: IPropagationCredentialsRepository, + mock_smb_exploiter: BruteForceExploiter, ) -> Plugin: monkeypatch.setattr( - "agent_plugins.exploiters.smb.src.plugin.BruteForceExploiter", MockSMBExploiter + "agent_plugins.exploiters.smb.src.plugin.BruteForceExploiter", + lambda *args, **kwargs: mock_smb_exploiter, ) return Plugin( @@ -67,9 +73,9 @@ def plugin( ) -def test_run__fails_on_bad_options(plugin: Plugin): +def test_run__fails_on_bad_options(plugin: Plugin, target_host: TargetHost): result = plugin.run( - host=TARGET_HOST, + host=target_host, servers=SERVERS, current_depth=1, options=BAD_SMB_OPTIONS_DICT, @@ -80,9 +86,22 @@ def test_run__fails_on_bad_options(plugin: Plugin): assert not result.propagation_success -def test_run__fails_if_no_open_smb_ports(plugin: Plugin): - host = TARGET_HOST.copy(deep=True) - host.ports_status = CLOSED_SMB_PORTS +@pytest.mark.parametrize( + "tcp_port_status", + ( + {SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED)}, + {SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED)}, + {}, + ), +) +def test_run__attempts_exploit_if_port_status_unknown( + plugin: Plugin, + mock_smb_exploiter: BruteForceExploiter, + target_host: TargetHost, + tcp_port_status: Dict[NetworkPort, PortScanData], +): + host = target_host + host.ports_status.tcp_ports = tcp_port_status result = plugin.run( host=host, servers=SERVERS, @@ -91,16 +110,37 @@ def test_run__fails_if_no_open_smb_ports(plugin: Plugin): interrupt=Event(), ) - assert not result.exploitation_success - assert not result.propagation_success + mock_smb_exploiter.exploit_host.assert_called_once() + assert result == EXPLOITER_RESULT_DATA -@pytest.mark.parametrize("port", [NetworkPort(445), NetworkPort(139)]) -def test_run__proceeds_if_open_smb_ports(plugin: Plugin, port: NetworkPort): - host = TARGET_HOST.copy(deep=True) - host.ports_status = TargetHostPorts( - tcp_ports={port: PortScanData(port=port, status=PortStatus.OPEN)} - ) +@pytest.mark.parametrize( + "tcp_port_status", + ( + {SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN)}, + {SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN)}, + { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN), + }, + { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED), + }, + { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN), + }, + ), +) +def test_run__attempts_exploit_if_port_status_open( + plugin: Plugin, + mock_smb_exploiter: BruteForceExploiter, + target_host: TargetHost, + tcp_port_status: Dict[NetworkPort, PortScanData], +): + host = target_host + host.ports_status.tcp_ports = tcp_port_status result = plugin.run( host=host, servers=SERVERS, @@ -109,12 +149,37 @@ def test_run__proceeds_if_open_smb_ports(plugin: Plugin, port: NetworkPort): interrupt=Event(), ) + mock_smb_exploiter.exploit_host.assert_called_once() assert result == EXPLOITER_RESULT_DATA -def test_run__returns_exploiter_result_data(plugin: Plugin): +def test_run__skips_exploit_if_port_status_closed( + plugin: Plugin, + mock_smb_exploiter: BruteForceExploiter, + target_host: TargetHost, +): + host = target_host + host.ports_status.tcp_ports = { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED), + } + + result = plugin.run( + host=host, + servers=SERVERS, + current_depth=1, + options={}, + interrupt=Event(), + ) + + mock_smb_exploiter.exploit_host.assert_not_called() + assert result.exploitation_success is False + assert result.propagation_success is False + + +def test_run__returns_exploiter_result_data(plugin: Plugin, target_host: TargetHost): result = plugin.run( - host=TARGET_HOST, + host=target_host, servers=SERVERS, current_depth=1, options={}, @@ -128,6 +193,7 @@ def test_run__exploit_host_raises_exception( monkeypatch, plugin: Plugin, propagation_credentials_repository: IPropagationCredentialsRepository, + target_host: TargetHost, ): monkeypatch.setattr( "agent_plugins.exploiters.smb.src.plugin.BruteForceExploiter", @@ -142,7 +208,7 @@ def test_run__exploit_host_raises_exception( propagation_credentials_repository=propagation_credentials_repository, ) result = plugin.run( - host=TARGET_HOST, + host=target_host, servers=SERVERS, current_depth=1, options={}, From facf76ad25fe135b96b4b0ededfd5d1f1ac3bb02 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 23 Mar 2023 19:53:41 -0400 Subject: [PATCH 03/10] Agent: Remote disused TargetHost.filter_selected_tcp_ports() --- monkey/infection_monkey/i_puppet/target_host.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/monkey/infection_monkey/i_puppet/target_host.py b/monkey/infection_monkey/i_puppet/target_host.py index 375d3b2ab53..2a43aa655c7 100644 --- a/monkey/infection_monkey/i_puppet/target_host.py +++ b/monkey/infection_monkey/i_puppet/target_host.py @@ -1,6 +1,6 @@ import pprint from ipaddress import IPv4Address -from typing import Callable, Dict, Optional, Sequence +from typing import Dict, Optional from pydantic import Field @@ -10,8 +10,6 @@ from . import PortScanData -FilterPortFunc = Callable[[PortScanData], bool] - class TargetHostPorts(MutableInfectionMonkeyBaseModel): tcp_ports: Dict[NetworkPort, PortScanData] = Field(default={}) @@ -24,15 +22,6 @@ class TargetHost(MutableInfectionMonkeyBaseModel): icmp: bool = Field(default=False) ports_status: TargetHostPorts = Field(default=TargetHostPorts()) - def filter_selected_tcp_ports( - self, ports: Sequence[NetworkPort], filter: FilterPortFunc - ) -> Sequence[PortScanData]: - return [ - p - for p in ports - if p in self.ports_status.tcp_ports and filter(self.ports_status.tcp_ports[p]) - ] - def __hash__(self): return hash(self.ip) From 56ce2b9290c9c39c4fd40bc7f1a3309eb1fab625 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 23 Mar 2023 20:13:37 -0400 Subject: [PATCH 04/10] Agent: Add TargetHostPorts.get_closed_tcp_ports() --- monkey/infection_monkey/i_puppet/target_host.py | 12 ++++++++++-- .../infection_monkey/i_puppet/test_target_host.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) create mode 100644 monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py diff --git a/monkey/infection_monkey/i_puppet/target_host.py b/monkey/infection_monkey/i_puppet/target_host.py index 2a43aa655c7..6739752e939 100644 --- a/monkey/infection_monkey/i_puppet/target_host.py +++ b/monkey/infection_monkey/i_puppet/target_host.py @@ -1,12 +1,12 @@ import pprint from ipaddress import IPv4Address -from typing import Dict, Optional +from typing import Dict, Optional, Set from pydantic import Field from common import OperatingSystem from common.base_models import MutableInfectionMonkeyBaseModel -from common.types import NetworkPort +from common.types import NetworkPort, PortStatus from . import PortScanData @@ -15,6 +15,14 @@ class TargetHostPorts(MutableInfectionMonkeyBaseModel): tcp_ports: Dict[NetworkPort, PortScanData] = Field(default={}) udp_ports: Dict[NetworkPort, PortScanData] = Field(default={}) + @property + def closed_tcp_ports(self) -> Set[NetworkPort]: + return { + port + for port, port_scan_data in self.tcp_ports.items() + if port_scan_data.status == PortStatus.CLOSED + } + class TargetHost(MutableInfectionMonkeyBaseModel): ip: IPv4Address diff --git a/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py new file mode 100644 index 00000000000..43a56580cc0 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py @@ -0,0 +1,15 @@ +from common.types import NetworkPort, PortStatus +from infection_monkey.i_puppet import PortScanData, TargetHostPorts + + +def test_closed_tcp_ports(): + expected_closed_ports = {NetworkPort(2), NetworkPort(4)} + tcp_ports = { + NetworkPort(1): PortScanData(port=NetworkPort(1), status=PortStatus.OPEN), + NetworkPort(2): PortScanData(port=NetworkPort(2), status=PortStatus.CLOSED), + NetworkPort(3): PortScanData(port=NetworkPort(3), status=PortStatus.OPEN), + NetworkPort(4): PortScanData(port=NetworkPort(4), status=PortStatus.CLOSED), + } + thp = TargetHostPorts(tcp_ports=tcp_ports) + + assert thp.closed_tcp_ports == expected_closed_ports From fd463eda1b260bf62c17ce214558641b5d43ccd4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 23 Mar 2023 20:14:28 -0400 Subject: [PATCH 05/10] SMB: Use get_closed_tcp_ports() in should_attempt_exploit() --- monkey/agent_plugins/exploiters/smb/src/plugin.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/monkey/agent_plugins/exploiters/smb/src/plugin.py b/monkey/agent_plugins/exploiters/smb/src/plugin.py index 1aa956d59b2..00cf9eab1e6 100644 --- a/monkey/agent_plugins/exploiters/smb/src/plugin.py +++ b/monkey/agent_plugins/exploiters/smb/src/plugin.py @@ -6,7 +6,7 @@ # common imports from common.event_queue import IAgentEventPublisher -from common.types import Event, PortStatus +from common.types import Event from common.utils.code_utils import del_key # dependencies to get rid of or internalize @@ -27,14 +27,8 @@ def should_attempt_exploit(host: TargetHost) -> bool: - for smb_port in SMB_PORTS: - if smb_port not in host.ports_status.tcp_ports: - return True - - if host.ports_status.tcp_ports[smb_port].status == PortStatus.OPEN: - return True - - return False + closed_tcp_ports = host.ports_status.closed_tcp_ports + return not all([p in closed_tcp_ports for p in SMB_PORTS]) logger = logging.getLogger(__name__) From 39d179d5823883e6f0cad976188d3b92151055fa Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 24 Mar 2023 08:35:48 -0400 Subject: [PATCH 06/10] Agent: Add PortScanDataDict --- monkey/infection_monkey/i_puppet/__init__.py | 2 +- .../infection_monkey/i_puppet/target_host.py | 9 ++- .../i_puppet/test_target_host.py | 68 +++++++++++++++++-- 3 files changed, 71 insertions(+), 8 deletions(-) diff --git a/monkey/infection_monkey/i_puppet/__init__.py b/monkey/infection_monkey/i_puppet/__init__.py index 9def052c20f..9a0dfaf4b68 100644 --- a/monkey/infection_monkey/i_puppet/__init__.py +++ b/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 diff --git a/monkey/infection_monkey/i_puppet/target_host.py b/monkey/infection_monkey/i_puppet/target_host.py index 6739752e939..f7714c8fd08 100644 --- a/monkey/infection_monkey/i_puppet/target_host.py +++ b/monkey/infection_monkey/i_puppet/target_host.py @@ -1,8 +1,9 @@ import pprint +from collections import UserDict from ipaddress import IPv4Address from typing import Dict, Optional, Set -from pydantic import Field +from pydantic import Field, validate_arguments from common import OperatingSystem from common.base_models import MutableInfectionMonkeyBaseModel @@ -11,6 +12,12 @@ from . import PortScanData +class PortScanDataDict(UserDict[NetworkPort, PortScanData]): + @validate_arguments + def __setitem__(self, key: NetworkPort, value: PortScanData): + super().__setitem__(key, value) + + class TargetHostPorts(MutableInfectionMonkeyBaseModel): tcp_ports: Dict[NetworkPort, PortScanData] = Field(default={}) udp_ports: Dict[NetworkPort, PortScanData] = Field(default={}) diff --git a/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py index 43a56580cc0..b9cb031a09f 100644 --- a/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py +++ b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py @@ -1,14 +1,70 @@ +import pytest + from common.types import NetworkPort, PortStatus -from infection_monkey.i_puppet import PortScanData, TargetHostPorts +from infection_monkey.i_puppet import PortScanData, PortScanDataDict, TargetHostPorts + + +def test_port_scan_data_dict__constructor(): + input_dict = { + NetworkPort(1): PortScanData(port=1, status=PortStatus.OPEN), + NetworkPort(2): PortScanData(port=2, status=PortStatus.CLOSED), + NetworkPort(3): PortScanData(port=3, status=PortStatus.OPEN), + } + expected_port_scan_data_dict = { + 1: PortScanData(port=1, status=PortStatus.OPEN), + 2: PortScanData(port=2, status=PortStatus.CLOSED), + 3: PortScanData(port=3, status=PortStatus.OPEN), + } + + port_scan_data_dict: PortScanDataDict = PortScanDataDict(input_dict) + + assert port_scan_data_dict == expected_port_scan_data_dict + + +def test_port_scan_data_dict__set(): + expected_port_scan_data_dict = { + 1: PortScanData(port=1, status=PortStatus.OPEN), + 2: PortScanData(port=2, status=PortStatus.CLOSED), + } + + port_scan_data_dict = PortScanDataDict() + port_scan_data_dict[1] = PortScanData(port=1, status=PortStatus.OPEN) + port_scan_data_dict[2] = PortScanData(port=2, status=PortStatus.CLOSED) + + assert port_scan_data_dict == expected_port_scan_data_dict + + +INVALID_PORTS = (-1, 65536, "string", None, "22.2") +VALID_PORT_SCAN_DATA = PortScanData(port=1, status=PortStatus.OPEN) + + +@pytest.mark.parametrize("invalid_port", INVALID_PORTS) +def test_port_scan_data_dict_constructor__invalid_port(invalid_port): + with pytest.raises((ValueError, TypeError)): + PortScanDataDict({invalid_port: VALID_PORT_SCAN_DATA}) + + +@pytest.mark.parametrize("invalid_port", INVALID_PORTS) +def test_port_scan_data_dict_update__invalid_port(invalid_port): + port_scan_data_dict = PortScanDataDict() + with pytest.raises((ValueError, TypeError)): + port_scan_data_dict.update({invalid_port: VALID_PORT_SCAN_DATA}) + + +@pytest.mark.parametrize("invalid_port", INVALID_PORTS) +def test_port_scan_data_dict_set__invalid_port(invalid_port): + port_scan_data_dict = PortScanDataDict() + with pytest.raises((ValueError, TypeError)): + port_scan_data_dict[invalid_port] = VALID_PORT_SCAN_DATA def test_closed_tcp_ports(): - expected_closed_ports = {NetworkPort(2), NetworkPort(4)} + expected_closed_ports = {2, 4} tcp_ports = { - NetworkPort(1): PortScanData(port=NetworkPort(1), status=PortStatus.OPEN), - NetworkPort(2): PortScanData(port=NetworkPort(2), status=PortStatus.CLOSED), - NetworkPort(3): PortScanData(port=NetworkPort(3), status=PortStatus.OPEN), - NetworkPort(4): PortScanData(port=NetworkPort(4), status=PortStatus.CLOSED), + 1: PortScanData(port=1, status=PortStatus.OPEN), + 2: PortScanData(port=2, status=PortStatus.CLOSED), + 3: PortScanData(port=3, status=PortStatus.OPEN), + 4: PortScanData(port=4, status=PortStatus.CLOSED), } thp = TargetHostPorts(tcp_ports=tcp_ports) From da5810f048a6f2fb2f179c183bf16b8eec789464 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 24 Mar 2023 09:18:09 -0400 Subject: [PATCH 07/10] Agent: Use PortScanDataDict for TargetHostPorts fields --- monkey/infection_monkey/i_puppet/i_puppet.py | 8 +-- .../infection_monkey/i_puppet/target_host.py | 14 ++-- .../master/ip_scan_results.py | 5 +- monkey/infection_monkey/master/ip_scanner.py | 8 +-- .../network_scanning/tcp_scanner.py | 16 ++--- monkey/infection_monkey/puppet/puppet.py | 6 +- .../hadoop/test_hadoop_exploiter.py | 24 +++++-- .../exploiters/smb/test_smb_plugin.py | 67 +++++++++++-------- .../i_puppet/test_target_host.py | 14 ++-- .../infection_monkey/master/mock_puppet.py | 5 +- 10 files changed, 99 insertions(+), 68 deletions(-) diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py index 78ce39b50f8..1e78a160b34 100644 --- a/monkey/infection_monkey/i_puppet/i_puppet.py +++ b/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: """ diff --git a/monkey/infection_monkey/i_puppet/target_host.py b/monkey/infection_monkey/i_puppet/target_host.py index f7714c8fd08..69c7731255e 100644 --- a/monkey/infection_monkey/i_puppet/target_host.py +++ b/monkey/infection_monkey/i_puppet/target_host.py @@ -1,12 +1,12 @@ import pprint from collections import UserDict from ipaddress import IPv4Address -from typing import Dict, Optional, Set +from typing import Optional, Set from pydantic import Field, validate_arguments from common import OperatingSystem -from common.base_models import MutableInfectionMonkeyBaseModel +from common.base_models import MutableInfectionMonkeyBaseModel, MutableInfectionMonkeyModelConfig from common.types import NetworkPort, PortStatus from . import PortScanData @@ -19,8 +19,11 @@ def __setitem__(self, key: NetworkPort, value: PortScanData): 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) @property def closed_tcp_ports(self) -> Set[NetworkPort]: @@ -32,6 +35,9 @@ def closed_tcp_ports(self) -> Set[NetworkPort]: class TargetHost(MutableInfectionMonkeyBaseModel): + class Config(MutableInfectionMonkeyModelConfig): + json_encoders = {PortScanDataDict: dict} + ip: IPv4Address operating_system: Optional[OperatingSystem] = Field(default=None) icmp: bool = Field(default=False) diff --git a/monkey/infection_monkey/master/ip_scan_results.py b/monkey/infection_monkey/master/ip_scan_results.py index 1c50797c79c..5b393518c28 100644 --- a/monkey/infection_monkey/master/ip_scan_results.py +++ b/monkey/infection_monkey/master/ip_scan_results.py @@ -1,8 +1,7 @@ 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 @@ -10,5 +9,5 @@ @dataclass class IPScanResults: ping_scan_data: PingScanData - port_scan_data: Dict[NetworkPort, PortScanData] + port_scan_data: PortScanDataDict fingerprint_data: Dict[FingerprinterName, FingerprintData] diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index f45ff068098..8e88e0434ae 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/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,7 +86,7 @@ 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( @@ -94,7 +94,7 @@ def _run_fingerprinters( 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 = {} diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index 6a6bbd8605f..217041de798 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -11,14 +11,14 @@ 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( @@ -26,7 +26,7 @@ def scan_tcp_ports( 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] diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 0f55934efbf..11313f859e5 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -13,7 +13,7 @@ IncompatibleOperatingSystemError, IPuppet, PingScanData, - PortScanData, + PortScanDataDict, TargetHost, ) from infection_monkey.puppet import PluginCompatabilityVerifier @@ -50,7 +50,7 @@ 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( @@ -58,7 +58,7 @@ def fingerprint( name: str, host: str, ping_scan_data: PingScanData, - port_scan_data: Dict[NetworkPort, PortScanData], + port_scan_data: PortScanDataDict, options: Dict, ) -> FingerprintData: try: diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_hadoop_exploiter.py b/monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_hadoop_exploiter.py index eddf400e1dd..10bf6b1498e 100644 --- a/monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_hadoop_exploiter.py +++ b/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, diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py index b5c73a174f2..704e150b312 100644 --- a/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py +++ b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py @@ -1,6 +1,5 @@ from ipaddress import IPv4Address from threading import Event -from typing import Dict from unittest.mock import MagicMock from uuid import UUID @@ -8,16 +7,22 @@ from agent_plugins.exploiters.smb.src.plugin import SMB_PORTS, Plugin from common import OperatingSystem -from common.types import NetworkPort, PortStatus +from common.types import PortStatus from infection_monkey.exploit.tools import BruteForceExploiter -from infection_monkey.i_puppet import ExploiterResultData, PortScanData, TargetHost, TargetHostPorts +from infection_monkey.i_puppet import ( + ExploiterResultData, + PortScanData, + PortScanDataDict, + TargetHost, + TargetHostPorts, +) from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository AGENT_ID = UUID("5c145d4e-ec61-44f7-998e-17477112f50f") BAD_SMB_OPTIONS_DICT = {"blah": "blah"} TARGET_IP = IPv4Address("1.1.1.1") OPEN_SMB_PORTS = TargetHostPorts( - tcp_ports={p: PortScanData(port=p, status=PortStatus.OPEN) for p in SMB_PORTS} + tcp_ports=PortScanDataDict({p: PortScanData(port=p, status=PortStatus.OPEN) for p in SMB_PORTS}) ) EMPTY_TARGET_HOST_PORTS = TargetHostPorts() SERVERS = ["10.10.10.10"] @@ -89,16 +94,16 @@ def test_run__fails_on_bad_options(plugin: Plugin, target_host: TargetHost): @pytest.mark.parametrize( "tcp_port_status", ( - {SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED)}, - {SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED)}, - {}, + PortScanDataDict({SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED)}), + PortScanDataDict({SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED)}), + PortScanDataDict({}), ), ) def test_run__attempts_exploit_if_port_status_unknown( plugin: Plugin, mock_smb_exploiter: BruteForceExploiter, target_host: TargetHost, - tcp_port_status: Dict[NetworkPort, PortScanData], + tcp_port_status: PortScanDataDict, ): host = target_host host.ports_status.tcp_ports = tcp_port_status @@ -117,27 +122,33 @@ def test_run__attempts_exploit_if_port_status_unknown( @pytest.mark.parametrize( "tcp_port_status", ( - {SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN)}, - {SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN)}, - { - SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED), - SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN), - }, - { - SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN), - SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED), - }, - { - SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN), - SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN), - }, + PortScanDataDict({SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN)}), + PortScanDataDict({SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN)}), + PortScanDataDict( + { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN), + } + ), + PortScanDataDict( + { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED), + } + ), + PortScanDataDict( + { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.OPEN), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.OPEN), + } + ), ), ) def test_run__attempts_exploit_if_port_status_open( plugin: Plugin, mock_smb_exploiter: BruteForceExploiter, target_host: TargetHost, - tcp_port_status: Dict[NetworkPort, PortScanData], + tcp_port_status: PortScanDataDict, ): host = target_host host.ports_status.tcp_ports = tcp_port_status @@ -159,10 +170,12 @@ def test_run__skips_exploit_if_port_status_closed( target_host: TargetHost, ): host = target_host - host.ports_status.tcp_ports = { - SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED), - SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED), - } + host.ports_status.tcp_ports = PortScanDataDict( + { + SMB_PORTS[0]: PortScanData(port=SMB_PORTS[0], status=PortStatus.CLOSED), + SMB_PORTS[1]: PortScanData(port=SMB_PORTS[1], status=PortStatus.CLOSED), + } + ) result = plugin.run( host=host, diff --git a/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py index b9cb031a09f..6460abcd3da 100644 --- a/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py +++ b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py @@ -60,12 +60,14 @@ def test_port_scan_data_dict_set__invalid_port(invalid_port): def test_closed_tcp_ports(): expected_closed_ports = {2, 4} - tcp_ports = { - 1: PortScanData(port=1, status=PortStatus.OPEN), - 2: PortScanData(port=2, status=PortStatus.CLOSED), - 3: PortScanData(port=3, status=PortStatus.OPEN), - 4: PortScanData(port=4, status=PortStatus.CLOSED), - } + tcp_ports = PortScanDataDict( + { + 1: PortScanData(port=1, status=PortStatus.OPEN), + 2: PortScanData(port=2, status=PortStatus.CLOSED), + 3: PortScanData(port=3, status=PortStatus.OPEN), + 4: PortScanData(port=4, status=PortStatus.CLOSED), + } + ) thp = TargetHostPorts(tcp_ports=tcp_ports) assert thp.closed_tcp_ports == expected_closed_ports diff --git a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py b/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py index fe462f82f89..3184474c756 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py +++ b/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py @@ -4,7 +4,7 @@ from common import OperatingSystem from common.agent_plugins import AgentPluginType from common.credentials import Credentials, LMHash, Password, SSHKeypair, Username -from common.types import Event, NetworkPort, NetworkProtocol, NetworkService, PortStatus +from common.types import Event, NetworkProtocol, NetworkService, PortStatus from infection_monkey.i_puppet import ( DiscoveredService, ExploiterResultData, @@ -13,6 +13,7 @@ IPuppet, PingScanData, PortScanData, + PortScanDataDict, TargetHost, ) @@ -74,7 +75,7 @@ def ping(self, host: str, timeout: float = 1) -> PingScanData: def scan_tcp_ports( self, host: str, ports: Sequence[int], timeout: float = 3 - ) -> Dict[NetworkPort, PortScanData]: + ) -> PortScanDataDict: logger.debug(f"run_scan_tcp_port({host}, {ports}, {timeout})") dot_1_results = { 22: PortScanData(port=22, status=PortStatus.CLOSED), From 8fcfa13cd739ec2a115cf7e69374e74e02ff20d9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 24 Mar 2023 09:22:21 -0400 Subject: [PATCH 08/10] Agent: Add PortScanDataDict.closed property --- monkey/infection_monkey/i_puppet/target_host.py | 8 ++++++++ .../infection_monkey/i_puppet/test_target_host.py | 13 ++++++------- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/i_puppet/target_host.py b/monkey/infection_monkey/i_puppet/target_host.py index 69c7731255e..25186be2722 100644 --- a/monkey/infection_monkey/i_puppet/target_host.py +++ b/monkey/infection_monkey/i_puppet/target_host.py @@ -17,6 +17,14 @@ class PortScanDataDict(UserDict[NetworkPort, PortScanData]): 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): class Config(MutableInfectionMonkeyModelConfig): diff --git a/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py index 6460abcd3da..28691a670f4 100644 --- a/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py +++ b/monkey/tests/unit_tests/infection_monkey/i_puppet/test_target_host.py @@ -1,7 +1,7 @@ import pytest from common.types import NetworkPort, PortStatus -from infection_monkey.i_puppet import PortScanData, PortScanDataDict, TargetHostPorts +from infection_monkey.i_puppet import PortScanData, PortScanDataDict def test_port_scan_data_dict__constructor(): @@ -62,12 +62,11 @@ def test_closed_tcp_ports(): expected_closed_ports = {2, 4} tcp_ports = PortScanDataDict( { - 1: PortScanData(port=1, status=PortStatus.OPEN), - 2: PortScanData(port=2, status=PortStatus.CLOSED), - 3: PortScanData(port=3, status=PortStatus.OPEN), - 4: PortScanData(port=4, status=PortStatus.CLOSED), + NetworkPort(1): PortScanData(port=1, status=PortStatus.OPEN), + NetworkPort(2): PortScanData(port=2, status=PortStatus.CLOSED), + NetworkPort(3): PortScanData(port=3, status=PortStatus.OPEN), + NetworkPort(4): PortScanData(port=4, status=PortStatus.CLOSED), } ) - thp = TargetHostPorts(tcp_ports=tcp_ports) - assert thp.closed_tcp_ports == expected_closed_ports + assert tcp_ports.closed == expected_closed_ports From 0be13fba0dc3725a0f4311f9809b792208b209cd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 24 Mar 2023 09:22:42 -0400 Subject: [PATCH 09/10] SMB: Use PortScanDict.closed property to check port status --- monkey/agent_plugins/exploiters/smb/src/plugin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/agent_plugins/exploiters/smb/src/plugin.py b/monkey/agent_plugins/exploiters/smb/src/plugin.py index 00cf9eab1e6..95ef1a8a16a 100644 --- a/monkey/agent_plugins/exploiters/smb/src/plugin.py +++ b/monkey/agent_plugins/exploiters/smb/src/plugin.py @@ -27,7 +27,7 @@ def should_attempt_exploit(host: TargetHost) -> bool: - closed_tcp_ports = host.ports_status.closed_tcp_ports + closed_tcp_ports = host.ports_status.tcp_ports.closed return not all([p in closed_tcp_ports for p in SMB_PORTS]) From 5ec66098e061c2af5956224b2c3056e2cbabbaec Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 24 Mar 2023 09:23:13 -0400 Subject: [PATCH 10/10] Agent: Remove disused TargetHostPorts.closed_tcp_ports() --- monkey/infection_monkey/i_puppet/target_host.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/monkey/infection_monkey/i_puppet/target_host.py b/monkey/infection_monkey/i_puppet/target_host.py index 25186be2722..b9d3556e3ba 100644 --- a/monkey/infection_monkey/i_puppet/target_host.py +++ b/monkey/infection_monkey/i_puppet/target_host.py @@ -33,14 +33,6 @@ class Config(MutableInfectionMonkeyModelConfig): tcp_ports: PortScanDataDict = Field(default_factory=PortScanDataDict) udp_ports: PortScanDataDict = Field(default_factory=PortScanDataDict) - @property - def closed_tcp_ports(self) -> Set[NetworkPort]: - return { - port - for port, port_scan_data in self.tcp_ports.items() - if port_scan_data.status == PortStatus.CLOSED - } - class TargetHost(MutableInfectionMonkeyBaseModel): class Config(MutableInfectionMonkeyModelConfig):