DevFlow/crates/df-nodes/src/script_node.rs

//! 脚本节点 — 执行 Shell 脚本或自定义命令

use async_trait::async_trait;
use df_workflow::node::{Node, NodeContext, NodeOutput, NodeResult, NodeSchema};

/// 脚本节点
pub struct ScriptNode;

// ── 运行时白/黑名单(通过 set_script_safety_config 注入,优先于 env var;env var 作回退) ──
use std::sync::OnceLock;

/// 运行时白名单配置(前端设置页写入,优先于环境变量)
static RUNTIME_SCRIPT_WHITELIST: OnceLock<String> = OnceLock::new();
/// 运行时黑名单配置(前端设置页写入,优先于环境变量)
static RUNTIME_SCRIPT_BLACKLIST: OnceLock<String> = OnceLock::new();

/// 注入运行时脚本安全配置(由 Tauri IPC 调用,前端设置页保存后触发)。
/// 空字符串表示「未配置」,回退到环境变量 / 默认行为。
pub fn set_script_safety_config(whitelist: &str, blacklist: &str) {
    let wl = whitelist.trim();
    let bl = blacklist.trim();
    if !wl.is_empty() { let _ = RUNTIME_SCRIPT_WHITELIST.set(wl.to_string()); }
    if !bl.is_empty() { let _ = RUNTIME_SCRIPT_BLACKLIST.set(bl.to_string()); }
}

#[async_trait]
impl Node for ScriptNode {
    async fn execute(&self, ctx: NodeContext) -> NodeResult {
        tracing::info!("ScriptNode 执行: {}", ctx.node_id);

        // 从 config 解析 command（必填）
        let command = ctx
            .config
            .get("command")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("ScriptNode 缺少必填参数: command"))?;

        // 解析可选参数
        let timeout_secs = ctx
            .config
            .get("timeout_secs")
            .and_then(|v| v.as_u64());

        let working_dir = ctx
            .config
            .get("working_dir")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        // 构建请求并执行
        let request = df_execute::shell::ShellRequest {
            command: command.to_string(),
            working_dir,
            env: std::collections::HashMap::new(),
            timeout_secs,
            shell_type: Default::default(),
        };

        // 命令执行安全:白/黑名单校验(运行时配置优先,回退环境变量/默认黑名单)。
        // 命令名取首词(shell 第一段,如 `rm -rf /` 取 `rm`),按 trim + 小写规范化比较。
        let cmd_name = command.split_whitespace().next().unwrap_or("").to_lowercase();
        if let Some(denied) = check_command_policy(&cmd_name) {
            tracing::warn!(
                command = %command,
                reason = %denied,
                "ScriptNode 命令被策略拒绝"
            );
            anyhow::bail!("脚本命令被策略拒绝: {} (命令: {})", denied, command);
        }

        // 危险关键词告警:匹配到仅告警不阻止执行(与硬拒绝的白/黑名单正交)
        let dangerous_keywords = ["rm -rf", "DROP TABLE", "Format", "del /f", "shutdown"];
        for &kw in &dangerous_keywords {
            if command.contains(kw) {
                tracing::warn!(
                    keyword = %kw,
                    command = %command,
                    "ScriptNode 即将执行包含危险关键词的命令"
                );
            }
        }

        tracing::info!("ScriptNode 执行命令: {}", command);
        let result = df_execute::shell::execute(request).await?;

        // 非零退出码视为执行失败
        let exit_code = result.exit_code.unwrap_or(-1);
        if exit_code != 0 {
            anyhow::bail!(
                "脚本执行失败 (exit_code={}): {}",
                exit_code,
                result.stderr.trim()
            );
        }

        tracing::info!(
            "ScriptNode 完成: 耗时 {}ms, exit_code={}",
            result.duration_ms,
            exit_code
        );

        Ok(NodeOutput::from_value(serde_json::json!({
            "stdout": result.stdout,
            "stderr": result.stderr,
            "exit_code": exit_code,
            "duration_ms": result.duration_ms,
        })))
    }

    fn schema(&self) -> NodeSchema {
        NodeSchema {
            params: serde_json::json!({
                "type": "object",
                "properties": {
                    "command": { "type": "string" },
                    "timeout_secs": { "type": "integer" },
                    "working_dir": { "type": "string" }
                },
                "required": ["command"]
            }),
            output: serde_json::json!({
                "type": "object",
                "properties": {
                    "stdout": { "type": "string" },
                    "stderr": { "type": "string" },
                    "exit_code": { "type": "integer" },
                    "duration_ms": { "type": "integer" }
                }
            }),
        }
    }

    fn node_type(&self) -> &str {
        "script"
    }
}

/// 默认黑名单:即使环境变量未设置也生效,挡最危险的破坏性命令。
/// 用户设置 `DF_SCRIPT_BLACKLIST` 时覆盖默认(不合并,用户显式控制)。
const DEFAULT_BLACKLIST: &[&str] = &["rm", "del", "format", "shutdown", "mkfs", "dd"];

/// 命令执行策略校验:从运行时配置 + 环境变量读取白/黑名单。
///
/// 优先级(高→低):运行时配置(前端设置页写入) > 环境变量 > 默认黑名单兜底。
/// 黑名单优先于白名单(命中黑名单总是拒绝,即便同时在白名单);命令名比较前 trim + 小写。
///
/// 返回 `Some(reason)` 表示拒绝;返回 `None` 表示放行。
fn check_command_policy(cmd_name: &str) -> Option<&'static str> {
    // 运行时配置优先(前端设置页写入,优先级最高)
    if let Some(raw) = RUNTIME_SCRIPT_BLACKLIST.get() {
        let list: Vec<&str> = raw.split(',').map(|s| s.trim()).filter(|s| !s.is_empty()).collect();
        if list.iter().any(|c| c.eq_ignore_ascii_case(cmd_name)) {
            return Some("命令在黑名单中(运行时配置)");
        }
    }
    // 运行时白名单非空时,命令须在名单内(命中即放行,跳过环境变量)
    if let Some(raw) = RUNTIME_SCRIPT_WHITELIST.get() {
        let list: Vec<&str> = raw.split(',').map(|s| s.trim()).filter(|s| !s.is_empty()).collect();
        if !list.is_empty() && !list.iter().any(|c| c.eq_ignore_ascii_case(cmd_name)) {
            return Some("命令不在白名单中(运行时配置)");
        }
        if !list.is_empty() {
            return None;
        }
    }

    // 回退环境变量(兼容旧部署);黑名单缺省用默认黑名单兜底
    let env_blacklist = std::env::var("DF_SCRIPT_BLACKLIST").ok();
    let blacklist: Vec<String> = match env_blacklist {
        Some(raw) => raw
            .split(',')
            .map(|s| s.trim().to_lowercase())
            .filter(|s| !s.is_empty())
            .collect(),
        None => DEFAULT_BLACKLIST.iter().map(|s| s.to_string()).collect(),
    };
    if blacklist.iter().any(|c| c == cmd_name) {
        return Some("命令在黑名单中");
    }

    // 环境变量白名单非空时,命令名须在名单内
    if let Ok(whitelist_raw) = std::env::var("DF_SCRIPT_WHITELIST") {
        let whitelist: Vec<String> = whitelist_raw
            .split(',')
            .map(|s| s.trim().to_lowercase())
            .filter(|s| !s.is_empty())
            .collect();
        if !whitelist.is_empty() && !whitelist.iter().any(|c| c == cmd_name) {
            return Some("命令不在白名单中");
        }
    }

    None
}