diff options
-rw-r--r-- | src/config.rs | 8 | ||||
-rw-r--r-- | src/main.rs | 40 |
2 files changed, 33 insertions, 15 deletions
diff --git a/src/config.rs b/src/config.rs index 44d500e..23c2a78 100644 --- a/src/config.rs +++ b/src/config.rs @@ -11,6 +11,14 @@ pub struct Config { #[arg(short, long)] pub resolve: bool, + /// filter hosts with open tcp port at <PORT> + #[arg(short, long)] + pub port: Option<u16>, + + /// wait <TIMEOUT> (ms) when checking for open ports in ms + #[arg(short, long, default_value_t = 250)] + pub timeout: u64, + /// include <NAME>. If <NAME> is a valid path or '-', hosts with be read from the file or /// stdin, respectivly. #[arg(short, long, id = "NAME")] diff --git a/src/main.rs b/src/main.rs index 7dd9390..a9fae7f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,9 @@ use std::{ io::{stdout, Write}, - net::IpAddr, + net::{IpAddr, SocketAddr, TcpStream}, str::FromStr, sync::{Arc, Mutex}, + time::Duration, }; use clap::Parser; @@ -48,29 +49,38 @@ async fn main() -> Result<(), anyhow::Error> { drop(tx); + let timeout = Duration::from_millis(config.timeout); + while let Some(ip_addr) = rx.recv().await { - join_set.spawn({ - let cache = cache.clone(); - async move { - let s = if config.resolve { - lookup_addr(&ip_addr).unwrap_or_else(|_| ip_addr.to_string()) - } else { - ip_addr.to_string() - }; + let cache = cache.clone(); + join_set.spawn(async move { + if !cache.lock()?.insert(ip_addr.to_string()) { + return Ok(()); + } + + if let Some(p) = config.port { + let addr = SocketAddr::from((ip_addr, p)); + if TcpStream::connect_timeout(&addr, timeout).is_err() { + return Ok(()); + } + } + if let Some(s) = config.resolve.then(|| lookup_addr(&ip_addr).ok()).flatten() { + let mut stdout = stdout().lock(); if cache.lock()?.insert(s.clone()) { - let mut stdout = stdout().lock(); writeln!(stdout, "{s}")?; } - Ok(()) - } + } else { + let mut stdout = stdout().lock(); + writeln!(stdout, "{ip_addr}")?; + }; + + Ok(()) }); } while let Some(res) = join_set.join_next().await { - if let Err(err) = res { - eprintln!("{err}") - } + res??; } Ok(()) |