Rust Network Port Scanner Improved

Portscanner2.rs

use clap::{Parser, Subcommand, ArgEnum};
use serde::{Serialize, Deserialize};
use std::fs::{self, File};
use std::io::{Write, BufReader};
use std::net::ToSocketAddrs;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::time::timeout;
use futures::stream::{FuturesUnordered, StreamExt};
use log::{info, error};
use simplelog::*;
use anyhow::{Result, Context};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

const HISTORY_FILE: &str = "scan_history.json";
const LOG_FILE: &str = "port_scanner.log";

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Scan ports on a target
    Scan {
        /// Target IP or hostname
        target: String,

        /// Port range, e.g. 1-1024
        port_range: String,

        /// Timeout per connection attempt in seconds [default: 1]
        #[arg(short, long, default_value_t = 1)]
        timeout: u64,

        /// Output format
        #[arg(short, long, default_value_t = OutputFormat::Plain)]
        output: OutputFormat,

        /// Dry run mode (show what would be scanned, don't scan)
        #[arg(long)]
        dry_run: bool,
    },

    /// Show history of last scans
    History,

    /// Undo last scan (delete last saved history)
    Undo,
}

#[derive(ArgEnum, Clone, Copy)]
enum OutputFormat {
    Plain,
    Json,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
struct ScanResult {
    target: String,
    open_ports: Vec<PortInfo>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
struct PortInfo {
    port: u16,
    service: Option<String>,
    fingerprint: Option<String>, // <-- new field for service fingerprinting
}

fn parse_port_range(range: &str) -> Result<(u16, u16)> {
    let parts: Vec<&str> = range.split('-').collect();
    if parts.len() != 2 {
        anyhow::bail!("Invalid port range format, expected start-end");
    }
    let start = parts[0].parse::<u16>()?;
    let end = parts[1].parse::<u16>()?;
    if start > end {
        anyhow::bail!("Start port must be <= end port");
    }
    Ok((start, end))
}

fn get_common_service(port: u16) -> Option<&'static str> {
    match port {
        20 => Some("FTP Data"),
        21 => Some("FTP Control"),
        22 => Some("SSH"),
        23 => Some("Telnet"),
        25 => Some("SMTP"),
        53 => Some("DNS"),
        80 => Some("HTTP"),
        110 => Some("POP3"),
        143 => Some("IMAP"),
        443 => Some("HTTPS"),
        3306 => Some("MySQL"),
        5432 => Some("PostgreSQL"),
        6379 => Some("Redis"),
        _ => None,
    }
}

async fn scan_port(addr: &str, port: u16, timeout_sec: u64) -> bool {
    let socket_str = format!("{}:{}", addr, port);
    let addrs_iter = match socket_str.to_socket_addrs() {
        Ok(mut iter) => iter,
        Err(_) => return false,
    };
    if let Some(socket_addr) = addrs_iter.into_iter().next() {
        let duration = Duration::from_secs(timeout_sec);
        let result = timeout(duration, TcpStream::connect(socket_addr)).await;
        result.is_ok() && result.unwrap().is_ok()
    } else {
        false
    }
}

async fn fingerprint_service(addr: &str, port: u16, timeout_sec: u64) -> Option<String> {
    let socket_str = format!("{}:{}", addr, port);
    let addrs_iter = socket_str.to_socket_addrs().ok()?;
    let socket_addr = addrs_iter.into_iter().next()?;

    let duration = Duration::from_secs(timeout_sec);
    let stream = timeout(duration, TcpStream::connect(socket_addr)).await.ok()??;

    // Probes for some common protocols
    let probe = match port {
        80 | 8080 | 8000 | 8888 => b"HEAD / HTTP/1.0\r\n\r\n",
        21 => b"FEAT\r\n",         // FTP
        25 => b"EHLO example.com\r\n",  // SMTP
        22 => b"",                // SSH usually sends banner immediately
        _ => b"",
    };

    let mut stream = stream;
    if !probe.is_empty() {
        if stream.write_all(probe).await.is_err() {
            return None;
        }
    }

    let mut buf = [0; 512];
    match timeout(duration, stream.read(&mut buf)).await {
        Ok(Ok(n)) if n > 0 => {
            let banner = String::from_utf8_lossy(&buf[..n]).to_string();

            // Basic banner checks
            if banner.contains("SSH") {
                Some(format!("SSH Service Banner: {}", banner.lines().next().unwrap_or("")))
            } else if banner.contains("HTTP") || banner.contains("html") {
                Some(format!("HTTP Service Banner: {}", banner.lines().next().unwrap_or("")))
            } else if banner.contains("FTP") {
                Some(format!("FTP Service Banner: {}", banner.lines().next().unwrap_or("")))
            } else if banner.contains("SMTP") {
                Some(format!("SMTP Service Banner: {}", banner.lines().next().unwrap_or("")))
            } else {
                Some(format!("Banner: {}", banner.lines().next().unwrap_or("")))
            }
        }
        _ => None,
    }
}

fn save_history(scan: &ScanResult) -> Result<()> {
    let json = serde_json::to_string_pretty(scan)?;
    let mut file = File::create(HISTORY_FILE)?;
    file.write_all(json.as_bytes())?;
    Ok(())
}

fn load_history() -> Result<ScanResult> {
    let file = File::open(HISTORY_FILE)?;
    let reader = BufReader::new(file);
    let scan: ScanResult = serde_json::from_reader(reader)?;
    Ok(scan)
}

fn delete_history() -> Result<()> {
    if std::path::Path::new(HISTORY_FILE).exists() {
        fs::remove_file(HISTORY_FILE)?;
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    CombinedLogger::init(
        vec![
            TermLogger::new(LevelFilter::Info, Config::default(), TerminalMode::Mixed, ColorChoice::Auto),
            WriteLogger::new(LevelFilter::Info, Config::default(), File::create(LOG_FILE)?),
        ]
    )?;

    let cli = Cli::parse();

    match cli.command {
        Commands::Scan { target, port_range, timeout, output, dry_run } => {
            info!("Starting scan on {} ports {}", target, port_range);
            println!("Starting scan on {} ports {}", target, port_range);

            let (start_port, end_port) = parse_port_range(&port_range)
                .context("Failed to parse port range")?;

            if dry_run {
                println!("Dry run mode: Would scan ports {} to {} on {}", start_port, end_port, target);
                info!("Dry run: no scanning performed");
                return Ok(());
            }

            let mut futures = FuturesUnordered::new();

            for port in start_port..=end_port {
                futures.push(scan_port(&target, port, timeout).then(move |is_open| async move {
                    (port, is_open)
                }));
            }

            let mut open_ports = Vec::new();
            while let Some((port, is_open)) = futures.next().await {
                if is_open {
                    let service = get_common_service(port).map(String::from);
                    // Fingerprint the service asynchronously
                    let fingerprint = fingerprint_service(&target, port, timeout).await;

                    open_ports.push(PortInfo { port, service, fingerprint });
                    println!("Port {} is open", port);
                    if let Some(fp) = &open_ports.last().unwrap().fingerprint {
                        println!("  Fingerprint: {}", fp);
                    }
                    info!("Port {} open with fingerprint {:?}", port, open_ports.last().unwrap().fingerprint);
                }
            }

            let scan_result = ScanResult {
                target: target.clone(),
                open_ports,
            };

            save_history(&scan_result)?;
            info!("Scan completed and saved");

            match output {
                OutputFormat::Plain => {
                    println!("\nScan result for {}", scan_result.target);
                    if scan_result.open_ports.is_empty() {
                        println!("No open ports found.");
                    } else {
                        for port_info in &scan_result.open_ports {
                            let mut line = format!("Port {}", port_info.port);
                            if let Some(service) = &port_info.service {
                                line.push_str(&format!(" - {}", service));
                            }
                            if let Some(fp) = &port_info.fingerprint {
                                line.push_str(&format!(" | {}", fp));
                            }
                            println!("{}", line);
                        }
                    }
                }
                OutputFormat::Json => {
                    let json = serde_json::to_string_pretty(&scan_result)?;
                    println!("{}", json);
                }
            }
        }

        Commands::History => {
            match load_history() {
                Ok(scan) => {
                    println!("Last scan result for: {}", scan.target);
                    if scan.open_ports.is_empty() {
                        println!("No open ports found.");
                    } else {
                        for port_info in scan.open_ports {
                            let mut line = format!("Port {}", port_info.port);
                            if let Some(service) = port_info.service {
                                line.push_str(&format!(" - {}", service));
                            }
                            if let Some(fp) = port_info.fingerprint {
                                line.push_str(&format!(" | {}", fp));
                            }
                            println!("{}", line);
                        }
                    }
                }
                Err(_) => {
                    println!("No scan history found.");
                }
            }
        }

        Commands::Undo => {
            delete_history()?;
            println!("Last scan history deleted.");
            info!("History deleted by user");
        }
    }

    Ok(())
}