terminal-mcp 0.1.5

// bash/mod.rs

pub mod macros;
pub mod ast;
pub mod rules;
mod tests;
mod deobf;
mod utils;

use crate::security::detect::{Detector, Rule, ShellContext};
use anyhow::Result;
use async_trait::async_trait;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;

use self::ast::{BashAstState, CurrentAst, extract_env_vars};

pub struct BashDetector {
    ctx: Arc<ShellContext>,
    rules: Vec<Arc<dyn Rule>>,
}

impl BashDetector {
    pub fn new(mut ctx: ShellContext, max_pending_bytes: usize) -> Self {
        ctx.extensions.insert(BashAstState::new(max_pending_bytes));
        ctx.extensions.insert(CurrentAst::new());

        let ctx = Arc::new(ctx);
        Self {
            ctx,
            rules: rules::get_all_rules(),
        }
    }
}

#[async_trait]
impl Detector for BashDetector {
    fn context(&self) -> &Arc<ShellContext> { &self.ctx }
    fn rules(&self) -> &[Arc<dyn Rule>] { &self.rules }

    async fn on_detect(&self, data: &str) -> Result<()> {
        let state = self.ctx.extensions.get::<BashAstState>()
            .ok_or_else(|| anyhow::anyhow!("BashAstState missing"))?;

        // 1. 提交输入，如果未能形成完整块（如遇到续行符或未闭合引号），则返回空 vec
        let blocks = state.push_and_commit(data).await;

        if blocks.is_empty() {
            return Ok(());
        }

        // ==========================================
        // 2. 反混淆：词法还原 + 执行汇聚点递归展开
        //    （eval / bash -c / base64|bash / $(...) 等）
        //    处理完成后，block.source / block.tree 才是 Rule 应消费的干净内容。
        // ==========================================
        let budget = AtomicUsize::new(deobf::MAX_DEOBF_TOTAL_BYTES);
        let mut all_blocks = Vec::with_capacity(blocks.len());
        for block in blocks {
            let expanded = deobf::deobfuscate_block(block, &self.ctx, state, 0, &budget).await;
            all_blocks.extend(expanded);
        }


        // ==========================================
        // 3. 基于清洗后的 AST 解析并更新环境变量和动态变量
        // ==========================================
        for block in &all_blocks {
            let extracted_vars = extract_env_vars(&block.tree, block.source.as_bytes());

            for update in extracted_vars {
                // 不管是否带 export，只要有赋值就存入跨语言的动态变量池
                self.ctx.var.set(&update.name, update.value.clone()).await;

                // 若包含 export，则同步覆盖到真实的 Env 环境变量快照池
                if update.is_export {
                    self.ctx.env_set(&update.name, update.value).await;
                }
            }
        }

        // ==========================================
        // 4. 将清洗 + 展开后的全部块推入当前 AST 扩展供 Rules 并发评估使用
        // ==========================================
        let current = self.ctx.extensions.get::<CurrentAst>()
            .ok_or_else(|| anyhow::anyhow!("CurrentAst missing"))?;
        *current.blocks.write().await = all_blocks;

        Ok(())
    }
}