terminal-mcp 0.1.6

// src/security/detect/mod.rs
// 该模块尚未完善和投入使用
#![allow(unused)]

pub mod bash;
pub mod cmd;
pub mod node;
pub mod powershell;
pub mod python;
mod utils;

use crate::security::audit::truncate_for_log;
use anyhow::Result;
use async_trait::async_trait;
use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::collections::VecDeque;
use std::env;
use std::sync::{Arc, LazyLock};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tokio::task::JoinSet;
use tokio::time::{timeout, Duration};

/// 规则并发评估的滑动窗口最大容量
const MAX_CONCURRENT_RULES: usize = 30;

static SHELL_DETECTION_LEVEL: LazyLock<Option<Severity>> = LazyLock::new(|| {
    match env::var("SHELL_DETECTION_LEVEL")
        .unwrap_or_else(|_| "medium".to_string())
        .to_lowercase()
        .as_str()
    {
        "critical" => Some(Severity::Critical),
        "high" => Some(Severity::High),
        "medium" => Some(Severity::Medium),
        "low" => Some(Severity::Low),
        "none" => None,
        _ => Some(Severity::Medium),
    }
});

static ON_DETECT_TIMEOUT: LazyLock<Duration> = LazyLock::new(|| {
    env::var("ON_DETECT_TIMEOUT_MS")
        .ok()
        .and_then(|v| v.parse().ok())
        .map(Duration::from_millis)
        .unwrap_or_else(|| Duration::from_millis(3000))
});

static RULE_TIMEOUT: LazyLock<Duration> = LazyLock::new(|| {
    env::var("RULE_TIMEOUT_MS")
        .ok()
        .and_then(|v| v.parse().ok())
        .map(Duration::from_millis)
        .unwrap_or_else(|| Duration::from_millis(3000))
});

#[derive(Default)]
pub struct Extensions {
    map: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
}

impl Extensions {
    pub fn new() -> Self {
        Self {
            map: HashMap::new(),
        }
    }
    /// 插入任意类型的数据
    pub fn insert<T: Send + Sync + 'static>(&mut self, val: T) {
        self.map.insert(TypeId::of::<T>(), Box::new(val));
    }
    /// 获取任意类型的数据的引用
    pub fn get<T: 'static>(&self) -> Option<&T> {
        self.map
            .get(&TypeId::of::<T>())
            .and_then(|boxed| boxed.as_ref().downcast_ref::<T>())
    }

    /// 获取任意类型数据的可变引用
    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
        self.map
            .get_mut(&TypeId::of::<T>())
            .and_then(|boxed| boxed.as_mut().downcast_mut::<T>())
    }
}

impl std::fmt::Debug for Extensions {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Extensions")
            .field("count", &self.map.len())
            .finish()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Severity {
    Low,
    Medium,
    High,
    Critical,
}

#[derive(Debug, Clone, PartialEq)]
pub struct RuleMetadata {
    pub name: String,
    pub description: String,
    /// 规则的静态默认危险等级（用于UI展示和兜底）
    pub default_severity: Severity,
}

/// 封装了单次规则命中的完整上下文
#[derive(Debug, Clone, PartialEq)]
pub struct ThreatHit {
    pub rule_meta: RuleMetadata,
    pub evidence: Option<String>,
    /// 最终确定的危险等级（结合了规则动态返回的覆盖值与默认值）
    pub final_severity: Severity,
}

#[derive(Debug, Clone, PartialEq)]
pub enum DetectResult {
    Safe,
    ThreatDetected(Vec<ThreatHit>),
    Unknown,
}

impl DetectResult {
    /// 聚合出本次命中中的最高危级别，方便调用方直接决策
    pub fn max_severity(&self) -> Option<Severity> {
        match self {
            DetectResult::ThreatDetected(hits) => {
                hits.iter().map(|hit| hit.final_severity).max()
            }
            _ => None,
        }
    }
}

/// 跨语言通用的动态变量值
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum VarValue {
    Null,
    Bool(bool),
    Int(i64),
    Float(f64),
    Str(String),
    List(Vec<VarValue>),
    Map(HashMap<String, VarValue>),
}

impl VarValue {
    pub fn as_str(&self) -> Option<&str> {
        match self {
            VarValue::Str(s) => Some(s),
            _ => None,
        }
    }
    pub fn as_i64(&self) -> Option<i64> {
        match self {
            VarValue::Int(i) => Some(*i),
            _ => None,
        }
    }
    pub fn is_truthy(&self) -> bool {
        match self {
            VarValue::Null => false,
            VarValue::Bool(b) => *b,
            VarValue::Int(i) => *i != 0,
            VarValue::Float(f) => *f != 0.0,
            VarValue::Str(s) => !s.is_empty(),
            VarValue::List(l) => !l.is_empty(),
            VarValue::Map(m) => !m.is_empty(),
        }
    }
}
impl From<String> for VarValue {
    fn from(s: String) -> Self { VarValue::Str(s) }
}
impl From<&str> for VarValue {
    fn from(s: &str) -> Self { VarValue::Str(s.to_string()) }
}
impl From<i64> for VarValue {
    fn from(i: i64) -> Self { VarValue::Int(i) }
}
impl From<bool> for VarValue {
    fn from(b: bool) -> Self { VarValue::Bool(b) }
}
#[derive(Debug, Default)]
pub struct VariableStore {
    inner: RwLock<HashMap<String, VarValue>>,
}

impl VariableStore {
    pub fn new() -> Self {
        Self { inner: RwLock::new(HashMap::new()) }
    }

    pub async fn set(
        &self,
        key: impl Into<String>,
        value: impl Into<VarValue>,
    ) {
        self.inner.write().await.insert(
            key.into(),
            value.into(),
        );
    }

    pub async fn get(&self, key: &str) -> Option<VarValue> {
        self.inner.read().await.get(key).cloned()
    }

    pub async fn remove(&self, key: &str) -> Option<VarValue> {
        self.inner.write().await.remove(key)
    }

    pub async fn snapshot(&self) -> HashMap<String, VarValue> {
        self.inner.read().await.clone()
    }

    /// 零拷贝闭包访问，避免整体克隆大 Map
    pub async fn with<R>(&self, f: impl FnOnce(&HashMap<String, VarValue>) -> R) -> R {
        let guard = self.inner.read().await;
        f(&guard)
    }

}

#[derive(Debug)]
pub struct ShellContext {
    pub shell_path: String,
    env: RwLock<HashMap<String, String>>,
    pub var: VariableStore,
    history: RwLock<VecDeque<String>>,         // 运行期追加，内部互斥
    max_history_size: usize,
    pub extensions: Extensions,
}

impl ShellContext {
    pub fn new(
        shell_path: impl Into<String>,
        env: HashMap<String, String>,
        max_history_size: usize,
    ) -> Self {

        Self {
            shell_path: shell_path.into(),
            env: RwLock::new(env),
            var: Default::default(),
            history: RwLock::new(VecDeque::with_capacity(max_history_size)),
            max_history_size,
            extensions: Extensions::new(),
        }
    }

    // ==========================================
    // 环境变量相关操作方法 (Environment operations)
    // ==========================================

    /// 获取所有环境变量的快照
    pub async fn env_snapshot(&self) -> HashMap<String, String> {
        self.env.read().await.clone()
    }

    /// 获取单个环境变量的值
    pub async fn env_get(&self, key: &str) -> Option<String> {
        self.env.read().await.get(key).cloned()
    }

    /// 动态设置或更新环境变量
    pub async fn env_set(&self, key: impl Into<String>, value: impl Into<String>) {
        self.env.write().await.insert(key.into(), value.into());
    }

    /// 动态移除环境变量，返回被移除的值（若存在）
    pub async fn env_remove(&self, key: &str) -> Option<String> {
        self.env.write().await.remove(key)
    }

    /// 提供只读闭包方式访问环境变量，避免克隆整个 HashMap
    pub async fn env_with<R>(&self, f: impl FnOnce(&HashMap<String, String>) -> R) -> R {
        let guard = self.env.read().await;
        f(&guard)
    }

    // ==========================================
    // 历史命令相关操作方法 (History operations)
    // ==========================================

    pub async fn push_history(&self, cmd: impl Into<String>) {
        if self.max_history_size == 0 {
            return;
        }
        let mut h = self.history.write().await;
        if h.len() >= self.max_history_size {
            h.pop_front();
        }
        h.push_back(cmd.into());
    }

    pub async fn history_snapshot(&self) -> Vec<String> {
        self.history.read().await.iter().cloned().collect()
    }

    pub async fn history_with<R>(&self, f: impl FnOnce(&VecDeque<String>) -> R) -> R {
        let guard = self.history.read().await;
        f(&guard)
    }

    pub async fn history_recent(&self, n: usize) -> Vec<String> {
        self.history_with(|h| h.iter().rev().take(n).rev().cloned().collect())
            .await
    }
}

pub enum EvaluateResult {
    /// 命中。参数1为证据(可选)，参数2为动态评估的危险等级(可选)
    /// 如果返回动态危险等级，将覆盖 RuleMetadata 中的 default_severity
    Hit(Option<String>, Option<Severity>),
    Miss,
}

impl EvaluateResult {
    /// 辅助方法：快速构建不覆盖默认等级的命中
    pub fn hit(evidence: impl Into<String>) -> Self {
        EvaluateResult::Hit(Some(evidence.into()), None)
    }

    /// 辅助方法：快速构建带有动态危险等级的命中
    pub fn hit_with_severity(evidence: impl Into<String>, severity: Severity) -> Self {
        EvaluateResult::Hit(Some(evidence.into()), Some(severity))
    }
}

#[async_trait]
pub trait Rule: Send + Sync {
    fn meta(&self) -> &RuleMetadata;
    async fn evaluate(&self, data: &str, ctx: &ShellContext) -> Result<EvaluateResult>;
}

#[async_trait]
pub trait Detector: Send + Sync {
    fn context(&self) -> &Arc<ShellContext>;
    fn rules(&self) -> &[Arc<dyn Rule>];

    async fn on_detect(&self, data: &str) -> Result<()>; // 用于更新上下文等等

    // 每个Detector中的detect一定是串行调用的
    async fn detect(&self, mut data: String, stop_on_first_hit: bool, append_enter: bool) -> DetectResult {
        if SHELL_DETECTION_LEVEL.is_none() {
            return DetectResult::Unknown;
        }
        let threshold_severity = SHELL_DETECTION_LEVEL.unwrap();

        let ctx = self.context();
        let mut hits = Vec::new();
        let mut evaluated = 0usize;

        let on_detect_timeout = *ON_DETECT_TIMEOUT;
        let rule_timeout = *RULE_TIMEOUT;

        if append_enter {
            data.push('\n');
        }
        let data_str = data.as_str();

        match timeout(on_detect_timeout, self.on_detect(data_str)).await {
            Ok(Ok(_)) => {}
            Ok(Err(e)) => {
                tracing::error!(
                    target: "security::on_detect",
                    shell = %ctx.shell_path,
                    input_len = data_str.len(),
                    input_preview = %truncate_for_log(data_str, 200),
                    error = %e,
                    "on_detect failed, skipped rule evaluation"
                );
                return DetectResult::Unknown;
            }
            Err(_) => {
                tracing::error!(
                    target: "security::on_detect",
                    shell = %ctx.shell_path,
                    input_len = data_str.len(),
                    input_preview = %truncate_for_log(data_str, 200),
                    "on_detect timed out after {:?}", on_detect_timeout
                );
                return DetectResult::Unknown;
            }
        }

        let mut set = JoinSet::new();
        let mut rules_iter = self.rules().iter().cloned();

        let data_arc: Arc<str> = Arc::from(data_str);

        // 规则内不进行上下文更新，通过on_detect时就准备好一切
        let spawn_rule = |set: &mut JoinSet<_>, rule: Arc<dyn Rule>, data_arc: Arc<str>, ctx: Arc<ShellContext>| {
            set.spawn(async move {
                let res = timeout(rule_timeout, rule.evaluate(&data_arc, &ctx)).await;
                (rule, res)
            });
        };

        // 初始填充滑动窗口，控制最大并发量
        for _ in 0..MAX_CONCURRENT_RULES {
            if let Some(rule) = rules_iter.next() {
                spawn_rule(&mut set, rule, Arc::clone(&data_arc), Arc::clone(ctx));
            }
        }

        while let Some(res) = set.join_next().await {
            match res {
                Ok((rule, timeout_res)) => {
                    match timeout_res {
                        Ok(Ok(EvaluateResult::Hit(evidence, override_severity))) => {
                            let meta = rule.meta().clone();

                            // 动态计算：如果有覆盖值则使用覆盖值，否则回退到静态默认值
                            let final_severity = override_severity.unwrap_or(meta.default_severity);
                            let is_over_threshold = final_severity >= threshold_severity;

                            hits.push(ThreatHit {
                                rule_meta: meta,
                                evidence,
                                final_severity,
                            });
                            evaluated += 1;

                            if stop_on_first_hit && is_over_threshold {
                                set.abort_all();
                                break;
                            }
                        }
                        Ok(Ok(EvaluateResult::Miss)) => {
                            evaluated += 1;
                        }
                        Ok(Err(err)) => {
                            tracing::error!(
                                target: "security::detect",
                                rule_name = %rule.meta().name,
                                rule_default_severity = ?rule.meta().default_severity,
                                shell = %ctx.shell_path,
                                input_len = data_str.len(),
                                input_preview = %truncate_for_log(data_str, 200),
                                error = %err,
                                error_debug = ?err,
                                "rule evaluate failed"
                            );
                        }
                        Err(_) => { // Timeout Error
                            tracing::warn!(
                                target: "security::detect",
                                rule_name = %rule.meta().name,
                                rule_default_severity = ?rule.meta().default_severity,
                                shell = %ctx.shell_path,
                                input_len = data_str.len(),
                                input_preview = %truncate_for_log(data_str, 200),
                                "rule evaluate timed out after {:?}", rule_timeout
                            );
                        }
                    }
                }
                Err(join_err) => {
                    if join_err.is_panic() {
                        tracing::error!("Rule evaluation task panicked: {}", join_err);
                    }
                }
            }

            // 窗口滑动：完成一个任务，再补充一个新任务
            if let Some(rule) = rules_iter.next() {
                spawn_rule(&mut set, rule, Arc::clone(&data_arc), Arc::clone(ctx));
            }
        }

        ctx.push_history(data).await;

        if !hits.is_empty() {
            DetectResult::ThreatDetected(hits)
        } else if evaluated == 0 {
            DetectResult::Unknown
        } else {
            DetectResult::Safe
        }
    }
}