1 files changed, 49 insertions, 40 deletions

diff --git a/src/main.rs b/src/main.rs
index 5fc8435..7dd9390 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,66 +1,75 @@
 use std::{
-    collections::HashSet,
     io::{stdout, Write},
+    net::IpAddr,
+    str::FromStr,
+    sync::{Arc, Mutex},
 };
 
 use clap::Parser;
 use config::Config;
-use tokio::sync::mpsc;
+use dns_lookup::{lookup_addr, lookup_host};
+use tokio::{sync::mpsc, task::JoinSet};
 
 mod config;
 
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
     let config = Config::parse();
+    let mut join_set = JoinSet::new();
 
+    let excluded = config.excluded()?;
+
+    let cache = Arc::new(Mutex::new(excluded));
     let (tx, mut rx) = mpsc::channel(100);
 
-    for host in config
-        .include_file
-        .into_iter()
-        .flat_map(|i| i.hosts())
-        .flatten()
-        .chain(config.include)
-    {
-        tx.send(host).await?;
-    }
+    for host in config.included()? {
+        match IpAddr::from_str(&host) {
+            Ok(ip) => tx.send(ip).await?,
+            Err(_) => {
+                if cache.lock().unwrap().insert(host.clone()) {
+                    let mut stdout = stdout().lock();
+                    writeln!(stdout, "{host}")?;
+                }
 
-    sshr::tmux::sessions(tx.clone(), &config.socket).await?;
+                if config.resolve {
+                    let cache = cache.clone();
+                    join_set.spawn(async move {
+                        for ip in lookup_host(&host)? {
+                            cache.lock()?.insert(ip.to_string());
+                        }
+                        Ok(())
+                    });
+                }
+            }
+        }
+    }
 
-    tokio::spawn(sshr::netlink::neighbours(tx.clone()));
+    join_set.spawn(sshr::netlink::neighbours(tx.clone()));
 
     drop(tx);
 
-    let mut cache = HashSet::new();
-
-    for host in config
-        .exclude_file
-        .into_iter()
-        .flat_map(|i| i.hosts())
-        .flatten()
-        .chain(config.exclude)
-    {
-        cache.insert(host);
-    }
-
-    let mut stdout = stdout();
-    while let Some(host) = rx.recv().await {
-        if !cache.insert(host.clone()) {
-            continue;
-        }
+    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 resolved = config.resolve.then_some(host.resolve().ok()).flatten();
-
-        if let Some(res) = resolved.clone() {
-            if !cache.insert(res) {
-                continue;
+                if cache.lock()?.insert(s.clone()) {
+                    let mut stdout = stdout().lock();
+                    writeln!(stdout, "{s}")?;
+                }
+                Ok(())
             }
-        }
+        });
+    }
 
-        if let Some(sshr::Host::Hostname(r)) = resolved {
-            writeln!(stdout, "{r}")?;
-        } else {
-            writeln!(stdout, "{host}")?;
+    while let Some(res) = join_set.join_next().await {
+        if let Err(err) = res {
+            eprintln!("{err}")
         }
     }