fuego 0.1.1

use crate::records::BannerRecord;
use std::error::Error;
use std::io::{ErrorKind, Read, Write};
use std::net::{IpAddr, SocketAddr, TcpStream};
use std::time::Duration;

pub struct Scnr {}

impl Scnr {
    /*

    basic scan

     */
    ///Timeout scan is a tcp connect scan with a timeout.
    /// More suitable if there are going to be a lot of closed ports.
    /// address is IPv6 or IPv4 format, ports is ports to scan and timeout is the
    /// amount of time you are willing to wait.
    pub async fn scan(
        address: IpAddr,
        ports: Vec<u16>,
        timeout: u64,
    ) -> Result<Vec<u16>, Box<dyn Error>> {
        let mut open_port_vector: Vec<u16> = vec![];
        let timeout = Duration::from_millis(timeout);

        for port in ports {
            let socket_address: SocketAddr = SocketAddr::new(address, port);

            if (try_connect(socket_address, timeout).await).is_ok() {
                open_port_vector.push(port)
            }
        }

        Ok(open_port_vector)
    }

    /**
      Props:
       address: IPv4/v6 address being scanned

       ports: port vector to scan,

       timeout: tcp stream connect timeout

       This function does a tcp connect try and then attempts to grab a banner from the
       open port.
    */

    pub async fn banner_scan(
        address: IpAddr,
        ports: Vec<u16>,
        timeout: u64,
    ) -> Result<Vec<BannerRecord>, Box<dyn Error>> {
        let mut banner_vector: Vec<BannerRecord> = vec![];
        let duration = Duration::from_millis(timeout);

        for port in ports {
            let socket_address: SocketAddr = SocketAddr::new(address, port);

            if let Ok(stream) = try_connect(socket_address, duration).await {
                match Grabber::stream_grab(stream, 150).await {
                    Ok(banner) => {
                        let banner = if banner.is_empty() {
                            "nil".to_string()
                        } else {
                            parse_banner(&banner)
                        };

                        banner_vector.push(BannerRecord { port, banner });
                    }
                    Err(_) => {
                        banner_vector.push(BannerRecord {
                            port,
                            banner: "nil".to_string(),
                        });
                    }
                }
            }
        }

        Ok(banner_vector)
    }
}

pub struct Grabber {}

impl Grabber {
    /*

    tcp send string

     */
    ///Sends a string with a connect timeout.
    /// Address is IPv4 and IPv6 &str.
    ///msg is the message to send.
    ///timeout is time to wait to connect.
    ///Returns string, useful for poking around.
    pub async fn string_grab(
        address: IpAddr,
        port: u16,
        msg: &str,
        timeout: u64,
    ) -> Result<String, Box<dyn Error>> {
        let socket_address: SocketAddr = SocketAddr::new(address, port);
        let timeout = Duration::from_millis(timeout);
        let rw_duration = Duration::from_millis(1000);
        let mut response_shell = String::new();
        let req_bytes = msg.as_bytes();
        let mut stream = match try_connect(socket_address, timeout).await {
            Ok(stream) => stream,
            Err(err) => return Err(Box::from(err.to_string())),
        };

        stream.set_read_timeout(Option::from(rw_duration))?;
        stream.set_write_timeout(Option::from(rw_duration))?;
        stream.write_all(req_bytes).expect("could not write bytes");
        stream
            .read_to_string(&mut response_shell)
            .expect("could not read string");
        Ok(parse_banner(&response_shell))
    }

    /*

    tcp send string
    thanks polonium https://github.com/creekorful/polonium
     */
    ///does a banner grab.
    /// Address is IPv6 or v6 &str.
    /// port is i32, its the port to grab a banner.
    /// timeout is for the read/write/connect.
    pub async fn banner_grab(
        address: IpAddr,
        port: u16,
        timeout: u64,
    ) -> Result<String, Box<dyn Error>> {
        let server_socket: SocketAddr = SocketAddr::new(address, port);

        let duration = Duration::from_millis(timeout);
        let rw_duration = Duration::from_millis(1000);

        let mut banner = String::new();
        let mut stream =
            TcpStream::connect_timeout(&server_socket, duration).expect("could not connect");

        stream.set_read_timeout(Option::from(rw_duration))?;
        stream.set_write_timeout(Option::from(rw_duration))?;

        let result = stream.read_to_string(&mut banner);
        if result.is_ok() && !banner.is_empty() {
            return Ok(banner);
        }

        // If timeout related error happens, do not fails
        // because we may need to talk first
        let error = result.err().unwrap();
        if error.kind() != ErrorKind::WouldBlock {
            return Err(error.into());
        }

        if banner.is_empty() {
            stream.write_all("GET / HTTP/1.1\n\n".as_ref())?;
            stream.read_to_string(&mut banner)?;
        }

        Ok(parse_banner(&banner))
    }
    /*
       Props:

       stream: the input stream being used to banner grab

       timeout: the read/write timeouts

       This is so we can grab a banner and also set a timeout without a way to handle the
       error of a timeout that might block
    */
    pub async fn stream_grab(
        mut stream: TcpStream,
        timeout: u64,
    ) -> Result<String, Box<dyn Error>> {
        let rw_duration = Duration::from_millis(timeout);
        let mut banner = String::new();

        stream.set_read_timeout(Option::from(rw_duration))?;
        stream.set_write_timeout(Option::from(rw_duration))?;

        let result = stream.read_to_string(&mut banner);
        if result.is_ok() && !banner.is_empty() {
            return Ok(banner);
        }

        let error = result.err().unwrap();
        if error.kind() != ErrorKind::WouldBlock {
            return Err(error.into());
        }

        if banner.is_empty() {
            stream.write_all("HEAD / HTTP/1.1\n\n".as_ref())?;
            stream.read_to_string(&mut banner)?;
        }

        Ok(banner)
    }
}

fn parse_banner(banner: &str) -> String {
    let b_vec = Vec::from_iter(banner.split(' '));
    String::from(b_vec[0])
}

async fn try_connect(
    socketaddr: SocketAddr,
    timeout: Duration,
) -> Result<TcpStream, std::io::Error> {
    match TcpStream::connect_timeout(&socketaddr, timeout) {
        Ok(stream) => Ok(stream),
        Err(err) => Err(err),
    }
}