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
11use async_trait::async_trait;
12use schemars::JsonSchema;
13use serde::{Deserialize, Serialize};
14use tokio::process::Command;
15
16use lc_core::tools::{BaseTool, Tool, ToolError};
17
18/// Python REPL 工具输入
19#[derive(Debug, Deserialize, JsonSchema)]
20pub struct PythonREPLInput {
21    /// 要执行的 Python 代码
22    pub code: String,
23    /// 超时时间（秒，默认 30）
24    pub timeout_seconds: Option<u64>,
25}
26
27/// Python REPL 工具输出
28#[derive(Debug, Serialize)]
29pub struct PythonREPLOutput {
30    /// 执行的代码
31    pub code: String,
32    /// 标准输出
33    pub stdout: String,
34    /// 标准错误
35    pub stderr: String,
36    /// 退出码
37    pub exit_code: i32,
38}
39
40/// Dangerous Python modules that are blocked for security.
41const BLOCKED_IMPORTS: &[&str] = &[
42    "os",
43    "subprocess",
44    "sys",
45    "shutil",
46    "signal",
47    "ctypes",
48    "multiprocessing",
49    "socket",
50    "http.server",
51    "xmlrpc",
52    "pickle",
53    "shelve",
54    "importlib",
55    "code",
56    "codeop",
57    "compileall",
58    "pty",
59    "commands",
60    "pdb",
61    "webbrowser",
62    "antigravity",
63];
64
65/// Check if Python code contains dangerous imports.
66fn contains_dangerous_import(code: &str) -> Option<String> {
67    for line in code.lines() {
68        let trimmed = line.trim();
69        if trimmed.starts_with('#') {
70            continue;
71        }
72        let code_part = trimmed.split('#').next().unwrap_or(trimmed);
73        if code_part.contains("import") {
74            for blocked in BLOCKED_IMPORTS {
75                if code_part.contains(&format!("import {}", blocked))
76                    || code_part.contains(&format!("from {} ", blocked))
77                    || code_part.contains(&format!("from {}.", blocked))
78                    || code_part.contains(&format!("from {}import", blocked))
79                {
80                    return Some(blocked.to_string());
81                }
82            }
83        }
84    }
85    None
86}
87
88/// Python 代码执行工具
89///
90/// 在本地 Python 环境中执行代码并返回结果。
91/// 适用于数学计算、数据处理等需要 Python 生态的场景。
92///
93/// # 安全警告
94/// 此工具默认**禁用**。必须调用 [`PythonREPLTool::with_dangerously_allow`]
95/// 才能执行代码。在生产环境中使用时应确保在沙箱环境中运行。
96pub struct PythonREPLTool {
97    python_path: String,
98    /// 是否允许执行代码（默认 false，必须显式 opt-in）
99    dangerously_allow: bool,
100    /// 是否启用危险 import 检查（默认 true）
101    check_dangerous_imports: bool,
102}
103
104impl PythonREPLTool {
105    pub fn new() -> Self {
106        Self {
107            python_path: Self::find_python(),
108            dangerously_allow: false,
109            check_dangerous_imports: true,
110        }
111    }
112
113    /// 使用自定义 Python 路径
114    pub fn with_python_path(path: impl Into<String>) -> Self {
115        Self {
116            python_path: path.into(),
117            dangerously_allow: false,
118            check_dangerous_imports: true,
119        }
120    }
121
122    /// 显式启用代码执行（默认禁用）
123    pub fn with_dangerously_allow(mut self, allow: bool) -> Self {
124        self.dangerously_allow = allow;
125        self
126    }
127
128    /// Disable dangerous import checking (default: enabled).
129    pub fn with_skip_dangerous_imports_check(mut self, skip: bool) -> Self {
130        self.check_dangerous_imports = !skip;
131        self
132    }
133
134    /// 自动查找系统 Python
135    fn find_python() -> String {
136        for candidate in &["python3", "python"] {
137            if std::process::Command::new(candidate)
138                .arg("--version")
139                .output()
140                .is_ok()
141            {
142                return candidate.to_string();
143            }
144        }
145        "python3".to_string()
146    }
147}
148
149impl Default for PythonREPLTool {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155#[async_trait]
156impl Tool for PythonREPLTool {
157    type Input = PythonREPLInput;
158    type Output = PythonREPLOutput;
159
160    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
161        if input.code.trim().is_empty() {
162            return Err(ToolError::InvalidInput(
163                "Python code must not be empty".to_string(),
164            ));
165        }
166
167        if !self.dangerously_allow {
168            return Err(ToolError::ExecutionFailed(
169                "PythonREPLTool is disabled by default for security. \
170                 Call .with_dangerously_allow(true) to enable execution."
171                    .to_string(),
172            ));
173        }
174
175        if self.check_dangerous_imports {
176            if let Some(blocked) = contains_dangerous_import(&input.code) {
177                return Err(ToolError::ExecutionFailed(format!(
178                    "Code contains dangerous import: '{}'. \
179                     This module is blocked for security. \
180                     Call .with_skip_dangerous_imports_check(true) to bypass (not recommended).",
181                    blocked
182                )));
183            }
184        }
185
186        let timeout_secs = input.timeout_seconds.unwrap_or(30);
187
188        let result = tokio::time::timeout(
189            std::time::Duration::from_secs(timeout_secs),
190            Command::new(&self.python_path)
191                .arg("-c")
192                .arg(&input.code)
193                .output(),
194        )
195        .await
196        .map_err(|_| {
197            ToolError::ExecutionFailed(format!(
198                "Python execution timed out after {} seconds",
199                timeout_secs
200            ))
201        })?
202        .map_err(|e| ToolError::ExecutionFailed(format!("Python execution failed: {}", e)))?;
203
204        let stdout = String::from_utf8_lossy(&result.stdout).to_string();
205        let stderr = String::from_utf8_lossy(&result.stderr).to_string();
206        let exit_code = result.status.code().unwrap_or(-1);
207
208        Ok(PythonREPLOutput {
209            code: input.code,
210            stdout,
211            stderr,
212            exit_code,
213        })
214    }
215}
216
217#[async_trait]
218impl BaseTool for PythonREPLTool {
219    fn name(&self) -> &str {
220        "python_repl"
221    }
222
223    fn description(&self) -> &str {
224        "Python code execution tool. Runs code in a local Python environment and returns results.
225
226Parameters:
227- code: Python code string to execute
228- timeout_seconds: Timeout in seconds (default: 30)
229
230Supports any Python syntax, including math, data processing, plotting, etc.
231
232SECURITY WARNING: Disabled by default. Must call .with_dangerously_allow(true) to enable.
233Only use in controlled/sandboxed environments.
234
235Examples:
236- Simple calc: {\"code\": \"print(1 + 2)\"}
237- List processing: {\"code\": \"print([x**2 for x in range(10)])\"}
238- Math: {\"code\": \"import math; print(math.pi)\"}"
239    }
240
241    async fn run(&self, input: String) -> Result<String, ToolError> {
242        let parsed: PythonREPLInput = serde_json::from_str(&input)
243            .map_err(|e| ToolError::InvalidInput(format!("JSON parse error: {}", e)))?;
244
245        let output = self.invoke(parsed).await?;
246
247        let mut result = String::new();
248        if !output.stdout.is_empty() {
249            result.push_str(&format!("stdout:\n{}\n", output.stdout));
250        }
251        if !output.stderr.is_empty() {
252            result.push_str(&format!("stderr:\n{}\n", output.stderr));
253        }
254        result.push_str(&format!("exit_code: {}", output.exit_code));
255
256        Ok(result)
257    }
258
259    fn args_schema(&self) -> Option<serde_json::Value> {
260        use schemars::schema_for;
261        serde_json::to_value(schema_for!(PythonREPLInput)).ok()
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    fn test_python_repl_tool_properties() {
271        let tool = PythonREPLTool::new();
272        assert_eq!(tool.name(), "python_repl");
273        assert!(tool.description().contains("Python"));
274        assert!(BaseTool::args_schema(&tool).is_some());
275    }
276
277    #[tokio::test]
278    async fn test_python_repl_empty_code() {
279        let tool = PythonREPLTool::new().with_dangerously_allow(true);
280        let result = tool.run(r#"{"code": ""}"#.to_string()).await;
281        assert!(result.is_err());
282    }
283
284    #[tokio::test]
285    async fn test_python_repl_disabled_by_default() {
286        let tool = PythonREPLTool::new();
287        let result = tool
288            .invoke(PythonREPLInput {
289                code: "print(1 + 2)".to_string(),
290                timeout_seconds: Some(10),
291            })
292            .await;
293        assert!(result.is_err());
294        let err_msg = result.unwrap_err().to_string();
295        assert!(
296            err_msg.contains("disabled by default"),
297            "Expected disabled error, got: {}",
298            err_msg
299        );
300    }
301
302    #[tokio::test]
303    async fn test_python_repl_basic_execution() {
304        let tool = PythonREPLTool::new().with_dangerously_allow(true);
305        let result = tool
306            .invoke(PythonREPLInput {
307                code: "print(1 + 2)".to_string(),
308                timeout_seconds: Some(10),
309            })
310            .await;
311
312        match result {
313            Ok(output) => {
314                if output.exit_code == 0 || !output.stdout.is_empty() {
315                    // Python available, verify functionality
316                } else {
317                    eprintln!(
318                        "Python may not be installed (exit_code={})",
319                        output.exit_code
320                    );
321                }
322            }
323            Err(e) => {
324                eprintln!("Python not available (may be expected): {}", e);
325            }
326        }
327    }
328
329    #[test]
330    fn test_dangerous_import_detection() {
331        assert!(contains_dangerous_import("import os").is_some());
332        assert!(contains_dangerous_import("import subprocess").is_some());
333        assert!(contains_dangerous_import("from sys import path").is_some());
334        assert!(contains_dangerous_import("from os.path import join").is_some());
335        // Safe imports should pass
336        assert!(contains_dangerous_import("import math").is_none());
337        assert!(contains_dangerous_import("import json").is_none());
338        assert!(contains_dangerous_import("from datetime import datetime").is_none());
339        // Comments should be ignored
340        assert!(contains_dangerous_import("# import os").is_none());
341    }
342
343    #[tokio::test]
344    async fn test_python_repl_blocks_dangerous_import() {
345        let tool = PythonREPLTool::new().with_dangerously_allow(true);
346        let result = tool
347            .invoke(PythonREPLInput {
348                code: "import os; print(os.getcwd())".to_string(),
349                timeout_seconds: Some(10),
350            })
351            .await;
352        assert!(result.is_err());
353        let err_msg = result.unwrap_err().to_string();
354        assert!(
355            err_msg.contains("dangerous import"),
356            "Expected dangerous import error, got: {}",
357            err_msg
358        );
359    }
360
361    #[tokio::test]
362    async fn test_python_repl_allows_safe_import() {
363        let tool = PythonREPLTool::new().with_dangerously_allow(true);
364        let result = tool
365            .invoke(PythonREPLInput {
366                code: "import math; print(math.pi)".to_string(),
367                timeout_seconds: Some(10),
368            })
369            .await;
370        match result {
371            Ok(output) => {
372                if output.exit_code == 0 {
373                    assert!(output.stdout.contains("3.14"));
374                }
375            }
376            Err(e) => {
377                assert!(
378                    !e.to_string().contains("dangerous import"),
379                    "math should not be blocked: {}",
380                    e
381                );
382            }
383        }
384    }
385
386    #[tokio::test]
387    async fn test_python_repl_with_error() {
388        let tool = PythonREPLTool::new().with_dangerously_allow(true);
389        let result = tool
390            .invoke(PythonREPLInput {
391                code: "print(undefined_var)".to_string(),
392                timeout_seconds: Some(10),
393            })
394            .await;
395
396        match result {
397            Ok(output) => {
398                if output.exit_code == 0 {
399                    // occasionally Python available but no error reported
400                }
401            }
402            Err(_) => {
403                // No Python available, skip
404            }
405        }
406    }
407}
lc_tools/python_repl.rs

lc_tools/
python_repl.rs
