Agent: Improve exception handling of tcp, ping and fingerprint scanners

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:
+            except Exception:
-                logger.debug(f"Exception encountered when running {plugin_type} {p}: {str(ex)}")
+                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")
 

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)
-                    ping_scan_data = self._puppet.ping(address.ip, icmp_timeout)
+                port_scan_data = self._puppet.scan_tcp_ports(address.ip, tcp_ports, tcp_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)}
 
                 fingerprint_data = {}
                 if IPScanner.port_scan_found_open_port(port_scan_data):
                     fingerprinters = options["fingerprinters"]
-                    try:
+                    fingerprint_data = self._run_fingerprinters(
-                        fingerprint_data = self._run_fingerprinters(
+                        address.ip, fingerprinters, ping_scan_data, port_scan_data, stop
-                            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, {})
 
                 scan_results = IPScanResults(ping_scan_data, port_scan_data, fingerprint_data)
                 results_callback(address, scan_results)

monkey/infection_monkey/network_scanning/__init__.py

@@ -1,2 +1,2 @@
-from .ping_scanner import ping
+from .ping_scanner import try_ping
-from .tcp_scanner import scan_tcp_ports
+from .tcp_scanner import try_scan_tcp_ports

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:
+def try_ping(host: str, timeout: float) -> PingScanData:
-    if "win32" == sys.platform:
+    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)

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)

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)
+        try:
-        return fingerprinter.get_host_fingerprint(host, ping_scan_data, port_scan_data, options)
+            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,

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
 

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