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

This commit is contained in:
vakarisz 2022-04-07 17:55:05 +03:00
parent 58d4c33959
commit 6157ffee76
8 changed files with 58 additions and 46 deletions

View File

@ -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")

View File

@ -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)

View File

@ -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

View File

@ -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)

View File

@ -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)

View File

@ -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,

View File

@ -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

View File

@ -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