kxf 0.1.0

use clap::Parser;
use nix::sys::signal::{Signal, kill};
use nix::unistd::Pid;
use std::collections::HashSet;
use std::error::Error;
use std::fs;
use std::io;
use std::path::Path;

#[derive(Parser)]
#[command(
    name = "kxf",
    version,
    about = "根据端口号查找并终止占用该端口的进程",
    arg_required_else_help = true
)]
struct Cli {
    /// 需要释放的端口号
    #[arg(value_name = "PORT")]
    port: u16,
}

#[derive(Debug, Clone)]
struct ProcessInfo {
    pid: u32,
    command: String,
}

fn main() {
    if let Err(err) = run() {
        eprintln!("错误: {err}");
        std::process::exit(1);
    }
}

fn run() -> Result<(), Box<dyn Error>> {
    let cli = Cli::parse();
    let inodes = collect_socket_inodes(cli.port)?;

    if inodes.is_empty() {
        println!("端口 {} 上没有找到任何进程。", cli.port);
        return Ok(());
    }

    let mut processes = find_processes_using_inodes(&inodes)?;
    if processes.is_empty() {
        println!("端口 {} 上没有找到任何进程。", cli.port);
        return Ok(());
    }
    processes.sort_by_key(|proc| proc.pid);

    let mut killed = Vec::new();
    let mut failed = Vec::new();

    for process in processes {
        match kill_process(process.pid) {
            Ok(_) => killed.push(process),
            Err(err) => failed.push((process, err.to_string())),
        }
    }

    if !killed.is_empty() {
        println!("成功杀死以下进程:");
        for proc in &killed {
            println!("  PID {:>5} | {}", proc.pid, proc.command);
        }
    }

    if !failed.is_empty() {
        println!("以下进程无法杀死，可能需要更高权限:");
        for (proc, err) in &failed {
            println!("  PID {:>5} | {} | 错误: {}", proc.pid, proc.command, err);
        }
    }

    if killed.is_empty() && failed.is_empty() {
        println!("端口 {} 上的进程在扫描与执行之间已经退出。", cli.port);
    }

    Ok(())
}

fn collect_socket_inodes(port: u16) -> io::Result<HashSet<u64>> {
    const PROC_NET_FILES: &[&str] = &[
        "/proc/net/tcp",
        "/proc/net/tcp6",
        "/proc/net/udp",
        "/proc/net/udp6",
    ];

    let mut inodes = HashSet::new();

    for path in PROC_NET_FILES {
        if let Err(err) = parse_proc_net_file(path, port, &mut inodes) {
            if err.kind() != io::ErrorKind::NotFound {
                return Err(err);
            }
        }
    }

    Ok(inodes)
}

fn parse_proc_net_file(path: &str, port: u16, inodes: &mut HashSet<u64>) -> io::Result<()> {
    let contents = fs::read_to_string(path)?;
    for line in contents.lines().skip(1) {
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() <= 9 {
            continue;
        }

        if port_matches(parts[1], port) {
            if let Ok(inode) = parts[9].parse::<u64>() {
                inodes.insert(inode);
            }
        }
    }

    Ok(())
}

fn port_matches(field: &str, port: u16) -> bool {
    match field.rsplit_once(':') {
        Some((_, port_hex)) => match u16::from_str_radix(port_hex, 16) {
            Ok(value) => value == port,
            Err(_) => false,
        },
        None => false,
    }
}

fn find_processes_using_inodes(inodes: &HashSet<u64>) -> io::Result<Vec<ProcessInfo>> {
    let mut processes = Vec::new();
    let proc_dir = fs::read_dir("/proc")?;

    for entry in proc_dir {
        let entry = match entry {
            Ok(e) => e,
            Err(_) => continue,
        };

        let pid = match entry
            .file_name()
            .to_str()
            .and_then(|s| s.parse::<u32>().ok())
        {
            Some(pid) => pid,
            None => continue,
        };

        let fd_dir = entry.path().join("fd");
        let fd_entries = match fs::read_dir(fd_dir) {
            Ok(it) => it,
            Err(_) => continue,
        };

        let mut matched = false;
        for fd_entry in fd_entries {
            let fd_entry = match fd_entry {
                Ok(e) => e,
                Err(_) => continue,
            };

            let target = match fs::read_link(fd_entry.path()) {
                Ok(t) => t,
                Err(_) => continue,
            };

            if let Some(inode) = parse_socket_inode(&target) {
                if inodes.contains(&inode) {
                    matched = true;
                    break;
                }
            }
        }

        if matched {
            let command = read_command(pid);
            processes.push(ProcessInfo { pid, command });
        }
    }

    Ok(processes)
}

fn parse_socket_inode(target: &Path) -> Option<u64> {
    let text = target.to_string_lossy();
    text.strip_prefix("socket:[")
        .and_then(|rest| rest.strip_suffix(']'))
        .and_then(|num| num.parse::<u64>().ok())
}

fn read_command(pid: u32) -> String {
    let cmdline_path = format!("/proc/{pid}/cmdline");
    if let Ok(bytes) = fs::read(&cmdline_path) {
        let parts: Vec<String> = bytes
            .split(|b| *b == 0)
            .filter(|segment| !segment.is_empty())
            .map(|segment| String::from_utf8_lossy(segment).to_string())
            .collect();

        if !parts.is_empty() {
            return parts.join(" ");
        }
    }

    let comm_path = format!("/proc/{pid}/comm");
    fs::read_to_string(comm_path)
        .map(|s| s.trim().to_string())
        .unwrap_or_else(|_| "<unknown>".to_string())
}

fn kill_process(pid: u32) -> Result<(), nix::errno::Errno> {
    kill(Pid::from_raw(pid as i32), Signal::SIGKILL)
}