Portnut/src/main.rs

use clap::Parser;
use std::io::prelude::*;
use std::io::{self, Read};
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<()> {
    let args = Args::parse();
    let address = args.address;
    let timeout = Duration::from_secs(args.timeout);
    let sleep = Duration::from_millis(args.wait);
    let ports = handle_ports(args.ports.clone(), &args.range);
    let threads = args.cthreads;
    
    //Reading the payload from stdin (if applicable)
    let mut payload = String::new();
    let aatty = atty::is(Stream::Stdin);

    if !aatty {
        io::stdin().read_to_string(&mut payload)?;
    } else {
        payload = "SAKAMOTO".to_string();
    }

    if args.stress {
        tcp_stress(payload, address, timeout, sleep, ports, threads)?;
    } else {
        tcp_scan(address, timeout, sleep, ports)?;
    }
    Ok(())

}

//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 = resolve_add(&address.as_str(), &i);
        match TcpStream::connect_timeout(&socket, timeout) {
            Ok(mut _s) => {
                println!("TCP/{} => ACCEPT", i);
                //s.write(b"Some padding")?;
            }
            Err(_) => {
                println!("TCP/{} => REJECT", i);
            }
        }
        
        thread::sleep(sleep.clone());
        //stream.write(b"Some great padding")?;
        //stream.read(&mut [0; 128])?;
    }
    Ok(())
}

///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![];

    for port in ports {
        tcp_scan(address.clone(), timeout, interval, vec![port])?;
        let mut threads_left = threads;

        loop {
            let add = address.clone();
            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 = 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) => {
                            //Remote port answers
                            //Writing
                            s.write_all(pay.as_bytes()).unwrap();
                            s.flush().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(_) => (),
                    }
                }
            }));
            *&mut threads_left -=1;
            if threads_left == 0 { break; }
        }
    }
    println!("\nStressing {} using {} concurrent threads per port for {:?} \n", address, threads, interval * 1000);
    for handle in handles {
        let _ = handle.join();
    }
    Ok(())
}

///A simple TCP port scanner / stresser.
///If stressing, you can pass a payload via a pipe.
#[derive(Debug, Parser)]
#[command(author, version, about, long_about = None)]
struct Args {
    ///IP address or hostname to scan
    #[arg(short, long)]
    address: String,
    ///Timeout for each connection in seconds
    #[arg(short, long, default_value_t = 1)]
    timeout: u64,
    ///Number of milliseconds to wait in between scans when scanning
    ///OR duration of stress when stress testing (in seconds)
    #[arg(short, long, default_value_t = 30)]
    wait: u64,
    ///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,
    ///How many threads per port when stressing
    ///(Concurrent Threads)
    #[arg(short, long, default_value_t = 5)]
    cthreads: u32,
}