From 6157ffee7696aff9e9be9651a322f4f0d64c91d6 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 7 Apr 2022 17:55:05 +0300 Subject: [PATCH] Agent: Improve exception handling of tcp, ping and fingerprint scanners --- .../master/automated_master.py | 7 +++-- monkey/infection_monkey/master/ip_scanner.py | 28 ++++--------------- .../network_scanning/__init__.py | 4 +-- .../network_scanning/ping_scanner.py | 14 ++++++++-- .../network_scanning/tcp_scanner.py | 11 +++++++- monkey/infection_monkey/puppet/puppet.py | 16 ++++++++--- .../network_scanning/test_ping.py | 16 +++++------ .../network_scanning/test_tcp_scanning.py | 8 +++--- 8 files changed, 58 insertions(+), 46 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 5586e808c..99f63522a 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -227,8 +227,11 @@ class AutomatedMaster(IMaster): for p in interruptible_iter(plugins, self._stop, interrupted_message): try: callback(p) - except Exception as ex: - logger.debug(f"Exception encountered when running {plugin_type} {p}: {str(ex)}") + except Exception: + logging.exception( + f"Got unhandled exception when running {plugin_type} plugin {p}. " + f"Plugin was passed to {callback}" + ) logger.info(f"Finished running {plugin_type}s") diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index bcd4f50c8..8c0ea5caa 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/monkey/infection_monkey/master/ip_scanner.py @@ -61,33 +61,15 @@ class IPScanner: address = addresses.get_nowait() logger.info(f"Scanning {address.ip}") - try: - ping_scan_data = self._puppet.ping(address.ip, icmp_timeout) - except Exception as ex: - logger.warning(f"Exception encountered when pinging {address.ip}: {str(ex)}") - ping_scan_data = PingScanData(False, None) - - try: - port_scan_data = self._puppet.scan_tcp_ports(address.ip, tcp_ports, tcp_timeout) - except Exception as ex: - logger.warning( - f"Exception encountered when scanning TCP ports on {address.ip}: {str(ex)}" - ) - port_scan_data = {-1: PortScanData(-1, PortStatus.CLOSED, None, None)} + ping_scan_data = self._puppet.ping(address.ip, icmp_timeout) + port_scan_data = self._puppet.scan_tcp_ports(address.ip, tcp_ports, tcp_timeout) fingerprint_data = {} if IPScanner.port_scan_found_open_port(port_scan_data): fingerprinters = options["fingerprinters"] - try: - fingerprint_data = self._run_fingerprinters( - address.ip, fingerprinters, ping_scan_data, port_scan_data, stop - ) - except Exception as ex: - logger.warning( - f"Exception encountered running fingerprinters on {address.ip}: " - f"{str(ex)}" - ) - fingerprint_data = FingerprintData(None, None, {}) + fingerprint_data = self._run_fingerprinters( + address.ip, fingerprinters, ping_scan_data, port_scan_data, stop + ) scan_results = IPScanResults(ping_scan_data, port_scan_data, fingerprint_data) results_callback(address, scan_results) diff --git a/monkey/infection_monkey/network_scanning/__init__.py b/monkey/infection_monkey/network_scanning/__init__.py index 8e97b0ec4..a1ec7240d 100644 --- a/monkey/infection_monkey/network_scanning/__init__.py +++ b/monkey/infection_monkey/network_scanning/__init__.py @@ -1,2 +1,2 @@ -from .ping_scanner import ping -from .tcp_scanner import scan_tcp_ports +from .ping_scanner import try_ping +from .tcp_scanner import try_scan_tcp_ports diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index 66ae79b2a..c2417d47e 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -6,16 +6,26 @@ import subprocess import sys from infection_monkey.i_puppet import PingScanData +from infection_monkey.utils.environment import is_windows_os TTL_REGEX = re.compile(r"TTL=([0-9]+)\b", re.IGNORECASE) LINUX_TTL = 64 # Windows TTL is 128 PING_EXIT_TIMEOUT = 10 +EMPTY_PING_SCAN = PingScanData(False, None) logger = logging.getLogger(__name__) -def ping(host: str, timeout: float) -> PingScanData: - if "win32" == sys.platform: +def try_ping(host: str, timeout: float) -> PingScanData: + try: + return _ping(host, timeout) + except Exception: + logging.exception("Unhandled exception occurred while running ping") + return EMPTY_PING_SCAN + + +def _ping(host: str, timeout: float) -> PingScanData: + if is_windows_os(): timeout = math.floor(timeout * 1000) ping_command_output = _run_ping_command(host, timeout) diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index 6fdded293..bdada64e8 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -11,11 +11,20 @@ from infection_monkey.utils.timer import Timer logger = logging.getLogger(__name__) POLL_INTERVAL = 0.5 +EMPTY_PORT_SCAN = {-1: PortScanData(-1, PortStatus.CLOSED, None, None)} -def scan_tcp_ports( +def try_scan_tcp_ports( host: str, ports_to_scan: Iterable[int], timeout: float ) -> Mapping[int, PortScanData]: + try: + return _scan_tcp_ports(host, ports_to_scan, timeout) + except Exception: + logging.exception("Unhandled exception occurred while trying to scan tcp ports") + return EMPTY_PORT_SCAN + + +def _scan_tcp_ports(host: str, ports_to_scan: Iterable[int], timeout: float): open_ports = _check_tcp_ports(host, ports_to_scan, timeout) return _build_port_scan_data(ports_to_scan, open_ports) diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 030008a04..853e63a95 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -18,6 +18,8 @@ from infection_monkey.model import VictimHost from .plugin_registry import PluginRegistry +EMPTY_FINGERPRINT = PingScanData(False, None) + logger = logging.getLogger() @@ -39,12 +41,12 @@ class Puppet(IPuppet): return pba.run(options) def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData: - return network_scanning.ping(host, timeout) + return network_scanning.try_ping(host, timeout) def scan_tcp_ports( self, host: str, ports: List[int], timeout: float = CONNECTION_TIMEOUT ) -> Dict[int, PortScanData]: - return network_scanning.scan_tcp_ports(host, ports, timeout) + return network_scanning.try_scan_tcp_ports(host, ports, timeout) def fingerprint( self, @@ -54,8 +56,14 @@ class Puppet(IPuppet): port_scan_data: Dict[int, PortScanData], options: Dict, ) -> FingerprintData: - fingerprinter = self._plugin_registry.get_plugin(name, PluginType.FINGERPRINTER) - return fingerprinter.get_host_fingerprint(host, ping_scan_data, port_scan_data, options) + try: + fingerprinter = self._plugin_registry.get_plugin(name, PluginType.FINGERPRINTER) + return fingerprinter.get_host_fingerprint(host, ping_scan_data, port_scan_data, options) + except Exception: + logging.exception( + f"Unhandled exception occurred " f"while trying to run {name} fingerprinter" + ) + return EMPTY_FINGERPRINT def exploit_host( self, diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping.py index 37e365682..2624997c4 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock import pytest -from infection_monkey.network_scanning import ping +from infection_monkey.network_scanning import _ping LINUX_SUCCESS_OUTPUT = """ PING 192.168.1.1 (192.168.1.1) 56(84) bytes of data. @@ -87,7 +87,7 @@ def set_os_windows(monkeypatch): @pytest.mark.usefixtures("set_os_linux") def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output): patch_subprocess_running_ping_with_ping_output(LINUX_SUCCESS_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = _ping("192.168.1.1", 1.0) assert result.response_received assert result.os == "linux" @@ -96,7 +96,7 @@ def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output): @pytest.mark.usefixtures("set_os_linux") def test_linux_ping_no_response(patch_subprocess_running_ping_with_ping_output): patch_subprocess_running_ping_with_ping_output(LINUX_NO_RESPONSE_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = _ping("192.168.1.1", 1.0) assert not result.response_received assert result.os is None @@ -105,7 +105,7 @@ def test_linux_ping_no_response(patch_subprocess_running_ping_with_ping_output): @pytest.mark.usefixtures("set_os_windows") def test_windows_ping_success(patch_subprocess_running_ping_with_ping_output): patch_subprocess_running_ping_with_ping_output(WINDOWS_SUCCESS_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = _ping("192.168.1.1", 1.0) assert result.response_received assert result.os == "windows" @@ -114,7 +114,7 @@ def test_windows_ping_success(patch_subprocess_running_ping_with_ping_output): @pytest.mark.usefixtures("set_os_windows") def test_windows_ping_no_response(patch_subprocess_running_ping_with_ping_output): patch_subprocess_running_ping_with_ping_output(WINDOWS_NO_RESPONSE_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = _ping("192.168.1.1", 1.0) assert not result.response_received assert result.os is None @@ -122,7 +122,7 @@ def test_windows_ping_no_response(patch_subprocess_running_ping_with_ping_output def test_malformed_ping_command_response(patch_subprocess_running_ping_with_ping_output): patch_subprocess_running_ping_with_ping_output(MALFORMED_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = _ping("192.168.1.1", 1.0) assert not result.response_received assert result.os is None @@ -130,7 +130,7 @@ def test_malformed_ping_command_response(patch_subprocess_running_ping_with_ping @pytest.mark.usefixtures("patch_subprocess_running_ping_to_raise_timeout_expired") def test_timeout_expired(): - result = ping("192.168.1.1", 1.0) + result = _ping("192.168.1.1", 1.0) assert not result.response_received assert result.os is None @@ -147,7 +147,7 @@ def ping_command_spy(monkeypatch): @pytest.fixture def assert_expected_timeout(ping_command_spy): def inner(timeout_flag, timeout_input, expected_timeout): - ping("192.168.1.1", timeout_input) + _ping("192.168.1.1", timeout_input) assert ping_command_spy.call_args is not None diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py index 725a3aaa0..5079c1ac9 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py @@ -1,7 +1,7 @@ import pytest from infection_monkey.i_puppet import PortStatus -from infection_monkey.network_scanning import scan_tcp_ports +from infection_monkey.network_scanning import try_scan_tcp_ports PORTS_TO_SCAN = [22, 80, 8080, 143, 445, 2222] @@ -20,7 +20,7 @@ def patch_check_tcp_ports(monkeypatch, open_ports_data): def test_tcp_successful(monkeypatch, patch_check_tcp_ports, open_ports_data): closed_ports = [8080, 143, 445] - port_scan_data = scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) + port_scan_data = try_scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) assert len(port_scan_data) == 6 for port in open_ports_data.keys(): @@ -37,7 +37,7 @@ def test_tcp_successful(monkeypatch, patch_check_tcp_ports, open_ports_data): @pytest.mark.parametrize("open_ports_data", [{}]) def test_tcp_empty_response(monkeypatch, patch_check_tcp_ports, open_ports_data): - port_scan_data = scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) + port_scan_data = try_scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) assert len(port_scan_data) == 6 for port in open_ports_data: @@ -49,6 +49,6 @@ def test_tcp_empty_response(monkeypatch, patch_check_tcp_ports, open_ports_data) @pytest.mark.parametrize("open_ports_data", [OPEN_PORTS_DATA]) def test_tcp_no_ports_to_scan(monkeypatch, patch_check_tcp_ports, open_ports_data): - port_scan_data = scan_tcp_ports("127.0.0.1", [], 0) + port_scan_data = try_scan_tcp_ports("127.0.0.1", [], 0) assert len(port_scan_data) == 0