Added range feature and and dns resolving

src/main.rs

@@ -1,10 +1,11 @@
 use clap::Parser;
 use std::io::prelude::*;
 use std::io::{self, Read};
-use std::net::{SocketAddr, TcpStream};
+use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
 use std::time::{Duration, Instant};
 use std::thread;
 use atty::Stream;
+use std::io::BufReader;
 
 
 fn main() -> std::io::Result<()> {
@@ -12,7 +13,7 @@ fn main() -> std::io::Result<()> {
     let address = args.address;
     let timeout = Duration::from_secs(args.timeout);
     let sleep = Duration::from_millis(args.wait);
-    let ports = args.ports;
+    let ports = handle_ports(args.ports.clone(), &args.range);
     let threads = args.cthreads;
     
     //Reading the payload from stdin (if applicable)
@@ -34,17 +35,33 @@ fn main() -> std::io::Result<()> {
 
 }
 
+//Returns a correct list of ports, according to range.
+//If range is true, we send a range of ports. Otherwise
+//we send the original back
+fn handle_ports(ports: Vec<u32>, range: &bool) -> Vec<u32> {
+    if *range {
+        if ports.iter().count() != 2 {
+            panic!("A range must be given only a start and an end value.");
+        }
+        let ret: Vec<u32> = (ports[0]..ports[1]).collect();
+        ret
+    } else {
+        ports
+    }
+}
+
+
+
 ///Scan the ports on address by attempting a tcp connection.
 fn tcp_scan(address: String, timeout: Duration, sleep: Duration, ports: Vec<u32>) -> std::io::Result<()> {
     println!("Scanning {}", address);
 
     for i in ports {
-        let socket = format!("{}:{}", &address, i);
-        let socket = socket.parse::<SocketAddr>().unwrap();
+        let socket = resolve_add(&address.as_str(), &i);
         match TcpStream::connect_timeout(&socket, timeout) {
-            Ok(mut s) => {
+            Ok(mut _s) => {
                 println!("TCP/{} => ACCEPT", i);
-                s.write(b"Some padding")?;
+                //s.write(b"Some padding")?;
             }
             Err(_) => {
                 println!("TCP/{} => REJECT", i);
@@ -58,8 +75,18 @@ fn tcp_scan(address: String, timeout: Duration, sleep: Duration, ports: Vec<u32>
     Ok(())
 }
 
-///Try to stress the remote machine by attempting to every given port at once and sending random
-///data
+///Takes either an hostname or an IP address and port, and resolve to a socket address.
+///DNS recordings can lead to multiple entries : we take the first one.
+fn resolve_add(host: &str, port: &u32) -> SocketAddr {
+    let details = format!("{}:{}", host, port);
+    let server: Vec<_> = details
+        .to_socket_addrs()
+        .expect("Can't resolve your address")
+        .collect();
+    return server[0];
+}
+
+///Try to stress the remote machine by attempting to every given port at once and sending payload
 fn tcp_stress(payload: String, address: String, timeout: Duration, interval: Duration, ports: Vec<u32>, threads: u32) -> std::io::Result<()> {
     let mut handles = vec![];
 
@@ -72,20 +99,33 @@ fn tcp_stress(payload: String, address: String, timeout: Duration, interval: Dur
             let pay = payload.clone();
             let int = interval * 1000;
             let end_time = Instant::now() + int;
+            //pay.push_str("\r\n");
 
 
             handles.push(thread::spawn( move || {
-                let socket = format!("{}:{}", &add, &port);
-                let socket = socket.parse::<SocketAddr>().unwrap();
+                let socket = resolve_add(&add.as_str(), &port);
                 loop {
                     //Check if it is time to stop
                     if Instant::now() >= end_time { break; }
                     match TcpStream::connect_timeout(&socket, timeout) {
                         Ok(mut s) => {
-                            s.write(pay.as_bytes()).unwrap();
+                            //Remote port answers
+                            //Writing
+                            s.write_all(pay.as_bytes()).unwrap();
                             s.flush().unwrap();
-                            //s.read(&mut [0; 128]).unwrap();
+
+                            //Reading
+                            s.set_read_timeout(Some(Duration::from_millis(500))).unwrap();
+                            let mut reader = BufReader::new(s.try_clone().unwrap());
+                            let mut answer = String::new();
+                            let count = reader.read_line(&mut answer).unwrap_or(0);
+                            if count > 0 {
+                                println!("Read {} bytes for port {} -> {}", count, &port, answer);
+                            }
+
+                            //Sleeping
                             thread::sleep(Duration::from_millis(10));
+
                         },
                         Err(_) => (),
                     }
@@ -107,7 +147,7 @@ fn tcp_stress(payload: String, address: String, timeout: Duration, interval: Dur
 #[derive(Debug, Parser)]
 #[command(author, version, about, long_about = None)]
 struct Args {
-    ///IP address to scan
+    ///IP address or hostname to scan
     #[arg(short, long)]
     address: String,
     ///Timeout for each connection in seconds
@@ -120,6 +160,9 @@ struct Args {
     ///Ports to stress / scan, separated by commas (22,80)
     #[arg(short, long, value_parser, num_args=1.., value_delimiter=',')]
     ports: Vec<u32>,
+    ///Set this flag to treat the ports as a range rather than a list
+    #[arg(long, short, action, default_value_t = false)]
+    range: bool,
     ///Set this flag to stress the ports instead of scanning them
     #[arg(long, short, action, default_value_t = false)]
     stress: bool,