terminal-mcp 0.1.6

// src/security/detect/powershell/mod.rs

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

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::{CurrentAst, PsAstState, extract_env_vars};

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

impl PowerShellDetector {
    pub fn new(mut ctx: ShellContext, max_pending_bytes: usize) -> Self {
        ctx.extensions.insert(PsAstState::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 PowerShellDetector {
    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::<PsAstState>()
            .ok_or_else(|| anyhow::anyhow!("PsAstState missing"))?;

        // 1. 提交输入；未形成完整语句（续行 / 未闭合块）则返回空 vec
        let blocks = state.push_and_commit(data).await;
        if blocks.is_empty() {
            return Ok(());
        }

        // 2. 反混淆：词法还原 + 执行汇聚点递归展开
        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 {
                self.ctx.var.set(&update.name, update.value.clone()).await;
                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(())
    }
}