netscli_core/
sweep.rs

1use crate::discover::{DiscoverEngine, Host};
2use crate::error::Result;
3use crate::scan::{PortResult, PortScanner};
4use futures::stream::{self, StreamExt};
5use ipnet::Ipv4Net;
6use serde::Serialize;
7use std::sync::{
8    atomic::{AtomicUsize, Ordering},
9    Arc,
10};
11
12#[derive(Debug, Clone, Serialize)]
13pub struct SweepEntry {
14    pub host: Host,
15    pub open_ports: Vec<PortResult>,
16}
17
18#[derive(Debug, Clone, Copy)]
19pub enum SweepPhase {
20    DiscoverPing,
21    DiscoverResolve,
22    Scan,
23}
24
25#[derive(Debug, Clone)]
26pub struct SweepProgress {
27    pub phase: SweepPhase,
28    pub completed: usize,
29    pub total: usize,
30    pub found: usize,
31    pub ip: std::net::IpAddr,
32}
33
34pub struct SweepEngine {
35    discover: DiscoverEngine,
36    scanner: PortScanner,
37    scan_timeout_ms: u64,
38    host_concurrency: usize,
39}
40
41impl SweepEngine {
42    pub fn new(concurrency: usize) -> Self {
43        Self::new_with_timeouts(
44            concurrency,
45            crate::DEFAULT_PING_TIMEOUT_MS,
46            crate::DEFAULT_SCAN_TIMEOUT_MS,
47            crate::DEFAULT_DNS_TIMEOUT_MS,
48        )
49    }
50
51    pub fn new_with_timeouts(
52        concurrency: usize,
53        ping_timeout_ms: u64,
54        scan_timeout_ms: u64,
55        dns_timeout_ms: u64,
56    ) -> Self {
57        let concurrency = concurrency.max(1);
58        Self {
59            discover: DiscoverEngine::new_with_timeouts(
60                concurrency,
61                ping_timeout_ms,
62                dns_timeout_ms,
63            ),
64            scanner: PortScanner::new(concurrency),
65            scan_timeout_ms: scan_timeout_ms.max(1),
66            host_concurrency: concurrency,
67        }
68    }
69
70    pub async fn sweep(
71        &self,
72        subnet: Ipv4Net,
73        ports: Vec<u16>,
74        resolve_hostnames: bool,
75    ) -> Result<Vec<SweepEntry>> {
76        self.sweep_with_progress(subnet, ports, resolve_hostnames, None)
77            .await
78    }
79
80    /// Sweep a network with bounded concurrency.
81    ///
82    /// Progress reporting is optional and is invoked after each host scan completes.
83    pub async fn sweep_with_progress(
84        &self,
85        subnet: Ipv4Net,
86        ports: Vec<u16>,
87        resolve_hostnames: bool,
88        progress: Option<Arc<dyn Fn(SweepProgress) + Send + Sync>>,
89    ) -> Result<Vec<SweepEntry>> {
90        let discover_progress = progress.clone().map(|cb| {
91            Arc::new(move |p: crate::discover::DiscoverProgress| {
92                let phase = match p.phase {
93                    crate::discover::DiscoverPhase::Ping => SweepPhase::DiscoverPing,
94                    crate::discover::DiscoverPhase::Resolve => SweepPhase::DiscoverResolve,
95                };
96                cb(SweepProgress {
97                    phase,
98                    completed: p.completed,
99                    total: p.total,
100                    found: p.found,
101                    ip: p.ip,
102                });
103            }) as Arc<dyn Fn(crate::discover::DiscoverProgress) + Send + Sync>
104        });
105
106        let hosts = self
107            .discover
108            .scan_subnet_with_progress(subnet, resolve_hostnames, discover_progress)
109            .await;
110        let total_hosts = hosts.len();
111        let completed = Arc::new(AtomicUsize::new(0));
112        let open_hosts = Arc::new(AtomicUsize::new(0));
113        let ports = Arc::new(ports);
114
115        let scan_timeout_ms = self.scan_timeout_ms;
116        let scanner = self.scanner.clone();
117
118        let mut entries = stream::iter(hosts)
119            .map(|h| {
120                let scanner = scanner.clone();
121                let ports = ports.clone();
122                let completed = completed.clone();
123                let open_hosts = open_hosts.clone();
124                let progress = progress.clone();
125                async move {
126                    let host_ip = h.ip;
127                    let open_ports = scanner
128                        .scan_host(host_ip, (*ports).clone(), scan_timeout_ms)
129                        .await
130                        .into_iter()
131                        .filter(|p| p.open)
132                        .collect();
133
134                    let entry = SweepEntry {
135                        host: h,
136                        open_ports,
137                    };
138
139                    let done = completed.fetch_add(1, Ordering::SeqCst) + 1;
140                    if let Some(cb) = &progress {
141                        let found = if entry.open_ports.is_empty() {
142                            open_hosts.load(Ordering::SeqCst)
143                        } else {
144                            open_hosts.fetch_add(1, Ordering::SeqCst) + 1
145                        };
146                        cb(SweepProgress {
147                            phase: SweepPhase::Scan,
148                            completed: done,
149                            total: total_hosts,
150                            found,
151                            ip: host_ip,
152                        });
153                    }
154
155                    entry
156                }
157            })
158            .buffer_unordered(self.host_concurrency)
159            .collect::<Vec<SweepEntry>>()
160            .await;
161
162        // Provide deterministic output order (useful for CLI/TUI and tests).
163        entries.sort_by_key(|e| e.host.ip);
164        Ok(entries)
165    }
166}
netscli_core/sweep.rs

netscli_core/
sweep.rs
