vkit 0.1.4

Fast Rust dev CLI: manage git worktrees, Node ports, run scripts, install & sync VS Code / Cursor extensions.
use std::collections::HashMap;
use std::process::Command;
use std::thread;
use std::time::Duration;

use anyhow::{anyhow, Result};
use nix::errno::Errno;
use nix::sys::signal::{kill, Signal};
use nix::unistd::Pid;

/// 一个监听端口条目:端口 + 进程 + 命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PortEntry {
    pub port: u16,
    pub pid: i32,
    /// lsof COMMAND(进程名)。
    pub command: String,
    /// ps 补全的完整启动命令。
    pub full_command: String,
    /// 进程工作目录(启动服务的目录),macOS lsof -d cwd 取得;取不到为空串。
    pub cwd: String,
}

/// 解析 `lsof -F pcn` 输出为端口条目(仅保留 node,按 port/pid 升序)。
///
/// `-F` 每行首字符是标签:`p`=pid、`c`=command、`n`=network name。
pub fn parse_lsof_output(stdout: &str) -> Vec<PortEntry> {
    let mut entries: Vec<PortEntry> = Vec::new();
    let mut pid: Option<i32> = None;
    let mut command: Option<String> = None;

    for line in stdout.lines() {
        if line.is_empty() {
            continue;
        }

        let (tag, value) = line.split_at(1);
        match tag {
            "p" => pid = value.parse::<i32>().ok(),
            "c" => command = Some(value.to_string()),
            "n" => {
                let (Some(pid), Some(command)) = (pid, command.as_ref()) else {
                    continue;
                };
                if let Some(port) = extract_port(value) {
                    entries.push(PortEntry {
                        port,
                        pid,
                        command: command.clone(),
                        full_command: String::new(),
                        cwd: String::new(),
                    });
                }
            }
            _ => {}
        }
    }

    entries.retain(|entry| entry.command.eq_ignore_ascii_case("node"));
    entries.sort_by(|a, b| a.port.cmp(&b.port).then(a.pid.cmp(&b.pid)));
    entries
}

fn extract_port(network_name: &str) -> Option<u16> {
    let colon = network_name.rfind(':')?;
    network_name[colon + 1..].parse::<u16>().ok()
}

/// 大小写不敏感的子串匹配过滤(命中 port / command / full_command 任一)。
pub fn filter_ports(entries: &[PortEntry], query: &str) -> Vec<PortEntry> {
    let query = query.trim().to_lowercase();
    if query.is_empty() {
        return entries.to_vec();
    }

    entries
        .iter()
        .filter(|entry| {
            entry.port.to_string().contains(&query)
                || entry.command.to_lowercase().contains(&query)
                || entry.full_command.to_lowercase().contains(&query)
                || entry.cwd.to_lowercase().contains(&query)
        })
        .cloned()
        .collect()
}

/// 按显示宽度截断文本,超出部分用 `…` 结尾(按字符计数)。
pub fn truncate_text(text: &str, max_length: usize) -> String {
    if max_length == 0 {
        return String::new();
    }

    if text.chars().count() <= max_length {
        return text.to_string();
    }

    if max_length == 1 {
        return "".to_string();
    }

    let head: String = text.chars().take(max_length - 1).collect();
    format!("{head}")
}

/// 扫描正在监听(TCP LISTEN)的 node 端口。
pub fn list_listening_ports() -> Result<Vec<PortEntry>> {
    let output = Command::new("lsof")
        .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn", "-c", "node"])
        .output();

    let output = match output {
        Ok(output) => output,
        Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
            return Err(anyhow!("未找到 lsof 命令。此工具仅支持 macOS。"));
        }
        Err(err) => return Err(anyhow!("执行 lsof 失败:{err}")),
    };

    let stdout = String::from_utf8_lossy(&output.stdout);
    let code = output.status.code().unwrap_or(-1);

    // lsof 退出码 1 表示「没有匹配结果」,视为正常。
    if !output.status.success() && code != 1 {
        let stderr = String::from_utf8_lossy(&output.stderr);
        let detail = {
            let stderr = stderr.trim();
            if stderr.is_empty() {
                stdout.trim()
            } else {
                stderr
            }
        };

        if detail.to_lowercase().contains("permission denied") {
            return Err(anyhow!("没有权限读取端口信息。请尝试用更高权限运行。"));
        }

        return Err(if detail.is_empty() {
            anyhow!("lsof 执行失败(退出码 {code})")
        } else {
            anyhow!("{detail}")
        });
    }

    let mut entries = parse_lsof_output(&stdout);
    if entries.is_empty() {
        return Ok(entries);
    }

    let full_commands = load_full_commands(&entries);
    for entry in &mut entries {
        entry.full_command = full_commands
            .get(&entry.pid)
            .cloned()
            .unwrap_or_else(|| entry.command.clone());
    }

    let cwds = load_cwds(&entries);
    for entry in &mut entries {
        entry.cwd = cwds.get(&entry.pid).cloned().unwrap_or_default();
    }

    Ok(entries)
}

/// 对去重 PID 批量跑一次 `ps` 补全完整命令;失败则返回空表(调用方回退进程名)。
fn load_full_commands(entries: &[PortEntry]) -> HashMap<i32, String> {
    let mut pids: Vec<i32> = entries.iter().map(|entry| entry.pid).collect();
    pids.sort_unstable();
    pids.dedup();

    let mut commands = HashMap::new();
    if pids.is_empty() {
        return commands;
    }

    let pid_arg = pids
        .iter()
        .map(|pid| pid.to_string())
        .collect::<Vec<_>>()
        .join(",");

    let output = Command::new("ps")
        .args(["-p", &pid_arg, "-o", "pid=", "-o", "command="])
        .output();

    let output = match output {
        Ok(output) if output.status.success() => output,
        _ => return commands,
    };

    let stdout = String::from_utf8_lossy(&output.stdout);
    for line in stdout.lines() {
        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }

        let mut parts = trimmed.splitn(2, char::is_whitespace);
        let Some(pid) = parts.next().and_then(|value| value.parse::<i32>().ok()) else {
            continue;
        };
        let command = parts.next().unwrap_or("").trim_start().to_string();
        commands.insert(pid, command);
    }

    commands
}

/// 对去重 PID 批量跑一次 `lsof -d cwd` 取进程工作目录;失败则返回空表(调用方回退空串)。
///
/// `-F pn` 输出里 `p` 行切换当前 pid,`n` 行记录该 pid 的路径。
fn load_cwds(entries: &[PortEntry]) -> HashMap<i32, String> {
    let mut pids: Vec<i32> = entries.iter().map(|entry| entry.pid).collect();
    pids.sort_unstable();
    pids.dedup();

    let mut cwds = HashMap::new();
    if pids.is_empty() {
        return cwds;
    }

    let pid_arg = pids
        .iter()
        .map(|pid| pid.to_string())
        .collect::<Vec<_>>()
        .join(",");

    let output = Command::new("lsof")
        .args(["-a", "-p", &pid_arg, "-d", "cwd", "-F", "pn"])
        .output();

    let output = match output {
        Ok(output) if output.status.success() => output,
        _ => return cwds,
    };

    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut pid: Option<i32> = None;
    for line in stdout.lines() {
        if line.is_empty() {
            continue;
        }
        let (tag, value) = line.split_at(1);
        match tag {
            "p" => pid = value.parse::<i32>().ok(),
            "n" => {
                if let Some(pid) = pid {
                    cwds.insert(pid, value.to_string());
                }
            }
            _ => {}
        }
    }

    cwds
}

/// 关闭进程:先 SIGTERM 优雅退出,3 秒未退再 SIGKILL 强制。
pub fn kill_pid(pid: i32) -> Result<()> {
    let target = Pid::from_raw(pid);

    match kill(target, Signal::SIGTERM) {
        Ok(()) => {}
        Err(Errno::ESRCH) => return Err(anyhow!("进程 {pid} 已不存在。")),
        Err(Errno::EPERM) => return Err(anyhow!("没有权限关闭进程 {pid}。")),
        Err(err) => return Err(anyhow!("无法关闭进程 {pid}:{err}")),
    }

    for _ in 0..30 {
        thread::sleep(Duration::from_millis(100));
        // 用 signal=None 探测进程是否仍存在。
        if let Err(Errno::ESRCH) = kill(target, None) {
            return Ok(());
        }
    }

    match kill(target, Signal::SIGKILL) {
        Ok(()) | Err(Errno::ESRCH) => Ok(()),
        Err(Errno::EPERM) => Err(anyhow!("没有权限关闭进程 {pid}。")),
        Err(err) => Err(anyhow!("无法强制关闭进程 {pid}:{err}")),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const SAMPLE: &str = "p90344\ncnode\nf20\nn*:3000\nf21\nn127.0.0.1:3001\np97605\ncnode\nf20\nn*:3002\np49594\ncnode\nf14\nn[::1]:5274\n";

    #[test]
    fn parse_extracts_ports() {
        let entries = parse_lsof_output(SAMPLE);
        assert_eq!(entries.len(), 4);
        assert_eq!(entries[0].port, 3000);
        assert_eq!(entries[0].pid, 90344);
        assert_eq!(entries[0].command, "node");
        assert_eq!(entries[1].port, 3001);
        assert_eq!(entries[1].pid, 90344);
        assert_eq!(entries[3].port, 5274);
        assert_eq!(entries[3].pid, 49594);
    }

    #[test]
    fn parse_empty_returns_empty() {
        assert!(parse_lsof_output("").is_empty());
    }

    #[test]
    fn parse_only_keeps_node() {
        let mixed = "p90344\ncnode\nf20\nn*:3000\np19758\ncssh\nf9\nn*:5432\n";
        let entries = parse_lsof_output(mixed);
        assert_eq!(entries.len(), 1);
        assert_eq!(entries[0].port, 3000);
    }

    fn entry(port: u16, pid: i32, command: &str, full: &str) -> PortEntry {
        PortEntry {
            port,
            pid,
            command: command.to_string(),
            full_command: full.to_string(),
            cwd: String::new(),
        }
    }

    #[test]
    fn filter_matches_port_and_name() {
        let entries = vec![
            entry(3000, 1, "node", "node server.js"),
            entry(5432, 2, "postgres", "/usr/local/bin/postgres"),
        ];
        assert_eq!(filter_ports(&entries, "3000").len(), 1);
        assert_eq!(filter_ports(&entries, "node").len(), 1);
        assert_eq!(filter_ports(&entries, "server").len(), 1);
        assert_eq!(filter_ports(&entries, "").len(), 2);
    }

    #[test]
    fn truncate_keeps_within_max() {
        assert_eq!(truncate_text("hello", 10), "hello");
        assert_eq!(truncate_text("hello world", 8), "hello w…");
        assert_eq!(truncate_text("hello", 0), "");
        assert_eq!(truncate_text("hello", 1), "");
    }
}