terminal-mcp 0.1.6

use crate::security::detect::bash::ast::{CurrentAst, get_command_name, language};
use crate::security::detect::utils::{
    node_extract_text, process_is_downloader, shell_is_unix,
};
use crate::security::detect::{Severity, ShellContext};
use std::sync::LazyLock;
use tree_sitter::{Node, Query, QueryCursor, StreamingIterator};

// 语言无关的路径/标志/命令名工具已提升至 detect/utils.rs，此处统一重导出，
// 保持 bash 各规则文件的 import 路径不变。
pub use crate::security::detect::utils::{
    cluster_has_flag, command_basename, is_block_device, normalize_target, path_has_marker,
};

/// 简单清理 Bash 字符串（去掉首尾的引号）
pub fn clean_bash_string(s: &str) -> String {
    let s = s.trim();
    if ((s.starts_with('"') && s.ends_with('"')) || (s.starts_with('\'') && s.ends_with('\'')))
        && s.len() >= 2
    {
        return s[1..s.len() - 1].to_string();
    }
    s.to_string()
}

/// 在节点树中 DFS 寻找下载器命令，并提取其中的 URL
pub fn extract_downloader_url(node: &Node, source: &[u8]) -> Option<String> {
    let mut stack = vec![*node];

    while let Some(n) = stack.pop() {
        if n.kind() == "command"
            && let Some(cmd_name) = get_command_name(&n, source)
            && process_is_downloader(cmd_name)
        {
            // 找到了下载器，遍历其同级节点寻找 URL 参数
            let mut arg_cursor = n.walk();
            for child in n.children(&mut arg_cursor) {
                if (child.kind() == "word" || child.kind() == "string")
                    && let Some(text) = node_extract_text(&child, source)
                {
                    // 清理两端的引号
                    let cleaned = text.trim_matches(|c| c == '\'' || c == '"');
                    if cleaned.starts_with("http://") || cleaned.starts_with("https://") {
                        return Some(cleaned.to_string());
                    }
                }
            }
        }

        let mut cursor = n.walk();
        for child in n.children(&mut cursor) {
            stack.push(child);
        }
    }
    None
}

pub fn detect_remote_execution(node: &Node, source: &[u8]) -> Option<String> {
    match node.kind() {
        "pipeline" => {
            // 场景 1: curl xxx | base64 -d | bash
            let mut cursor = node.walk();
            let mut commands = Vec::new();
            for child in node.children(&mut cursor) {
                // 仅收录管道顶层参与流转的命令或子 shell
                if child.kind() == "command" || child.kind() == "subshell" {
                    commands.push(child);
                }
            }

            if commands.len() >= 2 {
                let last_cmd = commands.last().unwrap();
                // 检查最终接收流的是否为 shell (Sink点)
                if let Some(last_name) = get_command_name(last_cmd, source)
                    && shell_is_unix(last_name)
                {
                    // 忽略中间管道，从上游寻找下载器 (Source点)
                    for cmd in commands.iter().take(commands.len() - 1) {
                        if let Some(url) = extract_downloader_url(cmd, source) {
                            return Some(url);
                        }
                    }
                }
            }
        }
        "command" => {
            // 场景 2: bash <(curl xxx) 或 bash -c "$(curl xxx)"
            if let Some(cmd_name) = get_command_name(node, source)
                && shell_is_unix(cmd_name)
            {
                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    // 检索参数中是否包含了进程替换或命令替换
                    if (child.kind() == "process_substitution"
                        || child.kind() == "command_substitution"
                        || child.kind() == "string")
                        && let Some(url) = extract_downloader_url(&child, source)
                    {
                        return Some(url);
                    }
                }
            }
        }
        _ => {}
    }
    None
}

// =============================================================================
// AST 规则公共工具（原 rules/helpers.rs 合并而来）
// =============================================================================

/// 覆盖 command 与 file_redirect 两类节点的公共查询。
/// file_redirect 是 command 在 redirected_statement 下的兄弟节点，
/// 不能只靠 `(command) @cmd` 拿到重定向目标，因此统一用两条 capture。
static COMMON_QUERY: LazyLock<Query> = LazyLock::new(|| {
    Query::new(&language(), "[(command) @cmd (file_redirect) @redirect]").expect("invalid query")
});

/// 允许出现在真实命令之前的"包装/前缀"命令。
pub const PREFIX_COMMANDS: &[&str] = &[
    "sudo", "env", "nohup", "exec", "command", "setsid", "nice", "ionice", "timeout", "stdbuf",
    "time",
];

/// 遍历当前所有反混淆块，对每个 command / file_redirect 节点运行 `analyze`，
/// 返回命中中最高严重度的一条。
pub async fn best_hit(
    ctx: &ShellContext,
    analyze: impl Fn(&Node, &[u8]) -> Option<(Severity, String)>,
) -> anyhow::Result<Option<(Severity, String)>> {
    let current = ctx
        .extensions
        .get::<CurrentAst>()
        .ok_or_else(|| anyhow::anyhow!("CurrentAst missing"))?;
    let blocks = current.blocks.read().await;

    let query = &*COMMON_QUERY;
    let mut best: Option<(Severity, String)> = None;

    for block in blocks.iter() {
        let source_bytes = block.source.as_bytes();
        let mut cursor = QueryCursor::new();
        let mut matches = cursor.matches(query, block.tree.root_node(), source_bytes);
        while let Some(m) = StreamingIterator::next(&mut matches) {
            for capture in m.captures {
                if let Some(hit) = analyze(&capture.node, source_bytes)
                    && best.as_ref().is_none_or(|(b, _)| hit.0 > *b)
                {
                    best = Some(hit);
                }
            }
        }
    }

    Ok(best)
}

/// 收集 `command` 节点在 command_name 之后的字符串类参数（不含重定向）。
pub fn collect_args<'a>(node: &Node, source: &'a [u8]) -> Vec<&'a str> {
    let mut args = Vec::new();
    let mut cursor = node.walk();
    let mut seen_name = false;

    for child in node.children(&mut cursor) {
        if child.kind() == "command_name" {
            seen_name = true;
            continue;
        }
        if !seen_name {
            continue;
        }
        if matches!(
            child.kind(),
            "word"
                | "string"
                | "raw_string"
                | "translated_string"
                | "ansi_c_string"
                | "concatenation"
                | "number"
                | "simple_expansion"
                | "expansion"
                | "command_substitution"
                | "process_substitution"
                | "arithmetic_expansion"
                | "brace_expression"
        ) && let Some(t) = node_extract_text(&child, source)
        {
            args.push(t);
        }
    }
    args
}

/// 解析 file_redirect 节点：返回 (操作符文本, 目标路径)。
/// 操作符为 `>` / `>>` / `>|` / `<` 等节点 kind。
pub fn redirect_target<'a>(node: &Node, source: &'a [u8]) -> Option<(&'a str, &'a str)> {
    let mut operator: Option<&str> = None;
    let mut dest: Option<&'a str> = None;

    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        match child.kind() {
            ">" | ">>" | ">|" | "<" | "<>" | "&>" | "&>>" | "<<<" => {
                operator = Some(child.kind());
            }
            "word" if operator.is_some() => {
                dest = node_extract_text(&child, source);
            }
            _ => {}
        }
    }

    match (operator, dest) {
        (Some(op), Some(dest)) => Some((op, dest)),
        _ => None,
    }
}

/// 若命令是前缀包装命令（sudo/env/timeout...），在参数中定位真实命令，
/// 返回 (真实命令名, 其后参数切片)。
/// 非前缀命令直接返回 (cmd_name, args)。
pub fn unwrap_command<'a>(
    cmd_name: &'a str,
    args: &'a [&'a str],
    known: &[&str],
) -> Option<(&'a str, &'a [&'a str])> {
    unwrap_command_where(cmd_name, args, |n| known.contains(&n))
}

/// `unwrap_command` 的谓词版本：`is_known` 自定义真实命令判定
/// （例如需要匹配 `mkfs.*` 等无法穷举的名字时）。
pub fn unwrap_command_where<'a>(
    cmd_name: &'a str,
    args: &'a [&'a str],
    is_known: impl Fn(&str) -> bool,
) -> Option<(&'a str, &'a [&'a str])> {
    if !PREFIX_COMMANDS.contains(&cmd_name) {
        return Some((cmd_name, args));
    }

    // 需要单独占用下一个参数作为"值"的选项（如 sudo -u root）
    let val_opts = [
        "-u",
        "--user",
        "-g",
        "--group",
        "-C",
        "--chdir",
        "-p",
        "--prompt",
        "-h",
        "--host",
        "-A",
        "--askpass",
        "-D",
        "--chroot",
        "-c",
        "--class",
        "-r",
        "--role",
        "-t",
        "--type",
        "-s",
        "--signal",
        "-n",
        "--adjustment",
        "-k",
        "--kill-after",
    ];

    let mut i = 0;
    while i < args.len() {
        let a = args[i];
        if val_opts.contains(&a) {
            i += 2;
            continue;
        }
        if a.starts_with('-') {
            i += 1;
            continue;
        }
        if is_known(a) {
            return Some((a, &args[i + 1..]));
        }
        i += 1;
    }
    None
}