1+ import errno
2+ import socket
3+ import time
4+ from dataclasses import dataclass , field
5+ from queue import Empty , Queue
6+ from threading import Lock , Thread
7+ from typing import Dict , Iterable , List , Optional
8+
9+ from .scanner import resolve_host
10+
11+ DEFAULT_DISCOVERY_PORTS = (443 , 80 , 22 )
12+
13+
14+ @dataclass
15+ class HostDiscoveryResult :
16+ target : str
17+ ip : Optional [str ]
18+ alive : bool
19+ reason : str
20+ latency_ms : Optional [float ] = None
21+
22+
23+ @dataclass
24+ class HostDiscoverySummary :
25+ alive : List [HostDiscoveryResult ] = field (default_factory = list )
26+ down : List [HostDiscoveryResult ] = field (default_factory = list )
27+ unresolved : List [HostDiscoveryResult ] = field (default_factory = list )
28+
29+
30+ def tcp_probe (ip : str , port : int , timeout : float ) -> tuple [bool , str , Optional [float ]]:
31+ """Probe a TCP port and return (is_host_reachable, outcome, latency_ms)."""
32+ start = time .perf_counter ()
33+ try :
34+ with socket .socket (socket .AF_INET , socket .SOCK_STREAM ) as sock :
35+ sock .settimeout (timeout )
36+ code = sock .connect_ex ((ip , port ))
37+ latency_ms = (time .perf_counter () - start ) * 1000.0
38+
39+ if code == 0 :
40+ return True , "open" , latency_ms
41+
42+ # RST/connection refused indicates the host is reachable.
43+ if code == errno .ECONNREFUSED :
44+ return True , "refused" , latency_ms
45+
46+ return False , f"code:{ code } , latency_ms
47+ except socket .timeout :
48+ return False , "timeout" , (time .perf_counter () - start ) * 1000.0
49+ except OSError as e :
50+ return False , f"error:{ e .errno } , (time .perf_counter () - start ) * 1000.0
51+
52+
53+ def discover_host (target : str , probe_ports : Iterable [int ], timeout : float ) -> HostDiscoveryResult :
54+ """Resolve and probe a single target to determine if it is likely reachable."""
55+ ip = resolve_host (target )
56+ if ip is None :
57+ return HostDiscoveryResult (target = target , ip = None , alive = False , reason = "dns_failed" )
58+
59+ for port in probe_ports :
60+ alive , outcome , latency_ms = tcp_probe (ip , port , timeout )
61+ if alive :
62+ return HostDiscoveryResult (
63+ target = target ,
64+ ip = ip ,
65+ alive = True ,
66+ reason = f"{ outcome } { port } ,
67+ latency_ms = latency_ms ,
68+ )
69+
70+ return HostDiscoveryResult (target = target , ip = ip , alive = False , reason = "all_probes_failed" )
71+
72+
73+ def _discovery_worker (
74+ probe_ports : List [int ],
75+ timeout : float ,
76+ q : Queue [str ],
77+ summary : HostDiscoverySummary ,
78+ lock : Lock ,
79+ ) -> None :
80+ while True :
81+ try :
82+ target = q .get_nowait ()
83+ except Empty :
84+ return
85+
86+ try :
87+ result = discover_host (target , probe_ports , timeout )
88+ with lock :
89+ if result .ip is None :
90+ summary .unresolved .append (result )
91+ elif result .alive :
92+ summary .alive .append (result )
93+ else :
94+ summary .down .append (result )
95+ finally :
96+ q .task_done ()
97+
98+
99+ def discover_hosts (
100+ targets : List [str ],
101+ probe_ports : List [int ] | None = None ,
102+ timeout : float = 0.3 ,
103+ workers : int = 100 ,
104+ ) -> HostDiscoverySummary :
105+ """
106+ Discover reachable hosts from a list of targets.
107+
108+ A host is considered reachable when any discovery probe returns an open
109+ or connection-refused response.
110+ """
111+ deduped_targets = [t for t in dict .fromkeys (targets ) if t ]
112+ ports = list (probe_ports ) if probe_ports else list (DEFAULT_DISCOVERY_PORTS )
113+ if not ports :
114+ raise ValueError ("Discovery probe port list cannot be empty" )
115+
116+ summary = HostDiscoverySummary ()
117+ if not deduped_targets :
118+ return summary
119+
120+ if workers <= 1 :
121+ for target in deduped_targets :
122+ result = discover_host (target , ports , timeout )
123+ if result .ip is None :
124+ summary .unresolved .append (result )
125+ elif result .alive :
126+ summary .alive .append (result )
127+ else :
128+ summary .down .append (result )
129+ return summary
130+
131+ q : Queue [str ] = Queue ()
132+ for target in deduped_targets :
133+ q .put (target )
134+
135+ lock = Lock ()
136+ threads : list [Thread ] = []
137+ w = min (max (1 , workers ), len (deduped_targets ))
138+ for _ in range (w ):
139+ t = Thread (target = _discovery_worker , args = (ports , timeout , q , summary , lock ), daemon = True )
140+ t .start ()
141+ threads .append (t )
142+
143+ q .join ()
144+ for t in threads :
145+ t .join (timeout = 0.1 )
146+
147+ # Preserve input order for deterministic output.
148+ alive_map = {r .target : r for r in summary .alive }
149+ down_map = {r .target : r for r in summary .down }
150+ unresolved_map = {r .target : r for r in summary .unresolved }
151+ summary .alive = [alive_map [t ] for t in deduped_targets if t in alive_map ]
152+ summary .down = [down_map [t ] for t in deduped_targets if t in down_map ]
153+ summary .unresolved = [unresolved_map [t ] for t in deduped_targets if t in unresolved_map ]
154+ return summary
155+
156+
157+ def to_resolved_target_map (summary : HostDiscoverySummary ) -> Dict [str , str ]:
158+ """Return target -> ip mapping for discovered reachable hosts."""
159+ return {r .target : r .ip for r in summary .alive if r .ip is not None }
0 commit comments