Agent: Add fingerprinting to IPScanner

monkey/infection_monkey/master/ip_scanner.py

@@ -5,7 +5,13 @@ from queue import Queue
 from threading import Event
 from typing import Callable, Dict, List
 
-from infection_monkey.i_puppet import IPuppet, PingScanData, PortScanData
+from infection_monkey.i_puppet import (
+    FingerprintData,
+    IPuppet,
+    PingScanData,
+    PortScanData,
+    PortStatus,
+)
 
 from .threading_utils import create_daemon_thread
 
@@ -13,7 +19,10 @@ logger = logging.getLogger()
 
 IP = str
 Port = int
-Callback = Callable[[IP, PingScanData, Dict[Port, PortScanData]], None]
+FingerprinterName = str
+Callback = Callable[
+    [IP, PingScanData, Dict[Port, PortScanData], Dict[FingerprinterName, FingerprintData]], None
+]
 
 
 class IPScanner:
@@ -53,7 +62,12 @@ class IPScanner:
                 tcp_ports = options["tcp"]["ports"]
                 port_scan_data = self._scan_tcp_ports(ip, tcp_ports, tcp_timeout, stop)
 
-                results_callback(ip, ping_scan_data, port_scan_data)
+                fingerprint_data = {}
+                if IPScanner._found_open_port(port_scan_data):
+                    fingerprinters = options["fingerprinters"]
+                    fingerprint_data = self._run_fingerprinters(ip, fingerprinters, stop)
+
+                results_callback(ip, ping_scan_data, port_scan_data, fingerprint_data)
 
             logger.debug(
                 f"Detected the stop signal, scanning thread {threading.get_ident()} exiting"
@@ -64,7 +78,9 @@ class IPScanner:
                 f"ips_to_scan queue is empty, scanning thread {threading.get_ident()} exiting"
             )
 
-    def _scan_tcp_ports(self, ip: str, ports: List[int], timeout: float, stop: Event):
+    def _scan_tcp_ports(
+        self, ip: str, ports: List[int], timeout: float, stop: Event
+    ) -> Dict[int, PortScanData]:
         port_scan_data = {}
 
         for p in ports:
@@ -74,3 +90,22 @@ class IPScanner:
             port_scan_data[p] = self._puppet.scan_tcp_port(ip, p, timeout)
 
         return port_scan_data
+
+    @staticmethod
+    def _found_open_port(port_scan_data: Dict[int, PortScanData]):
+        for psd in port_scan_data.values():
+            if psd.status == PortStatus.OPEN:
+                return True
+
+        return False
+
+    def _run_fingerprinters(self, ip: str, fingerprinters: List[str], stop: Event):
+        fingerprint_data = {}
+
+        for f in fingerprinters:
+            if stop.is_set():
+                break
+
+            fingerprint_data[f] = self._puppet.fingerprint(f, ip)
+
+        return fingerprint_data

monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py

@@ -4,7 +4,7 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from infection_monkey.i_puppet import PortScanData, PortStatus
+from infection_monkey.i_puppet import FingerprintData, PortScanData, PortStatus
 from infection_monkey.master import IPScanner
 from infection_monkey.puppet.mock_puppet import MockPuppet
 
@@ -29,6 +29,7 @@ def scan_config():
         "icmp": {
             "timeout_ms": 1000,
         },
+        "fingerprinters": {"HTTPFinger", "SMBFinger", "SSHFinger"},
     }
 
 
@@ -50,9 +51,16 @@ def assert_port_status(port_scan_data, expected_open_ports: Set[int]):
             assert psd.status == PortStatus.CLOSED
 
 
-def assert_scan_results_no_1(ip, ping_scan_data, port_scan_data):
-    assert ip == "10.0.0.1"
+def assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data):
+    if ip == "10.0.0.1":
+        assert_scan_results_no_1(ping_scan_data, port_scan_data, fingerprint_data)
+    elif ip == "10.0.0.3":
+        assert_scan_results_no_3(ping_scan_data, port_scan_data, fingerprint_data)
+    else:
+        assert_scan_results_host_down(ip, ping_scan_data, port_scan_data, fingerprint_data)
 
+
+def assert_scan_results_no_1(ping_scan_data, port_scan_data, fingerprint_data):
     assert ping_scan_data.response_received is True
     assert ping_scan_data.os == WINDOWS_OS
 
@@ -70,11 +78,22 @@ def assert_scan_results_no_1(ip, ping_scan_data, port_scan_data):
     assert psd_3389.service == "tcp-3389"
 
     assert_port_status(port_scan_data, {445, 3389})
+    assert_fingerprint_results_no_1(fingerprint_data)
 
 
-def assert_scan_results_no_3(ip, ping_scan_data, port_scan_data):
-    assert ip == "10.0.0.3"
+def assert_fingerprint_results_no_1(fingerprint_data):
+    assert len(fingerprint_data.keys()) == 3
+    assert fingerprint_data["SSHFinger"].services == {}
+    assert fingerprint_data["HTTPFinger"].services == {}
 
+    assert fingerprint_data["SMBFinger"].os_type == WINDOWS_OS
+    assert fingerprint_data["SMBFinger"].os_version == "vista"
+
+    assert len(fingerprint_data["SMBFinger"].services.keys()) == 1
+    assert fingerprint_data["SMBFinger"].services["tcp-445"]["name"] == "smb_service_name"
+
+
+def assert_scan_results_no_3(ping_scan_data, port_scan_data, fingerprint_data):
     assert ping_scan_data.response_received is True
     assert ping_scan_data.os == LINUX_OS
     assert len(port_scan_data.keys()) == 6
@@ -91,15 +110,36 @@ def assert_scan_results_no_3(ip, ping_scan_data, port_scan_data):
     assert psd_22.service == "tcp-22"
 
     assert_port_status(port_scan_data, {22, 443})
+    assert_fingerprint_results_no_3(fingerprint_data)
 
 
-def assert_scan_results_host_down(ip, ping_scan_data, port_scan_data):
+def assert_fingerprint_results_no_3(fingerprint_data):
+    assert len(fingerprint_data.keys()) == 3
+    assert fingerprint_data["SMBFinger"].services == {}
+
+    assert fingerprint_data["SSHFinger"].os_type == LINUX_OS
+    assert fingerprint_data["SSHFinger"].os_version == "ubuntu"
+
+    assert len(fingerprint_data["SSHFinger"].services.keys()) == 1
+    assert fingerprint_data["SSHFinger"].services["tcp-22"]["name"] == "SSH"
+    assert fingerprint_data["SSHFinger"].services["tcp-22"]["banner"] == "SSH BANNER"
+
+    assert len(fingerprint_data["HTTPFinger"].services.keys()) == 2
+    assert fingerprint_data["HTTPFinger"].services["tcp-80"]["name"] == "http"
+    assert fingerprint_data["HTTPFinger"].services["tcp-80"]["data"] == ("SERVER_HEADERS", False)
+    assert fingerprint_data["HTTPFinger"].services["tcp-443"]["name"] == "http"
+    assert fingerprint_data["HTTPFinger"].services["tcp-443"]["data"] == ("SERVER_HEADERS_2", True)
+
+
+def assert_scan_results_host_down(ip, ping_scan_data, port_scan_data, fingerprint_data):
     assert ip not in {"10.0.0.1", "10.0.0.3"}
 
     assert ping_scan_data.response_received is False
     assert len(port_scan_data.keys()) == 6
     assert_port_status(port_scan_data, set())
 
+    assert fingerprint_data == {}
+
 
 def test_scan_single_ip(callback, scan_config, stop):
     ips = ["10.0.0.1"]
@@ -109,8 +149,8 @@ def test_scan_single_ip(callback, scan_config, stop):
 
     callback.assert_called_once()
 
-    (ip, ping_scan_data, port_scan_data) = callback.call_args_list[0][0]
-    assert_scan_results_no_1(ip, ping_scan_data, port_scan_data)
+    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[0][0]
+    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
 
 
 def test_scan_multiple_ips(callback, scan_config, stop):
@@ -121,17 +161,17 @@ def test_scan_multiple_ips(callback, scan_config, stop):
 
     assert callback.call_count == 4
 
-    (ip, ping_scan_data, port_scan_data) = callback.call_args_list[0][0]
-    assert_scan_results_no_1(ip, ping_scan_data, port_scan_data)
+    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[0][0]
+    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
 
-    (ip, ping_scan_data, port_scan_data) = callback.call_args_list[1][0]
-    assert_scan_results_host_down(ip, ping_scan_data, port_scan_data)
+    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[1][0]
+    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
 
-    (ip, ping_scan_data, port_scan_data) = callback.call_args_list[2][0]
-    assert_scan_results_no_3(ip, ping_scan_data, port_scan_data)
+    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[2][0]
+    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
 
-    (ip, ping_scan_data, port_scan_data) = callback.call_args_list[3][0]
-    assert_scan_results_host_down(ip, ping_scan_data, port_scan_data)
+    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[3][0]
+    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
 
 
 def test_scan_lots_of_ips(callback, scan_config, stop):
@@ -182,3 +222,25 @@ def test_interrupt_port_scanning(callback, scan_config, stop):
     ns.scan(ips, scan_config, callback, stop)
 
     assert puppet.scan_tcp_port.call_count == 2
+
+
+def test_interrupt_fingerprinting(callback, scan_config, stop):
+    def stopable_fingerprint(port, *_):
+        # Block all threads here until 2 threads reach this barrier, then set stop
+        # and test that neither thread scans any more ports
+        stopable_fingerprint.barrier.wait()
+        stop.set()
+
+        return FingerprintData(None, None, {})
+
+    stopable_fingerprint.barrier = Barrier(2)
+
+    puppet = MockPuppet()
+    puppet.fingerprint = MagicMock(side_effect=stopable_fingerprint)
+
+    ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.4"]
+
+    ns = IPScanner(puppet, num_workers=2)
+    ns.scan(ips, scan_config, callback, stop)
+
+    assert puppet.fingerprint.call_count == 2