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(
238 std::time::Duration::from_secs(timeout_secs),
239 Command::new(&self.python_path)
240 .arg("-c")
241 .arg(&input.code)
242 .kill_on_drop(true)
243 .output(),
244 )
245 .await
246 .map_err(|_| {
247 ToolError::ExecutionFailed(format!(
248 "Python execution timed out after {} seconds",
249 timeout_secs
250 ))
251 })?
252 .map_err(|e| ToolError::ExecutionFailed(format!("Python execution failed: {}", e)))?;
253
254 let stdout = String::from_utf8_lossy(&result.stdout).to_string();
255 let stderr = String::from_utf8_lossy(&result.stderr).to_string();
256 let exit_code = result.status.code().unwrap_or(-1);
257
258 Ok(PythonREPLOutput {
259 code: input.code,
260 stdout,
261 stderr,
262 exit_code,
263 })
264 }
265}
266
267#[async_trait]
268impl BaseTool for PythonREPLTool {
269 fn name(&self) -> &str {
270 "python_repl"
271 }
272
273 fn description(&self) -> &str {
274 "Python code execution tool. Runs code in a local Python environment and returns results.
275
276Parameters:
277- code: Python code string to execute
278- timeout_seconds: Timeout in seconds (default: 30)
279
280Supports any Python syntax, including math, data processing, plotting, etc.
281
282SECURITY WARNING: Disabled by default. Must call .with_dangerously_allow(true) to enable.
283Only use in controlled/sandboxed environments.
284
285Examples:
286- Simple calc: {\"code\": \"print(1 + 2)\"}
287- List processing: {\"code\": \"print([x**2 for x in range(10)])\"}
288- Math: {\"code\": \"import math; print(math.pi)\"}"
289 }
290
291 async fn run(&self, input: String) -> Result<String, ToolError> {
292 let parsed: PythonREPLInput = serde_json::from_str(&input)
293 .map_err(|e| ToolError::InvalidInput(format!("JSON parse error: {}", e)))?;
294
295 let output = self.invoke(parsed).await?;
296
297 let mut result = String::new();
298 if !output.stdout.is_empty() {
299 result.push_str(&format!("stdout:\n{}\n", output.stdout));
300 }
301 if !output.stderr.is_empty() {
302 result.push_str(&format!("stderr:\n{}\n", output.stderr));
303 }
304 result.push_str(&format!("exit_code: {}", output.exit_code));
305
306 Ok(result)
307 }
308
309 fn args_schema(&self) -> Option<serde_json::Value> {
310 use schemars::schema_for;
311 serde_json::to_value(schema_for!(PythonREPLInput)).ok()
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use super::*;
318
319 #[test]
320 fn test_python_repl_tool_properties() {
321 let tool = PythonREPLTool::new();
322 assert_eq!(tool.name(), "python_repl");
323 assert!(tool.description().contains("Python"));
324 assert!(BaseTool::args_schema(&tool).is_some());
325 }
326
327 #[tokio::test]
328 async fn test_python_repl_empty_code() {
329 let tool = PythonREPLTool::new().with_dangerously_allow(true);
330 let result = tool.run(r#"{"code": ""}"#.to_string()).await;
331 assert!(result.is_err());
332 }
333
334 #[tokio::test]
335 async fn test_python_repl_disabled_by_default() {
336 let tool = PythonREPLTool::new();
337 let result = tool
338 .invoke(PythonREPLInput {
339 code: "print(1 + 2)".to_string(),
340 timeout_seconds: Some(10),
341 })
342 .await;
343 assert!(result.is_err());
344 let err_msg = result.unwrap_err().to_string();
345 assert!(
346 err_msg.contains("disabled by default"),
347 "Expected disabled error, got: {}",
348 err_msg
349 );
350 }
351
352 #[tokio::test]
353 async fn test_python_repl_basic_execution() {
354 let tool = PythonREPLTool::new().with_dangerously_allow(true);
355 let result = tool
356 .invoke(PythonREPLInput {
357 code: "print(1 + 2)".to_string(),
358 timeout_seconds: Some(10),
359 })
360 .await;
361
362 match result {
363 Ok(output) => {
364 if output.exit_code == 0 || !output.stdout.is_empty() {
365 } else {
367 eprintln!(
368 "Python may not be installed (exit_code={})",
369 output.exit_code
370 );
371 }
372 }
373 Err(e) => {
374 eprintln!("Python not available (may be expected): {}", e);
375 }
376 }
377 }
378
379 #[test]
380 fn test_dangerous_import_detection() {
381 assert!(contains_dangerous_code("import os").is_some());
382 assert!(contains_dangerous_code("import subprocess").is_some());
383 assert!(contains_dangerous_code("from sys import path").is_some());
384 assert!(contains_dangerous_code("from os.path import join").is_some());
385 assert!(contains_dangerous_code("import math").is_none());
387 assert!(contains_dangerous_code("import json").is_none());
388 assert!(contains_dangerous_code("from datetime import datetime").is_none());
389 assert!(contains_dangerous_code("# import os").is_none());
391 }
392
393 #[test]
394 fn test_dangerous_builtin_calls_detected() {
395 assert!(contains_dangerous_code("__import__('os').system('ls')").is_some());
397 assert!(contains_dangerous_code("importlib.import_module('os')").is_some());
398 assert!(contains_dangerous_code("eval('os')").is_some());
399 assert!(contains_dangerous_code("exec('import os')").is_some());
400 assert!(contains_dangerous_code("compile('import os', '<x>', 'exec')").is_some());
401 assert!(contains_dangerous_code("execfile('/tmp/x.py')").is_some());
402 }
403
404 #[test]
405 fn test_dangerous_builtin_calls_no_false_positive_on_words() {
406 assert!(contains_dangerous_code("print('evaluate the result')").is_none());
408 assert!(contains_dangerous_code("result = execute_query()").is_none());
409 assert!(contains_dangerous_code("print(len([1, 2, 3]))").is_none());
410 assert!(contains_dangerous_code("x = len('hello')").is_none());
411 }
412
413 #[tokio::test]
414 async fn test_python_repl_blocks_dangerous_import() {
415 let tool = PythonREPLTool::new().with_dangerously_allow(true);
416 let result = tool
417 .invoke(PythonREPLInput {
418 code: "import os; print(os.getcwd())".to_string(),
419 timeout_seconds: Some(10),
420 })
421 .await;
422 assert!(result.is_err());
423 let err_msg = result.unwrap_err().to_string();
424 assert!(
425 err_msg.contains("dangerous import"),
426 "Expected dangerous import error, got: {}",
427 err_msg
428 );
429 }
430
431 #[tokio::test]
432 async fn test_python_repl_allows_safe_import() {
433 let tool = PythonREPLTool::new().with_dangerously_allow(true);
434 let result = tool
435 .invoke(PythonREPLInput {
436 code: "import math; print(math.pi)".to_string(),
437 timeout_seconds: Some(10),
438 })
439 .await;
440 match result {
441 Ok(output) => {
442 if output.exit_code == 0 {
443 assert!(output.stdout.contains("3.14"));
444 }
445 }
446 Err(e) => {
447 assert!(
448 !e.to_string().contains("dangerous import"),
449 "math should not be blocked: {}",
450 e
451 );
452 }
453 }
454 }
455
456 #[tokio::test]
457 async fn test_python_repl_with_error() {
458 let tool = PythonREPLTool::new().with_dangerously_allow(true);
459 let result = tool
460 .invoke(PythonREPLInput {
461 code: "print(undefined_var)".to_string(),
462 timeout_seconds: Some(10),
463 })
464 .await;
465
466 match result {
467 Ok(output) => {
468 if output.exit_code == 0 {
469 }
471 }
472 Err(_) => {
473 }
475 }
476 }
477}