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 {
147 Self {
148 python_path: Self::find_python(),
149 dangerously_allow: false,
150 check_dangerous_imports: true,
151 }
152 }
153
154 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 pub fn with_dangerously_allow(mut self, allow: bool) -> Self {
165 self.dangerously_allow = allow;
166 self
167 }
168
169 pub fn with_skip_dangerous_imports_check(mut self, skip: bool) -> Self {
171 self.check_dangerous_imports = !skip;
172 self
173 }
174
175 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 } 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 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 assert!(contains_dangerous_code("# import os").is_none());
383 }
384
385 #[test]
386 fn test_dangerous_builtin_calls_detected() {
387 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 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 }
463 }
464 Err(_) => {
465 }
467 }
468 }
469}