Agent, UT: Add default return values in IPScanner + tests

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)

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