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::powershell::deobf::DeobfMeta;
use crate::security::detect::utils::node_extract_text;

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

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

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

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

impl PsAstState {
    pub fn new(max_pending_bytes: usize) -> Self {
        let mut parser = Parser::new();
        parser
            .set_language(&language())
            .expect("Error loading Powershell 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, "powershell pending buffer overflow, forced flush");
            } else {
                buf.clear();
            }
            *frag = 0;
            return Vec::new();
        }

        // PowerShell 未闭合块 / 反引号续行 / 悬垂管道时不提交
        if !buf.ends_with('\n') || ends_with_backtick(&buf) || ends_with_dangling_ps(&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);

        vec![CommittedBlock::new_plain(source, tree, fragment_count)]
    }

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

/// 行末是否为反引号续行（`` ` `` 至少一次）。
fn ends_with_backtick(buf: &str) -> bool {
    let line = buf.strip_suffix('\n').unwrap_or(buf);
    line.trim_end().ends_with('`')
}

/// 是否以悬垂运算符结尾（未完成的管道 / 逻辑链 / 块 / 括号 / 赋值）。
fn ends_with_dangling_ps(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('(') || line.ends_with('=')
        || line.ends_with(',')
        || (line.ends_with('&') && !line.ends_with("&&"))
}

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()
}

/// 提取 PowerShell 命令名（field `command_name`，含 `&`/`.` 调用运算符形态）。
/// 返回 (原始文本, 是否为调用运算符形式)。
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 matches!(child.kind(), "command_name" | "command_name_expr" | "path_command_name") {
            return node_extract_text(&child, source);
        }
    }
    None
}

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

/// 从已提交的 PowerShell AST 块中提取变量赋值与环境变量（$env:NAME）写入。
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(),
            "(assignment_expression value: (_) @value) @assign",
        )
        .expect("Failed to create powershell assignment query")
    });

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

    while let Some(m) = StreamingIterator::next(&mut matches) {
        let Some(assign) = capture_by_name(&ASSIGN_QUERY, m, "assign") else {
            continue;
        };
        let Some(value) = capture_by_name(&ASSIGN_QUERY, m, "value") else {
            continue;
        };

        // 左侧必须是变量：variable / braced_variable；$env:NAME 视为环境变量
        let mut ac = assign.walk();
        let mut name: Option<String> = None;
        let mut is_env = false;
        for child in assign.children(&mut ac) {
            if matches!(child.kind(), "variable" | "braced_variable") {
                if let Some(t) = node_extract_text(&child, source) {
                    let t = t.trim_start_matches('$');
                    if let Some(rest) = t.strip_prefix("env:") {
                        is_env = true;
                        name = Some(rest.to_string());
                    } else if !t.contains(':') {
                        name = Some(t.to_string());
                    }
                }
            }
        }

        let (Some(name), Some(value_str)) = (name, node_extract_text(&value, source)) else {
            continue;
        };

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

    results
}