1use async_trait::async_trait;
17use regex::Regex;
18use schemars::JsonSchema;
19use serde::{Deserialize, Serialize};
20use tokio::process::Command;
21
22use lc_core::tools::{BaseTool, Tool, ToolError};
23
24#[derive(Debug, Deserialize, JsonSchema)]
26pub struct PythonREPLInput {
27 pub code: String,
29 pub timeout_seconds: Option<u64>,
31}
32
33#[derive(Debug, Serialize)]
35pub struct PythonREPLOutput {
36 pub code: String,
38 pub stdout: String,
40 pub stderr: String,
42 pub exit_code: i32,
44}
45
46const BLOCKED_IMPORTS: &[&str] = &[
48 "os",
49 "subprocess",
50 "sys",
51 "shutil",
52 "signal",
53 "ctypes",
54 "multiprocessing",
55 "socket",
56 "http.server",
57 "xmlrpc",
58 "pickle",
59 "shelve",
60 "importlib",
61 "code",
62 "codeop",
63 "compileall",
64 "pty",
65 "commands",
66 "pdb",
67 "webbrowser",
68 "antigravity",
69];
70
71const DANGEROUS_BUILTIN_CALLS: &[&str] = &[
73 "__import__",
74 "import_module",
75 "eval",
76 "exec",
77 "execfile",
78 "compile",
79];
80
81static DANGEROUS_CALL_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
88 Regex::new(&format!(
89 r"\b(?:{})\s*\(",
90 DANGEROUS_BUILTIN_CALLS.join("|")
91 ))
92 .unwrap()
93});
94
95fn contains_dangerous_code(code: &str) -> Option<String> {
102 for line in code.lines() {
103 let trimmed = line.trim();
104 if trimmed.starts_with('#') {
105 continue;
106 }
107 let code_part = trimmed.split('#').next().unwrap_or(trimmed);
109
110 if code_part.contains("import") {
112 for blocked in BLOCKED_IMPORTS {
113 if code_part.contains(&format!("import {}", blocked))
114 || code_part.contains(&format!("from {} ", blocked))
115 || code_part.contains(&format!("from {}.", blocked))
116 || code_part.contains(&format!("from {}import", blocked))
117 {
118 return Some(blocked.to_string());
119 }
120 }
121 }
122
123 if let Some(call) = DANGEROUS_CALL_REGEX.find(code_part) {
125 return Some(call.as_str().to_string());
126 }
127 }
128 None
129}
130
131pub struct PythonREPLTool {
140 python_path: String,
141 dangerously_allow: bool,
143 check_dangerous_imports: bool,
145}
146
147impl PythonREPLTool {
148 pub fn new() -> Self {
150 Self {
151 python_path: Self::find_python(),
152 dangerously_allow: false,
153 check_dangerous_imports: true,
154 }
155 }
156
157 pub fn with_python_path(path: impl Into<String>) -> Self {
159 Self {
160 python_path: path.into(),
161 dangerously_allow: false,
162 check_dangerous_imports: true,
163 }
164 }
165
166 pub fn with_dangerously_allow(mut self, allow: bool) -> Self {
168 self.dangerously_allow = allow;
169 self
170 }
171
172 pub fn with_skip_dangerous_imports_check(mut self, skip: bool) -> Self {
174 self.check_dangerous_imports = !skip;
175 self
176 }
177
178 fn find_python() -> String {
180 for candidate in &["python3", "python"] {
181 if std::process::Command::new(candidate)
182 .arg("--version")
183 .output()
184 .is_ok()
185 {
186 return candidate.to_string();
187 }
188 }
189 "python3".to_string()
190 }
191}
192
193impl Default for PythonREPLTool {
194 fn default() -> Self {
195 Self::new()
196 }
197}
198
199#[async_trait]
200impl Tool for PythonREPLTool {
201 type Input = PythonREPLInput;
202 type Output = PythonREPLOutput;
203
204 async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
205 if input.code.trim().is_empty() {
206 return Err(ToolError::InvalidInput(
207 "Python code must not be empty".to_string(),
208 ));
209 }
210
211 if !self.dangerously_allow {
212 return Err(ToolError::ExecutionFailed(
213 "PythonREPLTool is disabled by default for security. \
214 Call .with_dangerously_allow(true) to enable execution."
215 .to_string(),
216 ));
217 }
218
219 if self.check_dangerous_imports {
220 if let Some(blocked) = contains_dangerous_code(&input.code) {
221 return Err(ToolError::ExecutionFailed(format!(
222 "Code contains dangerous import or builtin call: '{}'. \
223 Blocked by the noise-filter blacklist (note: this is not a security boundary; \
224 untrusted code must run in a sandbox). \
225 Call .with_skip_dangerous_imports_check(true) to bypass (not recommended).",
226 blocked
227 )));
228 }
229 }
230
231 let timeout_secs = input.timeout_seconds.unwrap_or(30);
232
233 let result = tokio::time::timeout(
234 std::time::Duration::from_secs(timeout_secs),
235 Command::new(&self.python_path)
236 .arg("-c")
237 .arg(&input.code)
238 .output(),
239 )
240 .await
241 .map_err(|_| {
242 ToolError::ExecutionFailed(format!(
243 "Python execution timed out after {} seconds",
244 timeout_secs
245 ))
246 })?
247 .map_err(|e| ToolError::ExecutionFailed(format!("Python execution failed: {}", e)))?;
248
249 let stdout = String::from_utf8_lossy(&result.stdout).to_string();
250 let stderr = String::from_utf8_lossy(&result.stderr).to_string();
251 let exit_code = result.status.code().unwrap_or(-1);
252
253 Ok(PythonREPLOutput {
254 code: input.code,
255 stdout,
256 stderr,
257 exit_code,
258 })
259 }
260}
261
262#[async_trait]
263impl BaseTool for PythonREPLTool {
264 fn name(&self) -> &str {
265 "python_repl"
266 }
267
268 fn description(&self) -> &str {
269 "Python code execution tool. Runs code in a local Python environment and returns results.
270
271Parameters:
272- code: Python code string to execute
273- timeout_seconds: Timeout in seconds (default: 30)
274
275Supports any Python syntax, including math, data processing, plotting, etc.
276
277SECURITY WARNING: Disabled by default. Must call .with_dangerously_allow(true) to enable.
278Only use in controlled/sandboxed environments.
279
280Examples:
281- Simple calc: {\"code\": \"print(1 + 2)\"}
282- List processing: {\"code\": \"print([x**2 for x in range(10)])\"}
283- Math: {\"code\": \"import math; print(math.pi)\"}"
284 }
285
286 async fn run(&self, input: String) -> Result<String, ToolError> {
287 let parsed: PythonREPLInput = serde_json::from_str(&input)
288 .map_err(|e| ToolError::InvalidInput(format!("JSON parse error: {}", e)))?;
289
290 let output = self.invoke(parsed).await?;
291
292 let mut result = String::new();
293 if !output.stdout.is_empty() {
294 result.push_str(&format!("stdout:\n{}\n", output.stdout));
295 }
296 if !output.stderr.is_empty() {
297 result.push_str(&format!("stderr:\n{}\n", output.stderr));
298 }
299 result.push_str(&format!("exit_code: {}", output.exit_code));
300
301 Ok(result)
302 }
303
304 fn args_schema(&self) -> Option<serde_json::Value> {
305 use schemars::schema_for;
306 serde_json::to_value(schema_for!(PythonREPLInput)).ok()
307 }
308}
309
310#[cfg(test)]
311mod tests {
312 use super::*;
313
314 #[test]
315 fn test_python_repl_tool_properties() {
316 let tool = PythonREPLTool::new();
317 assert_eq!(tool.name(), "python_repl");
318 assert!(tool.description().contains("Python"));
319 assert!(BaseTool::args_schema(&tool).is_some());
320 }
321
322 #[tokio::test]
323 async fn test_python_repl_empty_code() {
324 let tool = PythonREPLTool::new().with_dangerously_allow(true);
325 let result = tool.run(r#"{"code": ""}"#.to_string()).await;
326 assert!(result.is_err());
327 }
328
329 #[tokio::test]
330 async fn test_python_repl_disabled_by_default() {
331 let tool = PythonREPLTool::new();
332 let result = tool
333 .invoke(PythonREPLInput {
334 code: "print(1 + 2)".to_string(),
335 timeout_seconds: Some(10),
336 })
337 .await;
338 assert!(result.is_err());
339 let err_msg = result.unwrap_err().to_string();
340 assert!(
341 err_msg.contains("disabled by default"),
342 "Expected disabled error, got: {}",
343 err_msg
344 );
345 }
346
347 #[tokio::test]
348 async fn test_python_repl_basic_execution() {
349 let tool = PythonREPLTool::new().with_dangerously_allow(true);
350 let result = tool
351 .invoke(PythonREPLInput {
352 code: "print(1 + 2)".to_string(),
353 timeout_seconds: Some(10),
354 })
355 .await;
356
357 match result {
358 Ok(output) => {
359 if output.exit_code == 0 || !output.stdout.is_empty() {
360 } else {
362 eprintln!(
363 "Python may not be installed (exit_code={})",
364 output.exit_code
365 );
366 }
367 }
368 Err(e) => {
369 eprintln!("Python not available (may be expected): {}", e);
370 }
371 }
372 }
373
374 #[test]
375 fn test_dangerous_import_detection() {
376 assert!(contains_dangerous_code("import os").is_some());
377 assert!(contains_dangerous_code("import subprocess").is_some());
378 assert!(contains_dangerous_code("from sys import path").is_some());
379 assert!(contains_dangerous_code("from os.path import join").is_some());
380 assert!(contains_dangerous_code("import math").is_none());
382 assert!(contains_dangerous_code("import json").is_none());
383 assert!(contains_dangerous_code("from datetime import datetime").is_none());
384 assert!(contains_dangerous_code("# import os").is_none());
386 }
387
388 #[test]
389 fn test_dangerous_builtin_calls_detected() {
390 assert!(contains_dangerous_code("__import__('os').system('ls')").is_some());
392 assert!(contains_dangerous_code("importlib.import_module('os')").is_some());
393 assert!(contains_dangerous_code("eval('os')").is_some());
394 assert!(contains_dangerous_code("exec('import os')").is_some());
395 assert!(contains_dangerous_code("compile('import os', '<x>', 'exec')").is_some());
396 assert!(contains_dangerous_code("execfile('/tmp/x.py')").is_some());
397 }
398
399 #[test]
400 fn test_dangerous_builtin_calls_no_false_positive_on_words() {
401 assert!(contains_dangerous_code("print('evaluate the result')").is_none());
403 assert!(contains_dangerous_code("result = execute_query()").is_none());
404 assert!(contains_dangerous_code("print(len([1, 2, 3]))").is_none());
405 assert!(contains_dangerous_code("x = len('hello')").is_none());
406 }
407
408 #[tokio::test]
409 async fn test_python_repl_blocks_dangerous_import() {
410 let tool = PythonREPLTool::new().with_dangerously_allow(true);
411 let result = tool
412 .invoke(PythonREPLInput {
413 code: "import os; print(os.getcwd())".to_string(),
414 timeout_seconds: Some(10),
415 })
416 .await;
417 assert!(result.is_err());
418 let err_msg = result.unwrap_err().to_string();
419 assert!(
420 err_msg.contains("dangerous import"),
421 "Expected dangerous import error, got: {}",
422 err_msg
423 );
424 }
425
426 #[tokio::test]
427 async fn test_python_repl_allows_safe_import() {
428 let tool = PythonREPLTool::new().with_dangerously_allow(true);
429 let result = tool
430 .invoke(PythonREPLInput {
431 code: "import math; print(math.pi)".to_string(),
432 timeout_seconds: Some(10),
433 })
434 .await;
435 match result {
436 Ok(output) => {
437 if output.exit_code == 0 {
438 assert!(output.stdout.contains("3.14"));
439 }
440 }
441 Err(e) => {
442 assert!(
443 !e.to_string().contains("dangerous import"),
444 "math should not be blocked: {}",
445 e
446 );
447 }
448 }
449 }
450
451 #[tokio::test]
452 async fn test_python_repl_with_error() {
453 let tool = PythonREPLTool::new().with_dangerously_allow(true);
454 let result = tool
455 .invoke(PythonREPLInput {
456 code: "print(undefined_var)".to_string(),
457 timeout_seconds: Some(10),
458 })
459 .await;
460
461 match result {
462 Ok(output) => {
463 if output.exit_code == 0 {
464 }
466 }
467 Err(_) => {
468 }
470 }
471 }
472}