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