From 2b4bb6a49d9f2c77ec1044d4aa2548f70f7d48fe Mon Sep 17 00:00:00 2001
From: Peter Dwyer <peter.dwyer@clanwilliamhealth.com>
Date: Tue, 5 Jul 2022 18:26:39 +0100
Subject: [PATCH] redone to parse host and port combo.

---
 src/main.rs | 45 +++++++++++++++++++++++++++++++++------------
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 5237eab..3bf321c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,4 @@
-use std::net::{SocketAddr, IpAddr};
+use std::net::{SocketAddr, IpAddr, ToSocketAddrs};
 use std::net::TcpStream;
 use std::{process, thread};
 use std::fmt::{Debug};
@@ -43,18 +43,21 @@ struct Args {
 
 fn main() {
     let args = Args::parse();
-    let target_result = get_target_ip(&args.host, &args.v4, &args.v6);
+    let timeout = Duration::from_millis(args.timeout.into());
+    let delay = chrono::Duration::milliseconds(args.delay.into());
+    let target_result = get_target_socket(&args);
+
     if target_result.is_err() {
-        println!("Error: {}", target_result.unwrap_err());
+        println!("{}", target_result.unwrap_err());
         process::exit(0);
     }
-    let target_ip: IpAddr = target_result.unwrap();
-    let delay = chrono::Duration::milliseconds(1000);
-    let timeout = Duration::from_millis(args.timeout.into());
+
+    let target_sock = target_result.unwrap();
+
     if args.repeat {
         loop {
             let start = Local::now();
-            let _result = ping(&target_ip, &args.port, &timeout);
+            let _result = ping(&target_sock, &timeout);
             let end = Local::now();
             let delta = end - start;
             if delta < delay {
@@ -62,7 +65,26 @@ fn main() {
             }
         }
     } else {
-        let _result = ping(&target_ip, &args.port, &timeout);
+        let _result = ping(&target_sock, &timeout);
+    }
+}
+
+fn get_target_socket(args: &Args) -> Result<SocketAddr, String> {
+    if let Ok(sock_iter) = args.host.to_socket_addrs() {
+        for sock in sock_iter {
+            if !args.v4 && !args.v6 {
+                return Ok(sock)
+            } else if args.v4 && sock.is_ipv4() {
+                return Ok(sock)
+            } else if args.v6 && sock.is_ipv6() {
+                return Ok(sock);
+            }
+        }
+        Err(format!("Error: Cannot parse host"))
+    } else if let Ok(target_result) = get_target_ip(&args.host, &args.v4, &args.v6){
+        return Ok(SocketAddr::new(target_result, args.port))
+    } else {
+        Err(format!("Error: Unable to resolve Host"))
     }
 }
 
@@ -89,13 +111,12 @@ fn get_target_ip(s: &str, v4: &bool, v6: &bool) -> Result<IpAddr, String> {
     }
 }
 
-fn ping(host: &IpAddr, port: &u16, timeout: &Duration) -> Result<(), ()> {
-    let target = SocketAddr::from((*host, *port));
+fn ping(target: &SocketAddr, timeout: &Duration) -> Result<(), ()> {
     if let Ok(_stream) = TcpStream::connect_timeout(&target, *timeout) {
-        println!("{} - {} - Connection Successful", Local::now().format("%Y-%m-%d %H:%M:%S"), host);
+        println!("{} - {} - Connection Successful", Local::now().format("%Y-%m-%d %H:%M:%S"), target);
         Ok(())
     } else {
-        println!("{} - {} - Connection Failed", Local::now().format("%Y-%m-%d %H:%M:%S"), host);
+        println!("{} - {} - Connection Failed", Local::now().format("%Y-%m-%d %H:%M:%S"), target);
         Err(())
     }
 }
\ No newline at end of file