Skip to main content

lc_tools/
python_repl.rs

1// lc-tools/src/python_repl.rs
2//! Python 代码执行工具
3//!
4//! 调用系统 Python 解释器执行代码并返回结果。
5//! 需要系统中已安装 Python。
6//!
7//! # 安全警告
8//! 此工具默认**禁用**,必须调用 `with_dangerously_allow(true)` 显式启用。
9//! 启用后会执行任意 Python 代码,仅应在受控/沙箱环境中使用。
10//!
11//! 内置的"危险 import 黑名单"只是**噪音过滤,不是安全边界**——它挡不住
12//! `__import__`/`eval`/`exec`/字符串拼接等编码混淆,也会误伤字符串字面量。
13//! 真正的隔离必须走沙箱([`crate::sandbox`]),黑名单只用于减少误入沙箱的噪音。
14
15use async_trait::async_trait;
16use regex::Regex;
17use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19use tokio::process::Command;
20
21use lc_core::tools::{BaseTool, Tool, ToolError};
22
23/// Python REPL 工具输入
24#[derive(Debug, Deserialize, JsonSchema)]
25pub struct PythonREPLInput {
26    /// 要执行的 Python 代码
27    pub code: String,
28    /// 超时时间(秒,默认 30)
29    pub timeout_seconds: Option<u64>,
30}
31
32/// Python REPL 工具输出
33#[derive(Debug, Serialize)]
34pub struct PythonREPLOutput {
35    /// 执行的代码
36    pub code: String,
37    /// 标准输出
38    pub stdout: String,
39    /// 标准错误
40    pub stderr: String,
41    /// 退出码
42    pub exit_code: i32,
43}
44
45/// Dangerous Python modules that are blocked for security.
46const BLOCKED_IMPORTS: &[&str] = &[
47    "os",
48    "subprocess",
49    "sys",
50    "shutil",
51    "signal",
52    "ctypes",
53    "multiprocessing",
54    "socket",
55    "http.server",
56    "xmlrpc",
57    "pickle",
58    "shelve",
59    "importlib",
60    "code",
61    "codeop",
62    "compileall",
63    "pty",
64    "commands",
65    "pdb",
66    "webbrowser",
67    "antigravity",
68];
69
70/// 常见绕过 import 黑名单的危险内建调用(单词边界 + 函数调用形式)。
71const DANGEROUS_BUILTIN_CALLS: &[&str] = &[
72    "__import__",
73    "import_module",
74    "eval",
75    "exec",
76    "execfile",
77    "compile",
78];
79
80/// 匹配 `__import__(` / `import_module(` / `eval(` / `exec(` / `compile(` 等危险调用。
81///
82/// `\b` 保证不会误伤 `evaluate(` / `execute(` / `length(` 这类含子串的普通词;
83/// 但仍会误伤字符串字面量里的 `"eval(...)"` 字样——这是字符串级拦截的固有局限,
84/// 详见 [`contains_dangerous_code`] 的安全定位说明。
85static DANGEROUS_CALL_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
86    Regex::new(&format!(
87        r"\b(?:{})\s*\(",
88        DANGEROUS_BUILTIN_CALLS.join("|")
89    ))
90    .unwrap()
91});
92
93/// Check if Python code contains dangerous imports or builtin calls.
94///
95/// **安全定位**:这是噪音过滤层,**不是安全边界**。逐行子串/正则匹配永远可以被
96/// unicode 混淆、`"o"+"s"` 拼接、`().__class__` 反射等绕过,也会误伤字符串字面量。
97/// 不可信代码必须走沙箱([`crate::sandbox`])。
98fn contains_dangerous_code(code: &str) -> Option<String> {
99    for line in code.lines() {
100        let trimmed = line.trim();
101        if trimmed.starts_with('#') {
102            continue;
103        }
104        // 去掉行内注释(# 之后的部分),避免注释内容误伤。
105        let code_part = trimmed.split('#').next().unwrap_or(trimmed);
106
107        // 1) 危险 import 检查(BLOCKED_IMPORTS)
108        if code_part.contains("import") {
109            for blocked in BLOCKED_IMPORTS {
110                if code_part.contains(&format!("import {}", blocked))
111                    || code_part.contains(&format!("from {} ", blocked))
112                    || code_part.contains(&format!("from {}.", blocked))
113                    || code_part.contains(&format!("from {}import", blocked))
114                {
115                    return Some(blocked.to_string());
116                }
117            }
118        }
119
120        // 2) 危险内建调用检查(__import__ / import_module / eval / exec / compile 等常见绕过)
121        if let Some(call) = DANGEROUS_CALL_REGEX.find(code_part) {
122            return Some(call.as_str().to_string());
123        }
124    }
125    None
126}
127
128/// Python 代码执行工具
129///
130/// 在本地 Python 环境中执行代码并返回结果。
131/// 适用于数学计算、数据处理等需要 Python 生态的场景。
132///
133/// # 安全警告
134/// 此工具默认**禁用**。必须调用 [`PythonREPLTool::with_dangerously_allow`]
135/// 才能执行代码。在生产环境中使用时应确保在沙箱环境中运行。
136pub struct PythonREPLTool {
137    python_path: String,
138    /// 是否允许执行代码(默认 false,必须显式 opt-in)
139    dangerously_allow: bool,
140    /// 是否启用危险 import 检查(默认 true)
141    check_dangerous_imports: bool,
142}
143
144impl PythonREPLTool {
145    /// 创建 Python 代码执行工具(默认禁用执行)。
146    pub fn new() -> Self {
147        Self {
148            python_path: Self::find_python(),
149            dangerously_allow: false,
150            check_dangerous_imports: true,
151        }
152    }
153
154    /// 使用自定义 Python 路径
155    pub fn with_python_path(path: impl Into<String>) -> Self {
156        Self {
157            python_path: path.into(),
158            dangerously_allow: false,
159            check_dangerous_imports: true,
160        }
161    }
162
163    /// 显式启用代码执行(默认禁用)
164    pub fn with_dangerously_allow(mut self, allow: bool) -> Self {
165        self.dangerously_allow = allow;
166        self
167    }
168
169    /// Disable dangerous import checking (default: enabled).
170    pub fn with_skip_dangerous_imports_check(mut self, skip: bool) -> Self {
171        self.check_dangerous_imports = !skip;
172        self
173    }
174
175    /// 自动查找系统 Python
176    fn find_python() -> String {
177        for candidate in &["python3", "python"] {
178            if std::process::Command::new(candidate)
179                .arg("--version")
180                .output()
181                .is_ok()
182            {
183                return candidate.to_string();
184            }
185        }
186        "python3".to_string()
187    }
188}
189
190impl Default for PythonREPLTool {
191    fn default() -> Self {
192        Self::new()
193    }
194}
195
196#[async_trait]
197impl Tool for PythonREPLTool {
198    type Input = PythonREPLInput;
199    type Output = PythonREPLOutput;
200
201    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
202        if input.code.trim().is_empty() {
203            return Err(ToolError::InvalidInput(
204                "Python code must not be empty".to_string(),
205            ));
206        }
207
208        if !self.dangerously_allow {
209            return Err(ToolError::ExecutionFailed(
210                "PythonREPLTool is disabled by default for security. \
211                 Call .with_dangerously_allow(true) to enable execution."
212                    .to_string(),
213            ));
214        }
215
216        if self.check_dangerous_imports {
217            if let Some(blocked) = contains_dangerous_code(&input.code) {
218                return Err(ToolError::ExecutionFailed(format!(
219                    "Code contains dangerous import or builtin call: '{}'. \
220                     Blocked by the noise-filter blacklist (note: this is not a security boundary; \
221                     untrusted code must run in a sandbox). \
222                     Call .with_skip_dangerous_imports_check(true) to bypass (not recommended).",
223                    blocked
224                )));
225            }
226        }
227
228        let timeout_secs = input.timeout_seconds.unwrap_or(30);
229
230        let result = tokio::time::timeout(
231            std::time::Duration::from_secs(timeout_secs),
232            Command::new(&self.python_path)
233                .arg("-c")
234                .arg(&input.code)
235                .output(),
236        )
237        .await
238        .map_err(|_| {
239            ToolError::ExecutionFailed(format!(
240                "Python execution timed out after {} seconds",
241                timeout_secs
242            ))
243        })?
244        .map_err(|e| ToolError::ExecutionFailed(format!("Python execution failed: {}", e)))?;
245
246        let stdout = String::from_utf8_lossy(&result.stdout).to_string();
247        let stderr = String::from_utf8_lossy(&result.stderr).to_string();
248        let exit_code = result.status.code().unwrap_or(-1);
249
250        Ok(PythonREPLOutput {
251            code: input.code,
252            stdout,
253            stderr,
254            exit_code,
255        })
256    }
257}
258
259#[async_trait]
260impl BaseTool for PythonREPLTool {
261    fn name(&self) -> &str {
262        "python_repl"
263    }
264
265    fn description(&self) -> &str {
266        "Python code execution tool. Runs code in a local Python environment and returns results.
267
268Parameters:
269- code: Python code string to execute
270- timeout_seconds: Timeout in seconds (default: 30)
271
272Supports any Python syntax, including math, data processing, plotting, etc.
273
274SECURITY WARNING: Disabled by default. Must call .with_dangerously_allow(true) to enable.
275Only use in controlled/sandboxed environments.
276
277Examples:
278- Simple calc: {\"code\": \"print(1 + 2)\"}
279- List processing: {\"code\": \"print([x**2 for x in range(10)])\"}
280- Math: {\"code\": \"import math; print(math.pi)\"}"
281    }
282
283    async fn run(&self, input: String) -> Result<String, ToolError> {
284        let parsed: PythonREPLInput = serde_json::from_str(&input)
285            .map_err(|e| ToolError::InvalidInput(format!("JSON parse error: {}", e)))?;
286
287        let output = self.invoke(parsed).await?;
288
289        let mut result = String::new();
290        if !output.stdout.is_empty() {
291            result.push_str(&format!("stdout:\n{}\n", output.stdout));
292        }
293        if !output.stderr.is_empty() {
294            result.push_str(&format!("stderr:\n{}\n", output.stderr));
295        }
296        result.push_str(&format!("exit_code: {}", output.exit_code));
297
298        Ok(result)
299    }
300
301    fn args_schema(&self) -> Option<serde_json::Value> {
302        use schemars::schema_for;
303        serde_json::to_value(schema_for!(PythonREPLInput)).ok()
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn test_python_repl_tool_properties() {
313        let tool = PythonREPLTool::new();
314        assert_eq!(tool.name(), "python_repl");
315        assert!(tool.description().contains("Python"));
316        assert!(BaseTool::args_schema(&tool).is_some());
317    }
318
319    #[tokio::test]
320    async fn test_python_repl_empty_code() {
321        let tool = PythonREPLTool::new().with_dangerously_allow(true);
322        let result = tool.run(r#"{"code": ""}"#.to_string()).await;
323        assert!(result.is_err());
324    }
325
326    #[tokio::test]
327    async fn test_python_repl_disabled_by_default() {
328        let tool = PythonREPLTool::new();
329        let result = tool
330            .invoke(PythonREPLInput {
331                code: "print(1 + 2)".to_string(),
332                timeout_seconds: Some(10),
333            })
334            .await;
335        assert!(result.is_err());
336        let err_msg = result.unwrap_err().to_string();
337        assert!(
338            err_msg.contains("disabled by default"),
339            "Expected disabled error, got: {}",
340            err_msg
341        );
342    }
343
344    #[tokio::test]
345    async fn test_python_repl_basic_execution() {
346        let tool = PythonREPLTool::new().with_dangerously_allow(true);
347        let result = tool
348            .invoke(PythonREPLInput {
349                code: "print(1 + 2)".to_string(),
350                timeout_seconds: Some(10),
351            })
352            .await;
353
354        match result {
355            Ok(output) => {
356                if output.exit_code == 0 || !output.stdout.is_empty() {
357                    // Python available, verify functionality
358                } else {
359                    eprintln!(
360                        "Python may not be installed (exit_code={})",
361                        output.exit_code
362                    );
363                }
364            }
365            Err(e) => {
366                eprintln!("Python not available (may be expected): {}", e);
367            }
368        }
369    }
370
371    #[test]
372    fn test_dangerous_import_detection() {
373        assert!(contains_dangerous_code("import os").is_some());
374        assert!(contains_dangerous_code("import subprocess").is_some());
375        assert!(contains_dangerous_code("from sys import path").is_some());
376        assert!(contains_dangerous_code("from os.path import join").is_some());
377        // Safe imports should pass
378        assert!(contains_dangerous_code("import math").is_none());
379        assert!(contains_dangerous_code("import json").is_none());
380        assert!(contains_dangerous_code("from datetime import datetime").is_none());
381        // Comments should be ignored
382        assert!(contains_dangerous_code("# import os").is_none());
383    }
384
385    #[test]
386    fn test_dangerous_builtin_calls_detected() {
387        // 常见绕过:不经过 import 语句,直接调用内建/导入函数
388        assert!(contains_dangerous_code("__import__('os').system('ls')").is_some());
389        assert!(contains_dangerous_code("importlib.import_module('os')").is_some());
390        assert!(contains_dangerous_code("eval('os')").is_some());
391        assert!(contains_dangerous_code("exec('import os')").is_some());
392        assert!(contains_dangerous_code("compile('import os', '<x>', 'exec')").is_some());
393        assert!(contains_dangerous_code("execfile('/tmp/x.py')").is_some());
394    }
395
396    #[test]
397    fn test_dangerous_builtin_calls_no_false_positive_on_words() {
398        // 单词边界:不误伤 evaluate / execute / length 等含子串的普通词
399        assert!(contains_dangerous_code("print('evaluate the result')").is_none());
400        assert!(contains_dangerous_code("result = execute_query()").is_none());
401        assert!(contains_dangerous_code("print(len([1, 2, 3]))").is_none());
402        assert!(contains_dangerous_code("x = len('hello')").is_none());
403    }
404
405    #[tokio::test]
406    async fn test_python_repl_blocks_dangerous_import() {
407        let tool = PythonREPLTool::new().with_dangerously_allow(true);
408        let result = tool
409            .invoke(PythonREPLInput {
410                code: "import os; print(os.getcwd())".to_string(),
411                timeout_seconds: Some(10),
412            })
413            .await;
414        assert!(result.is_err());
415        let err_msg = result.unwrap_err().to_string();
416        assert!(
417            err_msg.contains("dangerous import"),
418            "Expected dangerous import error, got: {}",
419            err_msg
420        );
421    }
422
423    #[tokio::test]
424    async fn test_python_repl_allows_safe_import() {
425        let tool = PythonREPLTool::new().with_dangerously_allow(true);
426        let result = tool
427            .invoke(PythonREPLInput {
428                code: "import math; print(math.pi)".to_string(),
429                timeout_seconds: Some(10),
430            })
431            .await;
432        match result {
433            Ok(output) => {
434                if output.exit_code == 0 {
435                    assert!(output.stdout.contains("3.14"));
436                }
437            }
438            Err(e) => {
439                assert!(
440                    !e.to_string().contains("dangerous import"),
441                    "math should not be blocked: {}",
442                    e
443                );
444            }
445        }
446    }
447
448    #[tokio::test]
449    async fn test_python_repl_with_error() {
450        let tool = PythonREPLTool::new().with_dangerously_allow(true);
451        let result = tool
452            .invoke(PythonREPLInput {
453                code: "print(undefined_var)".to_string(),
454                timeout_seconds: Some(10),
455            })
456            .await;
457
458        match result {
459            Ok(output) => {
460                if output.exit_code == 0 {
461                    // occasionally Python available but no error reported
462                }
463            }
464            Err(_) => {
465                // No Python available, skip
466            }
467        }
468    }
469}