terminal-mcp 0.1.6

// src/security/detect/powershell/deobf.rs

// 在达到语句边界（CommittedBlock 生成）之后调用，对块内容做两阶段清洗：
//
//   Phase A - 词法级还原（不产生新的可执行代码，只还原字面量）
//     - 反引号转义:      i`e`x Invoke-Expression        -> iex Invoke-Expression
//     - 字符串拼接:      'i'+'e'+'x'                     -> 'iex'
//     - 双引号内变量替换: "$cmd $args"                    -> 还原为变量当前值
//
//   Phase B - 执行汇聚点抽取（识别"字符串在运行时会被当作代码执行"的位置，
//             解码/展开后作为新的 CommittedBlock 递归投喂回 Phase A）
//     - iex / Invoke-Expression 嵌套解释器
//     - powershell/pwsh -EncodedCommand（UTF-16LE base64）
//     - [Convert]::FromBase64String(...) 解码
//     - $(...) 子表达式内部脚本
//
// 设计约束与 bash 模块一致：Phase A 中需要访问 ShellContext（async）的变量解析，
// 必须先把 Node 转换为纯 String/enum 中间表示，再异步求值，避免 Node 跨越 .await。

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::LazyLock;
use tree_sitter::{Node, Query, QueryCursor, StreamingIterator, Tree};

use crate::security::detect::ShellContext;
use crate::security::detect::powershell::ast::{
    PsAstState, capture_by_name, get_command_name, language,
};
use crate::security::detect::utils::node_extract_text;

/// 反混淆递归展开的最大深度
pub const MAX_DEOBF_DEPTH: usize = 512;
/// 单次 on_detect 调用中，Phase B 允许递归处理的总字节预算
pub const MAX_DEOBF_TOTAL_BYTES: usize = 512 * 1024;

// =============================================================================
// 元数据
// =============================================================================

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum ObfuscationTechnique {
    BacktickEscape,
    PsStringConcatenation,
    PsVariableSubstitution,
    PsEncodedCommand,
    PsBase64Convert,
    PsInvokeExpression,
    NestedPsInvocation,
    PsSubExpression,
    PsCallOperator,
    /// 存在无法静态求值的动态内容，仅打标不展开
    UnresolvedDynamic,
}

#[derive(Debug, Clone, Default)]
pub struct DeobfMeta {
    /// 本块自身（Phase A）命中的还原手法
    pub techniques: Vec<ObfuscationTechnique>,
    /// 若本块是从父块解码/展开而来，记录完整来源链
    pub decode_chain: Vec<ObfuscationTechnique>,
    /// 仅在发生了实质性重写时才保留原始文本，用于审计日志
    pub raw_source: Option<String>,
}

// =============================================================================
// Phase A - 中间表示
// =============================================================================

#[derive(PartialEq)]
enum WordKind {
    Name,
    Arg,
}

struct WordSpec {
    start_byte: usize,
    end_byte: usize,
    kind: WordKind,
}

// =============================================================================
// Phase A - 同步阶段：AST -> 词区间列表（不访问 ShellContext）
// =============================================================================

/// 反引号转义还原：`` `x `` -> x，`` `n `` -> 换行（PowerShell 转义序列）。
fn unescape_ps(raw: &str) -> String {
    let mut out = String::with_capacity(raw.len());
    let mut chars = raw.chars();
    while let Some(c) = chars.next() {
        if c == '`' {
            match chars.next() {
                Some('n') => out.push('\n'),
                Some('r') => out.push('\r'),
                Some('t') => out.push('\t'),
                Some('a') => out.push('\x07'),
                Some('b') => out.push('\x08'),
                Some('f') => out.push('\x0c'),
                Some('v') => out.push('\x0b'),
                Some('0') => out.push('\0'),
                Some('`') => out.push('`'),
                Some(other) => out.push(other),
                None => out.push('`'),
            }
        } else {
            out.push(c);
        }
    }
    out
}

/// 判断一个字符串是否包含 `$name` / `${name}` 形式的变量引用（用于替换判定）。
fn contains_variable_ref(text: &str) -> bool {
    text.contains('$')
}

/// 双引号字符串内容中的变量替换（保留 `$()` 原始字节）。
/// 返回 (替换后文本, 是否发生了替换)。
async fn substitute_vars_inner(raw: &str, ctx: &ShellContext) -> (String, bool) {
    let mut out = String::with_capacity(raw.len());
    let mut chars = raw.chars().peekable();
    let mut changed = false;

    while let Some(c) = chars.next() {
        if c == '$' {
            let next = chars.peek().copied();
            match next {
                Some('(') => {
                    // $(...) 子表达式：原样拷贝，交由 Phase B 单独抽取
                    out.push_str("$(");
                    chars.next();
                    continue;
                }
                Some('{') => {
                    // ${name}
                    let mut body = String::new();
                    for ch in chars.by_ref() {
                        if ch == '}' {
                            break;
                        }
                        body.push(ch);
                    }
                    match lookup_variable(&body, ctx).await {
                        Some(v) => {
                            out.push_str(&v);
                            changed = true;
                        }
                        None => {
                            out.push_str("${");
                            out.push_str(&body);
                            out.push('}');
                        }
                    }
                }
                Some(c) if c.is_ascii_alphabetic() || c == '_' => {
                    let mut name = String::new();
                    name.push(c);
                    chars.next();
                    while let Some(ch) = chars.peek() {
                        if ch.is_ascii_alphanumeric() || *ch == '_' {
                            name.push(*ch);
                            chars.next();
                        } else {
                            break;
                        }
                    }
                    match lookup_variable(&name, ctx).await {
                        Some(v) => {
                            out.push_str(&v);
                            changed = true;
                        }
                        None => {
                            out.push('$');
                            out.push_str(&name);
                        }
                    }
                }
                _ => {
                    out.push('$');
                }
            }
        } else {
            out.push(c);
        }
    }

    (out, changed)
}

/// 变量解析：优先变量池，其次 env 池（支持 $env:NAME 作用域写法）。
async fn lookup_variable(name: &str, ctx: &ShellContext) -> Option<String> {
    let (scope, key) = match name.split_once(':') {
        Some(("env", k)) => ("env", k.to_string()),
        Some((_, k)) => ("var", k.to_string()),
        None => ("var", name.to_string()),
    };

    if scope == "env" {
        if let Some(v) = ctx.env_get(&key).await {
            return Some(v);
        }
        return None;
    }

    if let Some(v) = ctx.var.get(name).await
        && let Some(s) = v.as_str()
    {
        return Some(s.to_string());
    }
    // var 池未命中时，回退到同名 env
    ctx.env_get(name).await
}

/// 收集 command 节点中需要参与词法重写的区间（跳过 redirection / 赋值等）。
fn collect_word_specs(tree: &Tree, source: &[u8]) -> Vec<(WordSpec, String)> {
    static QUERY: LazyLock<Query> =
        LazyLock::new(|| Query::new(&language(), "(command) @cmd").expect("invalid query"));

    let mut cursor = QueryCursor::new();
    let mut specs: Vec<(WordSpec, String)> = Vec::new();

    let mut matches = cursor.matches(&QUERY, tree.root_node(), source);
    while let Some(m) = StreamingIterator::next(&mut matches) {
        let Some(cmd) = capture_by_name(&QUERY, m, "cmd") else {
            continue;
        };
        if has_ancestor_kind(&cmd, "sub_expression") {
            continue;
        }

        let mut cc = cmd.walk();
        let mut seen_elements = false;
        for child in cmd.children(&mut cc) {
            if child.kind() == "command_name"
                || child.kind() == "command_name_expr"
                || child.kind() == "path_command_name"
            {
                if let Some(t) = node_extract_text(&child, source) {
                    specs.push((
                        WordSpec {
                            start_byte: child.start_byte(),
                            end_byte: child.end_byte(),
                            kind: WordKind::Name,
                        },
                        t.to_string(),
                    ));
                }
            } else if child.kind() == "command_elements" {
                seen_elements = true;
                let mut ec = child.walk();
                for el in child.children(&mut ec) {
                    match el.kind() {
                        "generic_token" | "command_parameter" | "string_literal"
                        | "expandable_string_literal" | "verbatim_string_characters"
                        | "verbatim_here_string_characters" | "expandable_here_string_literal" => {
                            if let Some(t) = node_extract_text(&el, source) {
                                specs.push((
                                    WordSpec {
                                        start_byte: el.start_byte(),
                                        end_byte: el.end_byte(),
                                        kind: WordKind::Arg,
                                    },
                                    t.to_string(),
                                ));
                            }
                        }
                        _ => {
                            // argument_list / parenthesized_expression / redirection 等
                            // 由 Phase B 或规则单独处理，不在此处重写
                        }
                    }
                }
            } else if !seen_elements && child.kind() != "command_invokation_operator" {
                // 不做处理
            }
        }
    }
    specs
}

fn has_ancestor_kind(node: &Node, kind: &str) -> bool {
    let mut cur = node.parent();
    while let Some(p) = cur {
        if p.kind() == kind {
            return true;
        }
        cur = p.parent();
    }
    false
}

/// 将字符串拼接表达式折叠为单字符串字面量：`'i'+'e'+'x'` -> `'iex'`。
/// 仅当所有子节点都是字符串字面量时折叠；命中返回 (节点区间, 折叠结果)。
fn fold_additive_strings(tree: &Tree, source: &[u8]) -> Vec<(usize, usize, String)> {
    static QUERY: LazyLock<Query> = LazyLock::new(|| {
        Query::new(
            &language(),
            "[(additive_expression) @add (additive_argument_expression) @add]",
        )
        .expect("invalid query")
    });

    let mut cursor = QueryCursor::new();
    let mut out = Vec::new();
    let mut matches = cursor.matches(&QUERY, tree.root_node(), source);
    while let Some(m) = StreamingIterator::next(&mut matches) {
        let Some(node) = capture_by_name(&QUERY, m, "add") else {
            continue;
        };
        if has_ancestor_kind(&node, "sub_expression") {
            continue;
        }
        if let Some(folded) = fold_one(&node, source) {
            out.push((node.start_byte(), node.end_byte(), folded));
        }
    }
    out
}

/// 递归折叠：node 的所有子节点必须是字符串字面量，返回拼接后的内容。
fn fold_one(node: &Node, source: &[u8]) -> Option<String> {
    let mut parts = Vec::new();
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        match child.kind() {
            "string_literal" => {
                let t = node_extract_text(&child, source)?;
                parts.push(clean_ps_string(t));
            }
            "additive_expression" | "additive_argument_expression" => {
                let t = fold_one(&child, source)?;
                parts.push(t);
            }
            _ => return None,
        }
    }
    if parts.is_empty() {
        None
    } else {
        Some(parts.concat())
    }
}

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

/// 对树中所有顶层命令做规范化重写，返回 (新文本, 命中的手法列表)。
/// 若没有任何变化，返回的文本与传入的 source 完全一致。
async fn build_normalized_source(
    tree: &Tree,
    source: &[u8],
    ctx: &ShellContext,
) -> (String, Vec<ObfuscationTechnique>) {
    let words = collect_word_specs(tree, source);
    let concats = fold_additive_strings(tree, source);

    let mut edits: Vec<(usize, usize, String)> = Vec::new();
    let mut techs: Vec<ObfuscationTechnique> = Vec::new();

    for (spec, text) in &words {
        let mut new_text = unescape_ps(text);
        if spec.kind == WordKind::Name {
            if new_text != *text {
                techs.push(ObfuscationTechnique::BacktickEscape);
            }
        } else if new_text != *text {
            techs.push(ObfuscationTechnique::BacktickEscape);
        }

        if matches!(spec.kind, WordKind::Arg) && contains_variable_ref(text) {
            let inner = trim_quote_pair(&new_text);
            let (subbed, changed) = substitute_vars_inner(&inner, ctx).await;
            if changed {
                // 保留原引号风格
                let re = match new_text.chars().next() {
                    Some('"') => format!("\"{subbed}\""),
                    _ => subbed,
                };
                new_text = re;
                techs.push(ObfuscationTechnique::PsVariableSubstitution);
            }
        }

        if new_text != *text {
            edits.push((spec.start_byte, spec.end_byte, new_text));
        }
    }

    for (start, end, folded) in concats {
        edits.push((start, end, folded));
        techs.push(ObfuscationTechnique::PsStringConcatenation);
    }

    if edits.is_empty() {
        return (String::from_utf8_lossy(source).into_owned(), Vec::new());
    }

    edits.sort_by_key(|e| e.0);
    let mut out = String::with_capacity(source.len());
    let mut last = 0usize;
    for (start, end, repl) in edits {
        if start < last {
            continue;
        }
        out.push_str(std::str::from_utf8(&source[last..start]).unwrap_or(""));
        out.push_str(&repl);
        last = end;
    }
    out.push_str(std::str::from_utf8(&source[last..]).unwrap_or(""));

    (out, techs)
}

fn trim_quote_pair(s: &str) -> &str {
    let s = s.trim();
    if s.len() >= 2
        && ((s.starts_with('"') && s.ends_with('"'))
            || (s.starts_with('\'') && s.ends_with('\'')))
    {
        &s[1..s.len() - 1]
    } else {
        s
    }
}

// =============================================================================
// Phase B - 执行汇聚点抽取（全同步）
// =============================================================================

fn base64_decode_bytes(input: &str) -> Option<Vec<u8>> {
    fn val(c: u8) -> Option<u8> {
        match c {
            b'A'..=b'Z' => Some(c - b'A'),
            b'a'..=b'z' => Some(c - b'a' + 26),
            b'0'..=b'9' => Some(c - b'0' + 52),
            b'+' => Some(62),
            b'/' => Some(63),
            _ => None,
        }
    }
    let bytes: Vec<u8> = input.bytes().filter(|&b| b != b'=').collect();
    if bytes.is_empty() {
        return None;
    }

    let mut out = Vec::with_capacity(bytes.len() * 3 / 4 + 3);
    let mut buf: u32 = 0;
    let mut bits: u32 = 0;
    for b in bytes {
        let v = val(b)?;
        buf = (buf << 6) | v as u32;
        bits += 6;
        if bits >= 8 {
            bits -= 8;
            out.push(((buf >> bits) & 0xFF) as u8);
        }
    }
    Some(out)
}

/// PowerShell base64 解码：优先 UTF-16LE（官方 -EncodedCommand 编码），回退 UTF-8。
fn decode_ps_base64(s: &str) -> Option<String> {
    let clean: String = s.chars().filter(|c| !c.is_whitespace()).collect();
    let bytes = base64_decode_bytes(&clean)?;
    if bytes.len() >= 2 {
        if let Some(text) = decode_utf16le(&bytes) {
            if !text.trim().is_empty() {
                return Some(text);
            }
        }
    }
    String::from_utf8(bytes).ok()
}

fn decode_utf16le(bytes: &[u8]) -> Option<String> {
    let units: Vec<u16> = bytes
        .chunks_exact(2)
        .map(|c| u16::from_le_bytes([c[0], c[1]]))
        .collect();
    // 常见编码器会在开头带 BOM（0xFEFF），剥掉
    let units = if units.first() == Some(&0xFEFF) {
        &units[1..]
    } else {
        &units[..]
    };
    String::from_utf16(units).ok()
}

/// 提取 command 节点 command_name 之后的第一个"可执行"参数节点。
fn extract_sink_argument<'a>(cmd: &Node<'a>) -> Option<Node<'a>> {
    let mut cc = cmd.walk();
    for child in cmd.children(&mut cc) {
        if child.kind() == "command_elements" {
            let mut ec = child.walk();
            for el in child.children(&mut ec) {
                match el.kind() {
                    "command_parameter" | "command_argument_sep" | "redirection" => continue,
                    _ => return Some(el),
                }
            }
        }
    }
    None
}

/// 收集 command_elements 的直接子节点（含原始文本），便于做参数级判断。
fn collect_elements<'a>(cmd: &Node<'a>, source: &[u8]) -> Vec<(String, Node<'a>)> {
    let mut out = Vec::new();
    let mut cc = cmd.walk();
    for child in cmd.children(&mut cc) {
        if child.kind() == "command_elements" {
            let mut ec = child.walk();
            for el in child.children(&mut ec) {
                let text = node_extract_text(&el, source).unwrap_or("").to_string();
                out.push((text, el));
            }
        }
    }
    out
}

/// 在 command_elements 中定位 `-EncodedCommand` / `-enc` 之后的值节点。
fn extract_encoded_command_arg<'a>(
    elements: &[(String, Node<'a>)],
) -> Option<Node<'a>> {
    for (i, (text, _)) in elements.iter().enumerate() {
        let lower = text.to_lowercase();
        if lower == "-enc" || lower == "-encodedcommand" || lower == "-e" {
            return elements.get(i + 1).map(|(_, n)| *n);
        }
    }
    None
}

/// 遍历树，收集执行汇聚点 payload。
fn walk_for_sinks(tree: &Tree, source: &[u8], out: &mut Vec<(String, ObfuscationTechnique)>) {
    let mut stack = vec![tree.root_node()];
    while let Some(n) = stack.pop() {
        match n.kind() {
            "command" => {
                if let Some(name) = get_command_name(&n, source) {
                    let lower = name.to_lowercase();
                    let elements = collect_elements(&n, source);
                    match lower.as_str() {
                        "iex" | "invoke-expression" => {
                            if let Some(arg) = extract_sink_argument(&n)
                                && let Some(text) = node_extract_text(&arg, source)
                                && !text.trim().is_empty()
                            {
                                out.push((
                                    text.to_string(),
                                    ObfuscationTechnique::PsInvokeExpression,
                                ));
                            }
                        }
                        "powershell" | "pwsh" | "powershell.exe" | "pwsh.exe" => {
                            // 优先识别 -EncodedCommand；否则作为嵌套 PS 调用抽取首参
                            if let Some(arg) = extract_encoded_command_arg(&elements)
                                && let Some(text) = node_extract_text(&arg, source)
                                && let Some(decoded) = decode_ps_base64(text)
                                && !decoded.trim().is_empty()
                            {
                                out.push((decoded, ObfuscationTechnique::PsEncodedCommand));
                            } else if let Some(arg) = extract_sink_argument(&n)
                                && let Some(text) = node_extract_text(&arg, source)
                                && !text.trim().is_empty()
                            {
                                out.push((text.to_string(), ObfuscationTechnique::NestedPsInvocation));
                            }
                        }
                        _ => {}
                    }
                }
            }
            "sub_expression" => {
                if let Some(text) = node_extract_text(&n, source) {
                    let inner = text
                        .trim_start_matches("$(")
                        .trim_end_matches(')')
                        .to_string();
                    if !inner.trim().is_empty() {
                        out.push((inner, ObfuscationTechnique::PsSubExpression));
                    }
                }
            }
            "invokation_expression" => {
                // [Convert]::FromBase64String('...')
                if let Some(decoded) = extract_base64_convert(&n, source) {
                    out.push((decoded, ObfuscationTechnique::PsBase64Convert));
                }
            }
            _ => {}
        }

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

/// 从 invokation_expression 中识别 [Convert]::FromBase64String('...') 并解码。
fn extract_base64_convert(node: &Node, source: &[u8]) -> Option<String> {
    let text = node_extract_text(node, source)?;
    // 文本级识别，避免深挖 member_access 结构
    let lower = text.to_lowercase();
    if !lower.contains("frombase64string") {
        return None;
    }
    // 提取第一个字符串字面量参数
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if child.kind() == "argument_list" {
            let mut ac = child.walk();
            for arg in child.children(&mut ac) {
                if matches!(arg.kind(), "string_literal" | "expandable_string_literal") {
                    if let Some(t) = node_extract_text(&arg, source) {
                        let inner = clean_ps_string(t);
                        if let Some(decoded) = decode_ps_base64(&inner) {
                            return Some(decoded);
                        }
                    }
                }
            }
        }
    }
    None
}

// =============================================================================
// 主函数
// =============================================================================

fn consume_budget(budget: &AtomicUsize, amount: usize) -> bool {
    loop {
        let cur = budget.load(Ordering::Relaxed);
        if amount > cur {
            return false;
        }
        if budget
            .compare_exchange_weak(cur, cur - amount, Ordering::Relaxed, Ordering::Relaxed)
            .is_ok()
        {
            return true;
        }
    }
}

/// 对单个 CommittedBlock 做完整的反混淆处理。
/// 返回值中第一个元素恒为"清洗后的原块"，之后是所有递归展开出的子块。
pub async fn deobfuscate_block(
    mut block: crate::security::detect::powershell::ast::CommittedBlock,
    ctx: &ShellContext,
    ast_state: &PsAstState,
    depth: usize,
    budget: &AtomicUsize,
) -> Vec<crate::security::detect::powershell::ast::CommittedBlock> {
    use crate::security::detect::powershell::ast::CommittedBlock;

    if depth > MAX_DEOBF_DEPTH {
        tracing::warn!(
            target: "security::deobf",
            depth,
            "max deobfuscation depth exceeded, stop expanding further"
        );
        return vec![block];
    }

    // ---------------- Phase A ----------------
    let (normalized_src, techs) =
        build_normalized_source(&block.tree, block.source.as_bytes(), ctx).await;

    if normalized_src != block.source {
        match ast_state.reparse(&normalized_src).await {
            Some(new_tree) if !new_tree.root_node().has_error() => {
                block
                    .deobf
                    .raw_source
                    .get_or_insert_with(|| block.source.clone());
                block.source = normalized_src;
                block.tree = new_tree;
            }
            _ => {
                tracing::debug!(
                    target: "security::deobf",
                    "normalized source failed to reparse cleanly, fallback to original"
                );
            }
        }
    }
    block.deobf.techniques.extend(techs);

    // ---------------- Phase B ----------------
    let mut payloads: Vec<(String, ObfuscationTechnique)> = Vec::new();
    walk_for_sinks(&block.tree, block.source.as_bytes(), &mut payloads);

    let parent_chain = block.deobf.decode_chain.clone();
    let mut result = vec![block];

    for (payload_text, tech) in payloads {
        if !consume_budget(budget, payload_text.len()) {
            tracing::warn!(
                target: "security::deobf",
                payload_len = payload_text.len(),
                "deobf byte budget exceeded, dropping remaining payload"
            );
            continue;
        }

        let Some(tree) = ast_state.reparse(&payload_text).await else {
            continue;
        };
        if tree.root_node().has_error() {
            continue;
        }

        let mut chain = parent_chain.clone();
        chain.push(tech);

        let child_block = CommittedBlock {
            source: payload_text,
            tree,
            is_decoded_payload: true,
            fragment_count: 0,
            deobf: DeobfMeta {
                techniques: Vec::new(),
                decode_chain: chain,
                raw_source: None,
            },
        };

        let expanded = Box::pin(deobfuscate_block(child_block, ctx, ast_state, depth + 1, budget))
            .await;
        result.extend(expanded);
    }

    result
}