Resolve DNS target

1 files changed, 21 insertions, 5 deletions

diff --git a/src/main.rs b/src/main.rs
index d7fae43..99db51b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,5 @@
 use clap::Parser;
-use std::net::IpAddr;
+use std::net::{IpAddr, ToSocketAddrs};
 use tokio::sync::mpsc;
 use tokio::sync::mpsc::error::TryRecvError;
 use tokio::time::sleep;
@@ -109,10 +109,26 @@ async fn main() {
     let args = Args::parse();
 
     let target: IpAddr = if let Some(target_str) = args.target.clone() {
-        parse_ip(&target_str).unwrap_or_else(|e| {
-            eprintln!("Error: {}", e);
-            std::process::exit(1);
-        })
+        if let Ok(ip) = parse_ip(&target_str) {
+            ip
+        } else {
+            // Try DNS resolution
+            let dns_str = format!("{}:0", target_str);
+            match dns_str.to_socket_addrs() {
+                Ok(mut addrs) => {
+                    if let Some(addr) = addrs.next() {
+                        addr.ip()
+                    } else {
+                        eprintln!("Error: no DNS addresses found");
+                        std::process::exit(1);
+                    }
+                }
+                Err(e) => {
+                    eprintln!("Error: DNS resolution failed: {}", e);
+                    std::process::exit(1);
+                }
+            }
+        }
     } else {
         eprintln!("Error: target address is required");
         std::process::exit(1);