From 9495027678278f28a73eac86d8035cdb86ee96a3 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 5 Apr 2022 17:56:12 +0530 Subject: [PATCH 1/7] Agent: Catch exception when running plugins in master --- monkey/infection_monkey/master/automated_master.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 0ef31129c..5586e808c 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -225,8 +225,10 @@ class AutomatedMaster(IMaster): interrupted_message = f"Received a stop signal, skipping remaining {plugin_type}s" for p in interruptible_iter(plugins, self._stop, interrupted_message): - # TODO: Catch exceptions to prevent thread from crashing - callback(p) + try: + callback(p) + except Exception as ex: + logger.debug(f"Exception encountered when running {plugin_type} {p}: {str(ex)}") logger.info(f"Finished running {plugin_type}s") From f9a6afffb410dff9cd6ae9b8159555bcf905a526 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 5 Apr 2022 18:27:53 +0530 Subject: [PATCH 2/7] Agent: Catch exceptions in IPScanner so thread doesn't crash --- monkey/infection_monkey/master/ip_scanner.py | 26 +++++++++++++++----- 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index 2702642c9..a01f23022 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/monkey/infection_monkey/master/ip_scanner.py @@ -61,16 +61,30 @@ class IPScanner: address = addresses.get_nowait() logger.info(f"Scanning {address.ip}") - # TODO: Catch exceptions to prevent thread from crashing - ping_scan_data = self._puppet.ping(address.ip, icmp_timeout) - port_scan_data = self._puppet.scan_tcp_ports(address.ip, tcp_ports, tcp_timeout) + try: + ping_scan_data = self._puppet.ping(address.ip, icmp_timeout) + except Exception as ex: + logger.debug(f"Exception encountered when pinging {address.ip}: {str(ex)}") + + try: + port_scan_data = self._puppet.scan_tcp_ports(address.ip, tcp_ports, tcp_timeout) + except Exception as ex: + logger.debug( + f"Exception encountered when scanning TCP ports on {address.ip}: {str(ex)}" + ) fingerprint_data = {} if IPScanner.port_scan_found_open_port(port_scan_data): fingerprinters = options["fingerprinters"] - fingerprint_data = self._run_fingerprinters( - address.ip, fingerprinters, ping_scan_data, port_scan_data, stop - ) + try: + fingerprint_data = self._run_fingerprinters( + address.ip, fingerprinters, ping_scan_data, port_scan_data, stop + ) + except Exception as ex: + logger.debug( + f"Exception encountered running fingerprinters on {address.ip}: " + f"{str(ex)}" + ) scan_results = IPScanResults(ping_scan_data, port_scan_data, fingerprint_data) results_callback(address, scan_results) From e3fc5cf5e54de51c7b5e071c2fd3939fac5a2b2c Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 7 Apr 2022 14:03:40 +0530 Subject: [PATCH 3/7] Agent, UT: Add default return values in IPScanner + tests --- monkey/infection_monkey/master/ip_scanner.py | 3 ++ .../master/test_ip_scanner.py | 43 ++++++++++++++++++- 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index a01f23022..70c10205d 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/monkey/infection_monkey/master/ip_scanner.py @@ -65,6 +65,7 @@ class IPScanner: ping_scan_data = self._puppet.ping(address.ip, icmp_timeout) except Exception as ex: logger.debug(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) @@ -72,6 +73,7 @@ class IPScanner: logger.debug( 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): @@ -85,6 +87,7 @@ class IPScanner: 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) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py b/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py index 9fafdaee2..e5b0381e2 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py @@ -5,7 +5,7 @@ from unittest.mock import MagicMock import pytest from tests.unit_tests.infection_monkey.master.mock_puppet import MockPuppet -from infection_monkey.i_puppet import FingerprintData, PortScanData, PortStatus +from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData, PortStatus from infection_monkey.master import IPScanner from infection_monkey.network import NetworkAddress @@ -202,6 +202,47 @@ def test_scan_lots_of_ips(callback, scan_config, stop): assert callback.call_count == 255 +def test_exception_in_pinging(callback, scan_config, stop): + addresses = [NetworkAddress("10.0.0.1", "d1")] + + puppet = MockPuppet() + puppet.ping = MagicMock(side_effect=Exception("Exception raised during pinging.")) + + ns = IPScanner(puppet, num_workers=1) + ns.scan(addresses, scan_config, callback, stop) + + (_, scan_results) = callback.call_args_list[0][0] + assert scan_results.ping_scan_data == PingScanData(False, None) + + +def test_exception_in_port_scanning(callback, scan_config, stop): + addresses = [NetworkAddress("10.0.0.1", "d1")] + + puppet = MockPuppet() + puppet.scan_tcp_ports = MagicMock( + side_effect=Exception("Exception raised when scanning TCP ports.") + ) + + ns = IPScanner(puppet, num_workers=1) + ns.scan(addresses, scan_config, callback, stop) + + (_, scan_results) = callback.call_args_list[0][0] + assert scan_results.port_scan_data == {-1: PortScanData(-1, PortStatus.CLOSED, None, None)} + + +def test_exception_in_fingerprinting(callback, scan_config, stop): + addresses = [NetworkAddress("10.0.0.1", "d1")] + + puppet = MockPuppet() + puppet.fingerprint = MagicMock(side_effect=Exception("Exception raised during fingerprinting.")) + + ns = IPScanner(puppet, num_workers=1) + ns.scan(addresses, scan_config, callback, stop) + + (_, scan_results) = callback.call_args_list[0][0] + assert scan_results.fingerprint_data == FingerprintData(None, None, {}) + + def test_stop_after_callback(scan_config, stop): def _callback(*_): # Block all threads here until 2 threads reach this barrier, then set stop From 58d4c339593ec6dc13884e94f9e6a471b922c22d Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 7 Apr 2022 14:04:30 +0530 Subject: [PATCH 4/7] Agent: Change exceptions' log level to warning in IPScanner --- monkey/infection_monkey/master/ip_scanner.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index 70c10205d..bcd4f50c8 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/monkey/infection_monkey/master/ip_scanner.py @@ -64,13 +64,13 @@ class IPScanner: try: ping_scan_data = self._puppet.ping(address.ip, icmp_timeout) except Exception as ex: - logger.debug(f"Exception encountered when pinging {address.ip}: {str(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.debug( + logger.warning( f"Exception encountered when scanning TCP ports on {address.ip}: {str(ex)}" ) port_scan_data = {-1: PortScanData(-1, PortStatus.CLOSED, None, None)} @@ -83,7 +83,7 @@ class IPScanner: address.ip, fingerprinters, ping_scan_data, port_scan_data, stop ) except Exception as ex: - logger.debug( + logger.warning( f"Exception encountered running fingerprinters on {address.ip}: " f"{str(ex)}" ) From 6157ffee7696aff9e9be9651a322f4f0d64c91d6 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 7 Apr 2022 17:55:05 +0300 Subject: [PATCH 5/7] 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 From 45c6cac60c9e671b40415857b7d6aec40fdd4ddf Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 8 Apr 2022 10:36:14 +0300 Subject: [PATCH 6/7] Agent: Improve method naming and exception handling --- monkey/infection_monkey/master/automated_master.py | 2 +- monkey/infection_monkey/network_scanning/__init__.py | 4 ++-- monkey/infection_monkey/network_scanning/ping_scanner.py | 4 ++-- monkey/infection_monkey/network_scanning/tcp_scanner.py | 4 ++-- monkey/infection_monkey/puppet/puppet.py | 6 +++--- .../network_scanning/test_tcp_scanning.py | 8 ++++---- 6 files changed, 14 insertions(+), 14 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 99f63522a..af263af6e 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -228,7 +228,7 @@ class AutomatedMaster(IMaster): try: callback(p) except Exception: - logging.exception( + logger.exception( f"Got unhandled exception when running {plugin_type} plugin {p}. " f"Plugin was passed to {callback}" ) diff --git a/monkey/infection_monkey/network_scanning/__init__.py b/monkey/infection_monkey/network_scanning/__init__.py index a1ec7240d..8e97b0ec4 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 try_ping -from .tcp_scanner import try_scan_tcp_ports +from .ping_scanner import ping +from .tcp_scanner import scan_tcp_ports diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index c2417d47e..16fb2df96 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -16,11 +16,11 @@ EMPTY_PING_SCAN = PingScanData(False, None) logger = logging.getLogger(__name__) -def try_ping(host: str, timeout: float) -> PingScanData: +def ping(host: str, timeout: float) -> PingScanData: try: return _ping(host, timeout) except Exception: - logging.exception("Unhandled exception occurred while running ping") + logger.exception("Unhandled exception occurred while running ping") return EMPTY_PING_SCAN diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index bdada64e8..d0c6e3e7a 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -14,13 +14,13 @@ POLL_INTERVAL = 0.5 EMPTY_PORT_SCAN = {-1: PortScanData(-1, PortStatus.CLOSED, None, None)} -def try_scan_tcp_ports( +def 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") + logger.exception("Unhandled exception occurred while trying to scan tcp ports") return EMPTY_PORT_SCAN diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 853e63a95..c7953d83c 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -41,12 +41,12 @@ class Puppet(IPuppet): return pba.run(options) def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData: - return network_scanning.try_ping(host, timeout) + return network_scanning.ping(host, timeout) def scan_tcp_ports( self, host: str, ports: List[int], timeout: float = CONNECTION_TIMEOUT ) -> Dict[int, PortScanData]: - return network_scanning.try_scan_tcp_ports(host, ports, timeout) + return network_scanning.scan_tcp_ports(host, ports, timeout) def fingerprint( self, @@ -60,7 +60,7 @@ class Puppet(IPuppet): 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( + logger.exception( f"Unhandled exception occurred " f"while trying to run {name} fingerprinter" ) return EMPTY_FINGERPRINT 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 5079c1ac9..725a3aaa0 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 try_scan_tcp_ports +from infection_monkey.network_scanning import 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 = try_scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) + port_scan_data = 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 = try_scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) + port_scan_data = 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 = try_scan_tcp_ports("127.0.0.1", [], 0) + port_scan_data = scan_tcp_ports("127.0.0.1", [], 0) assert len(port_scan_data) == 0 From e1b52428d128658e9cc55f41e4a34b320d93dd76 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 8 Apr 2022 12:00:06 +0300 Subject: [PATCH 7/7] UT: Add some tests for exception handling --- .../master/test_ip_scanner.py | 43 +------------------ .../{test_ping.py => test_ping_scanner.py} | 24 +++++++---- ...st_tcp_scanning.py => test_tcp_scanner.py} | 13 +++++- .../infection_monkey/puppet/test_puppet.py | 10 ++++- 4 files changed, 36 insertions(+), 54 deletions(-) rename monkey/tests/unit_tests/infection_monkey/network_scanning/{test_ping.py => test_ping_scanner.py} (88%) rename monkey/tests/unit_tests/infection_monkey/network_scanning/{test_tcp_scanning.py => test_tcp_scanner.py} (83%) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py b/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py index e5b0381e2..9fafdaee2 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py @@ -5,7 +5,7 @@ from unittest.mock import MagicMock import pytest from tests.unit_tests.infection_monkey.master.mock_puppet import MockPuppet -from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData, PortStatus +from infection_monkey.i_puppet import FingerprintData, PortScanData, PortStatus from infection_monkey.master import IPScanner from infection_monkey.network import NetworkAddress @@ -202,47 +202,6 @@ def test_scan_lots_of_ips(callback, scan_config, stop): assert callback.call_count == 255 -def test_exception_in_pinging(callback, scan_config, stop): - addresses = [NetworkAddress("10.0.0.1", "d1")] - - puppet = MockPuppet() - puppet.ping = MagicMock(side_effect=Exception("Exception raised during pinging.")) - - ns = IPScanner(puppet, num_workers=1) - ns.scan(addresses, scan_config, callback, stop) - - (_, scan_results) = callback.call_args_list[0][0] - assert scan_results.ping_scan_data == PingScanData(False, None) - - -def test_exception_in_port_scanning(callback, scan_config, stop): - addresses = [NetworkAddress("10.0.0.1", "d1")] - - puppet = MockPuppet() - puppet.scan_tcp_ports = MagicMock( - side_effect=Exception("Exception raised when scanning TCP ports.") - ) - - ns = IPScanner(puppet, num_workers=1) - ns.scan(addresses, scan_config, callback, stop) - - (_, scan_results) = callback.call_args_list[0][0] - assert scan_results.port_scan_data == {-1: PortScanData(-1, PortStatus.CLOSED, None, None)} - - -def test_exception_in_fingerprinting(callback, scan_config, stop): - addresses = [NetworkAddress("10.0.0.1", "d1")] - - puppet = MockPuppet() - puppet.fingerprint = MagicMock(side_effect=Exception("Exception raised during fingerprinting.")) - - ns = IPScanner(puppet, num_workers=1) - ns.scan(addresses, scan_config, callback, stop) - - (_, scan_results) = callback.call_args_list[0][0] - assert scan_results.fingerprint_data == FingerprintData(None, None, {}) - - def test_stop_after_callback(scan_config, stop): def _callback(*_): # Block all threads here until 2 threads reach this barrier, then set stop 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_scanner.py similarity index 88% rename from monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping.py rename to monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index 2624997c4..88c9dbeca 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -4,7 +4,8 @@ from unittest.mock import MagicMock import pytest -from infection_monkey.network_scanning import _ping +from infection_monkey.network_scanning import ping +from infection_monkey.network_scanning.ping_scanner import EMPTY_PING_SCAN LINUX_SUCCESS_OUTPUT = """ PING 192.168.1.1 (192.168.1.1) 56(84) bytes of data. @@ -87,7 +88,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 +97,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 +106,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 +115,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 +123,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 +131,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 +148,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 @@ -174,3 +175,10 @@ def test_linux_timeout(assert_expected_timeout): timeout = 1.42379 assert_expected_timeout(timeout_flag, timeout, str(math.ceil(timeout))) + + +def test_exception_handling(monkeypatch): + monkeypatch.setattr( + "infection_monkey.network_scanning.ping_scanner._ping", MagicMock(side_effect=Exception) + ) + assert ping("abc", 10) == EMPTY_PING_SCAN 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_scanner.py similarity index 83% rename from monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py rename to monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index 725a3aaa0..837b3da0d 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_scanner.py @@ -1,7 +1,10 @@ +from unittest.mock import MagicMock + import pytest from infection_monkey.i_puppet import PortStatus from infection_monkey.network_scanning import scan_tcp_ports +from infection_monkey.network_scanning.tcp_scanner import EMPTY_PORT_SCAN PORTS_TO_SCAN = [22, 80, 8080, 143, 445, 2222] @@ -36,7 +39,6 @@ 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) assert len(port_scan_data) == 6 @@ -48,7 +50,14 @@ 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) assert len(port_scan_data) == 0 + + +def test_exception_handling(monkeypatch): + monkeypatch.setattr( + "infection_monkey.network_scanning.tcp_scanner._scan_tcp_ports", + MagicMock(side_effect=Exception), + ) + assert scan_tcp_ports("abc", [123], 123) == EMPTY_PORT_SCAN diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py index 70c98d252..39273faee 100644 --- a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py +++ b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py @@ -1,8 +1,8 @@ import threading from unittest.mock import MagicMock -from infection_monkey.i_puppet import PluginType -from infection_monkey.puppet.puppet import Puppet +from infection_monkey.i_puppet import PingScanData, PluginType +from infection_monkey.puppet.puppet import EMPTY_FINGERPRINT, Puppet def test_puppet_run_payload_success(): @@ -41,3 +41,9 @@ def test_puppet_run_multiple_payloads(): p.run_payload(payload3_name, {}, threading.Event()) payload_3.run.assert_called_once() + + +def test_fingerprint_exception_handling(monkeypatch): + p = Puppet() + p._plugin_registry.get_plugin = MagicMock(side_effect=Exception) + assert p.fingerprint("", "", PingScanData("windows", False), {}, {}) == EMPTY_FINGERPRINT