fix(scan): bound the subnet and stop waiting on devices one at a time

netdiff/cli.py

@@ -13,7 +13,7 @@
 from . import audit as audit_rules
 from . import mdns, oui, store, upnp
 from .diff import diff, summarise
-from .scan import discover, grab_banner
+from .scan import discover, grab_banners
 
 DEFAULT_PORTS = (22, 80, 443, 445, 554, 1883, 3389, 5000, 8080, 8443)
 
@@ -180,11 +180,9 @@ def placeholders(text):
         resolve_names=not args.no_resolve,
         services={} if args.no_mdns else mdns.discover(),
     )
-    banners = {
-        (device.ip, port): grab_banner(device.ip, port)
-        for device in devices
-        for port in device.ports
-    }
+    banners = grab_banners(
+        (device.ip, port) for device in devices for port in device.ports
+    )
     gateway = None if args.no_upnp else upnp.probe_gateway(args.subnet)
     findings = audit_rules.audit(devices, gateway, banners)
 
@@ -339,7 +337,14 @@ def build_parser() -> argparse.ArgumentParser:
 
 def main(argv=None) -> int:
     args = build_parser().parse_args(argv)
-    return args.func(args)
+    try:
+        return args.func(args)
+    except ValueError as exc:
+        # A subnet is user input and it is parsed deep in `discover`, so a typo
+        # or a /8 arrives here rather than at the argparse layer. Say what is
+        # wrong instead of printing a traceback at someone.
+        print(exc, file=sys.stderr)
+        return 2
 
 
 if __name__ == "__main__":

netdiff/scan.py

@@ -26,6 +26,17 @@
 NUDGE_PORT = 9
 ARP_SETTLE_SECONDS = 1.0
 
+# ARP does not cross routers, so a subnet larger than this is not a broadcast
+# segment - it is a typo, or someone pointing the tool at 10.0.0.0/8 to see what
+# happens. That is 16.7M addresses materialised as strings before a single
+# packet moves, so refuse it up front rather than swapping to death.
+MAX_HOSTS = 65536
+
+# Concurrency for per-device work. Each worker is blocked on a socket timeout,
+# not on the CPU, so the useful number is set by how long we are willing to wait
+# rather than by core count.
+SCAN_WORKERS = 32
+
 # `arp -an` on macOS/BSD/Linux: "? (192.168.1.1) at ab:cd:ef:12:34:56 on en0"
 _ARP_LINE = re.compile(
     r"\((?P<ip>\d+\.\d+\.\d+\.\d+)\)\s+at\s+(?P<mac>[0-9a-fA-F:]{11,17})"
@@ -133,6 +144,21 @@ def scan_ports(ip: str, ports, timeout: float = 0.3) -> tuple[int, ...]:
     return tuple(sorted(open_ports))
 
 
+def grab_banners(pairs, timeout: float = 2.0) -> dict:
+    """Banner for each (ip, port), gathered concurrently.
+
+    Sequentially this is the slowest thing the audit does: every port that is
+    open but silent - a TLS port never greets - costs the full timeout, one
+    after another.
+    """
+    pairs = list(pairs)
+    if not pairs:
+        return {}
+    with ThreadPoolExecutor(max_workers=SCAN_WORKERS) as pool:
+        banners = pool.map(lambda pair: grab_banner(*pair, timeout=timeout), pairs)
+        return dict(zip(pairs, banners))
+
+
 def grab_banner(ip: str, port: int, timeout: float = 2.0) -> str:
     """Read what a service volunteers about itself.
 
@@ -180,6 +206,11 @@ def discover(
     plain dict that a test can hand over without a socket.
     """
     network = ipaddress.ip_network(subnet, strict=False)
+    if network.num_addresses > MAX_HOSTS:
+        raise ValueError(
+            f"{subnet} holds {network.num_addresses} addresses; "
+            f"netdiff scans one broadcast segment, up to {MAX_HOSTS}"
+        )
     hosts = [str(h) for h in network.hosts()]
     nudge(hosts)
     time.sleep(settle)
@@ -189,16 +220,20 @@ def discover(
         ip: mac for ip, mac in table.items() if ipaddress.ip_address(ip) in network
     }
 
-    devices = []
-    for ip, mac in in_subnet.items():
-        devices.append(
-            Device(
-                mac=mac,
-                ip=ip,
-                vendor=lookup_vendor(mac) if lookup_vendor else "",
-                hostname=resolve_hostname(ip) if resolve_names else "",
-                services=(services or {}).get(ip, ""),
-                ports=scan_ports(ip, ports) if ports else (),
-            )
+    def observe(item):
+        ip, mac = item
+        return Device(
+            mac=mac,
+            ip=ip,
+            vendor=lookup_vendor(mac) if lookup_vendor else "",
+            hostname=resolve_hostname(ip) if resolve_names else "",
+            services=(services or {}).get(ip, ""),
+            ports=scan_ports(ip, ports) if ports else (),
         )
+
+    # Reverse DNS and the port scan are both waits, not work, and neither depends
+    # on any other device. Sequentially the scan took the sum of every timeout on
+    # the network; concurrently it takes the worst single device.
+    with ThreadPoolExecutor(max_workers=SCAN_WORKERS) as pool:
+        devices = list(pool.map(observe, in_subnet.items()))
     return sorted(devices, key=lambda d: ipaddress.ip_address(d.ip))

tests/test_scan_and_store.py

@@ -2,10 +2,11 @@
 captured command output, and the database is a temp file."""
 
 import sqlite3
+import time
 
 import pytest
 
-from netdiff import store
+from netdiff import cli, scan, store
 from netdiff.audit import Finding
 from netdiff.oui import is_randomised, lookup
 from netdiff.scan import Device, normalise_mac, parse_arp_output, read_arp_table
@@ -86,6 +87,67 @@ def test_randomised_macs_are_labelled_not_treated_as_unknown_vendors():
     assert lookup("a2:bb:cc:dd:ee:ff") == "randomised"
 
 
+# --- bounds and concurrency -------------------------------------------------
+# A scan is almost entirely waiting on timeouts. Sequentially that cost is the
+# sum over every device on the network, which is why these are timed rather than
+# only checked for their return value - a pool that quietly stops being used
+# still returns the right answer.
+
+
+def test_a_subnet_larger_than_one_broadcast_segment_is_refused(monkeypatch):
+    """And refused before anything is sent, not after 16.7M datagrams."""
+    monkeypatch.setattr(scan, "nudge", lambda *a, **k: pytest.fail("sent packets"))
+    with pytest.raises(ValueError, match="broadcast segment"):
+        scan.discover("10.0.0.0/8")
+
+
+def test_a_subnet_that_fits_is_not_refused(monkeypatch):
+    monkeypatch.setattr(scan, "nudge", lambda *a, **k: None)
+    monkeypatch.setattr(scan, "read_arp_table", dict)
+    assert scan.discover("192.168.1.0/24", settle=0) == []
+
+
+def test_a_bad_subnet_is_a_message_not_a_traceback(capsys):
+    assert cli.main(["scan", "not-a-subnet"]) == 2
+    assert capsys.readouterr().err.strip(), "the reason has to reach the user"
+
+
+def test_devices_are_scanned_concurrently(monkeypatch):
+    table = {f"192.168.1.{n}": f"aa:bb:cc:00:00:{n:02x}" for n in range(1, 13)}
+    monkeypatch.setattr(scan, "nudge", lambda *a, **k: None)
+    monkeypatch.setattr(scan, "read_arp_table", lambda: table)
+    monkeypatch.setattr(scan, "scan_ports", lambda ip, ports: time.sleep(0.2) or ())
+
+    start = time.monotonic()
+    devices = scan.discover(
+        "192.168.1.0/24", ports=(22,), settle=0, resolve_names=False
+    )
+    elapsed = time.monotonic() - start
+
+    assert [d.ip for d in devices] == sorted(
+        table, key=lambda ip: int(ip.split(".")[3])
+    )
+    assert elapsed < 1.0, f"12 devices x 0.2s took {elapsed:.1f}s - run sequentially?"
+
+
+def test_banners_are_gathered_concurrently(monkeypatch):
+    pairs = [("192.168.1.10", port) for port in range(8000, 8012)]
+    monkeypatch.setattr(
+        scan, "grab_banner", lambda ip, port, timeout=2.0: time.sleep(0.2) or f"{port}"
+    )
+
+    start = time.monotonic()
+    banners = scan.grab_banners(pairs)
+    elapsed = time.monotonic() - start
+
+    assert banners == {pair: str(pair[1]) for pair in pairs}
+    assert elapsed < 1.0, f"12 ports x 0.2s took {elapsed:.1f}s - run sequentially?"
+
+
+def test_no_open_ports_means_no_banner_work():
+    assert scan.grab_banners([]) == {}
+
+
 @pytest.fixture
 def conn(tmp_path):
     return store.connect(tmp_path / "history.db")