terminal-mcp 0.1.6

use std::sync::LazyLock;
use tokio::sync::{Mutex, RwLock};
use tree_sitter::{Language, Node, Parser, Query, QueryCursor, QueryMatch, StreamingIterator, Tree};
use crate::security::detect::bash::utils::clean_bash_string;
use crate::security::detect::bash::deobf::DeobfMeta;
use crate::security::detect::utils::{node_extract_text, process_is_downloader, shell_is_unix};

pub fn language() -> Language {
    tree_sitter_bash::LANGUAGE.into()
}

/// 一个已达到命令边界、语法合法的代码块。
///
/// `source` / `tree` 在经过 `deobf::deobfuscate_block` 处理后，
/// 始终代表"反混淆后的干净文本"，Rule 应只消费这两个字段。
/// 若需要审计原始输入，使用 `deobf.raw_source`。
#[derive(Debug)]
pub struct CommittedBlock {
    pub source: String,
    pub tree: Tree,
    pub is_heredoc_body: bool,
    /// 是否是由 Phase B（eval/bash -c/base64|bash 等）递归展开出来的子块
    pub is_decoded_payload: bool,
    pub fragment_count: usize,
    /// 反混淆过程的元数据（命中的手法、解码链、原始文本）
    pub deobf: DeobfMeta,
}

impl CommittedBlock {
    /// 构造一个尚未经过反混淆处理的"原始"块
    pub fn new_plain(source: String, tree: Tree, is_heredoc_body: bool, fragment_count: usize) -> Self {
        Self {
            source,
            tree,
            is_heredoc_body,
            is_decoded_payload: false,
            fragment_count,
            deobf: DeobfMeta::default(),
        }
    }
}

pub struct BashAstState {
    parser: Mutex<Parser>,
    pending: RwLock<String>,
    fragment_counter: RwLock<usize>,
    max_pending_bytes: usize,
}

impl BashAstState {
    pub fn new(max_pending_bytes: usize) -> Self {
        let mut parser = Parser::new();
        parser.set_language(&language()).expect("Error loading Bash grammar");
        Self {
            parser: Mutex::new(parser),
            pending: RwLock::new(String::new()),
            fragment_counter: RwLock::new(0),
            max_pending_bytes,
        }
    }

    pub async fn push_and_commit(&self, data: &str) -> Vec<CommittedBlock> {
        let mut buf = self.pending.write().await;
        buf.push_str(data);

        let mut frag = self.fragment_counter.write().await;
        *frag += 1;

        if buf.len() > self.max_pending_bytes {
            if let Some(idx) = buf.rfind('\n') {
                let forced: String = buf.drain(..=idx).collect();
                tracing::warn!(command = %forced, "bash pending buffer overflow, forced flush");
            } else {
                buf.clear();
            }
            *frag = 0;
            return Vec::new();
        }

        if !buf.ends_with('\n') || ends_with_line_continuation(&buf) || ends_with_dangling_operator(&buf) {
            return Vec::new();
        }

        let mut parser = self.parser.lock().await;
        let Some(tree) = parser.parse(buf.as_str(), None) else { return Vec::new() };
        if tree.root_node().has_error() {
            return Vec::new();
        }

        let source = std::mem::take(&mut *buf);
        let fragment_count = std::mem::replace(&mut *frag, 0);

        let mut blocks: Vec<CommittedBlock> = extract_heredoc_bodies(&tree, source.as_bytes())
            .into_iter()
            .filter_map(|body| {
                parser.parse(&body, None).map(|t| CommittedBlock::new_plain(body, t, true, 0))
            })
            .collect();

        blocks.push(CommittedBlock::new_plain(source, tree, false, fragment_count));
        blocks
    }

    /// 供 deobf 模块调用：用同一把解析器锁，重新解析反混淆/解码后的文本。
    /// 复用锁而不是新建 Parser，避免每次解码都重新加载 grammar。
    pub async fn reparse(&self, text: &str) -> Option<Tree> {
        let mut parser = self.parser.lock().await;
        parser.parse(text, None)
    }
}

fn ends_with_line_continuation(buf: &str) -> bool {
    let line = buf.strip_suffix('\n').unwrap_or(buf);
    line.chars().rev().take_while(|&c| c == '\\').count() % 2 == 1
}

fn ends_with_dangling_operator(buf: &str) -> bool {
    let line = buf.trim_end_matches('\n').trim_end();
    line.ends_with('|') || line.ends_with("&&") || line.ends_with("||")
        || (line.ends_with('&') && !line.ends_with("&&"))
}

fn extract_heredoc_bodies(tree: &Tree, source: &[u8]) -> Vec<String> {
    static QUERY: LazyLock<Query> =
        LazyLock::new(|| Query::new(&language(), "(heredoc_body) @body").expect("invalid query"));
    let mut cursor = QueryCursor::new();
    let mut matches = cursor.matches(&QUERY, tree.root_node(), source);
    let mut out = Vec::new();
    while let Some(m) = matches.next() {
        if let Some(n) = m.captures.first().map(|c| c.node)
            && let Ok(text) = n.utf8_text(source)
        {
            out.push(text.to_string());
        }
    }
    out
}

pub struct CurrentAst {
    pub blocks: RwLock<Vec<CommittedBlock>>,
}

impl CurrentAst {
    pub fn new() -> Self { Self { blocks: RwLock::new(Vec::new()) } }
}

/// 按 capture 名称取节点，避免位置错位导致的 bug
pub fn capture_by_name<'a>(query: &Query, m: &QueryMatch<'a, 'a>, name: &str) -> Option<Node<'a>> {
    let idx = query.capture_index_for_name(name)?;
    m.captures.iter().find(|c| c.index == idx).map(|c| c.node)
}

/// 一次性把所有具名 capture 收集成 map，方便多 capture 关联判断
pub fn captures_map<'a>(
    query: &Query,
    m: &QueryMatch<'a, 'a>,
) -> std::collections::HashMap<String, Node<'a>> {
    let names = query.capture_names();

    m.captures
        .iter()
        .map(|c| (names[c.index as usize].to_string(), c.node))
        .collect()
}

/// 提取命令节点的名字（跳过参数，处理嵌套）
pub fn get_command_name<'a>(node: &Node, source: &'a [u8]) -> Option<&'a str> {
    if node.kind() != "command" { return None; }
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if child.kind() == "command_name" {
            let mut wc = child.walk();
            for w in child.children(&mut wc) {
                if w.kind() == "word" || w.kind() == "string" {
                    return node_extract_text(&w, source);
                }
            }
            return node_extract_text(&child, source);
        }
    }
    None
}

pub struct EnvUpdate {
    pub name: String,
    pub value: String,
    pub is_export: bool,
}

/// 从已提交的 Bash AST 块中提取变量赋值和环境变量导出
pub fn extract_env_vars(tree: &Tree, source: &[u8]) -> Vec<EnvUpdate> {
    let mut results = Vec::new();

    static ASSIGN_QUERY: LazyLock<Query> = LazyLock::new(|| {
        Query::new(
            &language(),
            "(variable_assignment name: (variable_name) @name value: (_) @value)"
        ).expect("Failed to create bash variable assignment query")
    });

    let mut cursor = QueryCursor::new();
    let mut matches = cursor.matches(&ASSIGN_QUERY, tree.root_node(), source);

    while let Some(m) = matches.next() {
        let name_node = capture_by_name(&ASSIGN_QUERY, m, "name");
        let value_node = capture_by_name(&ASSIGN_QUERY, m, "value");

        if let (Some(n), Some(v)) = (name_node, value_node)
            && let (Some(name_str), Some(value_str)) = (
            node_extract_text(&n, source),
            node_extract_text(&v, source),
        )
        {
            let mut is_export = false;
            let mut current = n.parent();

            while let Some(parent) = current {
                if parent.kind() == "declaration_command" {
                    let mut pc = parent.walk();
                    for child in parent.children(&mut pc) {
                        if child.kind() == "command_name" || child.kind() == "word" {
                            if let Some(cmd_name) = node_extract_text(&child, source)
                                && cmd_name == "export"
                            {
                                is_export = true;
                            }
                            break;
                        }
                    }
                    break;
                }

                if parent.kind() == "command" || parent.kind() == "pipeline" {
                    break;
                }
                current = parent.parent();
            }

            results.push(EnvUpdate {
                name: name_str.to_string(),
                value: clean_bash_string(value_str),
                is_export,
            });
        }
    }

    results
}