1use std::collections::BTreeMap;
13use std::ffi::OsString;
14use std::fs;
15use std::path::{Path, PathBuf};
16use std::process::Command;
17use std::sync::Arc;
18use std::time::{Duration, SystemTime};
19
20use byob_sandbox::{Limits, Policy, RunOptions, Sandbox, TempDir};
21use byob_tools::{Access, FileChange, Tool, ToolContext, ToolOutput, ToolSpec};
22use serde_json::Value;
23
24use crate::PluginRef;
25use crate::discover::Plugin;
26use crate::manifest::{ManifestTool, ToolKind, path_inside};
27use crate::names::{ToolNames, namespaced};
28
29pub const SCRIPT_TIMEOUT: Duration = Duration::from_secs(60);
31pub const MAX_SCRIPT_OUTPUT: usize = 256 * 1024;
33const MAX_SCANNED_FILES: usize = 5000;
35
36pub fn default_python() -> Vec<OsString> {
38 vec![OsString::from(if cfg!(windows) { "python" } else { "python3" })]
39}
40
41#[derive(Debug, Clone)]
43pub struct ScriptRun {
44 pub program: Vec<OsString>,
46 pub script: PathBuf,
47 pub args: Vec<OsString>,
49 pub cwd: PathBuf,
50 pub stdin: Vec<u8>,
51 pub env: BTreeMap<String, String>,
53 pub timeout: Duration,
54 pub max_output: usize,
56}
57
58#[derive(Debug, Clone, Default, PartialEq, Eq)]
59pub struct ScriptOutcome {
60 pub stdout: String,
61 pub stderr: String,
62 pub exit_code: Option<i32>,
64 pub timed_out: bool,
65 pub stopped: bool,
66 pub truncated: bool,
68}
69
70impl ScriptOutcome {
71 pub fn success(&self) -> bool {
72 self.exit_code == Some(0) && !self.timed_out && !self.stopped
73 }
74}
75
76pub fn run_script_unsandboxed(run: &ScriptRun, cancelled: &dyn Fn() -> bool) -> Result<ScriptOutcome, String> {
79 let Some((program, lead)) = run.program.split_first() else { return Err("no interpreter configured".into()) };
80 let mut cmd = Command::new(program);
81 cmd.args(lead).arg(&run.script).args(&run.args).current_dir(&run.cwd).envs(&run.env);
82 execute(cmd, run, cancelled)
83}
84
85pub fn run_script_sandboxed(
89 run: &ScriptRun,
90 sandbox: &Sandbox,
91 plugin_dir: &Path,
92 network: bool,
93 cancelled: &dyn Fn() -> bool,
94) -> Result<ScriptOutcome, String> {
95 let Some((program, lead)) = run.program.split_first() else { return Err("no interpreter configured".into()) };
96 let cwd = run.cwd.canonicalize().map_err(|e| format!("{}: {e}", run.cwd.display()))?;
97 let tmp = TempDir::new("byob-plugin-")?;
98 let mut policy = Policy::for_workspace(&cwd, tmp.path())
99 .network(network)
100 .limits(Limits::for_timeout(run.timeout.as_secs()))
101 .allow_read(plugin_dir)
102 .allow_path_dirs(std::env::var_os("PATH").as_deref())
103 .allow_unsandboxed(true);
105 for p in byob_sandbox::interpreter_read_paths(program) {
106 policy = policy.allow_read(p);
107 }
108 let mut args: Vec<OsString> = lead.to_vec();
109 args.push(run.script.clone().into_os_string());
110 args.extend(run.args.iter().cloned());
111 let mut cmd = sandbox.command(&policy, program, &args);
112 cmd.current_dir(&cwd);
113 byob_sandbox::clean_env(&mut cmd, tmp.path());
114 cmd.envs(&run.env);
115 let out = execute(cmd, run, cancelled)?;
116 drop(tmp);
117 Ok(out)
118}
119
120fn execute(cmd: Command, run: &ScriptRun, cancelled: &dyn Fn() -> bool) -> Result<ScriptOutcome, String> {
121 let opts = RunOptions { stdin: run.stdin.clone(), timeout: run.timeout, max_output: run.max_output };
122 let o = byob_sandbox::run(cmd, &opts, cancelled)?;
123 if let Some(e) = o.sandbox_error() {
124 return Err(format!("the sandbox couldn't start: {e}"));
125 }
126 Ok(ScriptOutcome {
127 stdout: String::from_utf8_lossy(&o.stdout.bytes()).into_owned(),
128 stderr: String::from_utf8_lossy(&o.stderr.bytes()).into_owned(),
129 exit_code: o.exit_code,
130 timed_out: o.timed_out,
131 stopped: o.cancelled,
132 truncated: o.stdout.truncated() || o.stderr.truncated(),
133 })
134}
135
136pub fn fill_title(template: &str, args: &Value) -> String {
138 let mut out = String::new();
139 let mut rest = template;
140 while let Some(open) = rest.find('{') {
141 out.push_str(&rest[..open]);
142 let after = &rest[open + 1..];
143 let key_len = after.find(|c: char| !(c.is_alphanumeric() || c == '_')).unwrap_or(after.len());
144 if key_len > 0 && after[key_len..].starts_with('}') {
145 let key = &after[..key_len];
146 match args.get(key) {
147 Some(Value::String(s)) => out.push_str(s),
148 Some(v) => out.push_str(&v.to_string()),
149 None => out.push_str(&rest[open..open + key_len + 2]),
150 }
151 rest = &after[key_len + 1..];
152 } else {
153 out.push('{');
154 rest = after;
155 }
156 }
157 out.push_str(rest);
158 out.chars().take(120).collect()
159}
160
161fn snapshot(dir: &Path) -> BTreeMap<PathBuf, Option<SystemTime>> {
163 let mut m = BTreeMap::new();
164 let mut stack = vec![(dir.to_path_buf(), 0u8)];
165 while let Some((d, depth)) = stack.pop() {
166 let Ok(rd) = fs::read_dir(&d) else { continue };
167 for e in rd.flatten() {
168 if m.len() >= MAX_SCANNED_FILES {
169 return m;
170 }
171 if e.file_name().to_string_lossy().starts_with('.') {
172 continue;
173 }
174 let Ok(meta) = e.metadata() else { continue };
175 if meta.is_dir() {
176 if depth < 2 {
177 stack.push((e.path(), depth + 1));
178 }
179 } else {
180 m.insert(e.path(), meta.modified().ok());
181 }
182 }
183 }
184 m
185}
186
187pub struct PythonTool {
189 name: String,
190 plugin: PluginRef,
191 decl: ManifestTool,
192 script: String,
193 interpreter: Vec<OsString>,
194 sandbox: Option<Sandbox>,
195 timeout: Duration,
196}
197
198impl PythonTool {
199 pub fn new(name: String, plugin: PluginRef, decl: ManifestTool, interpreter: Vec<OsString>) -> Option<Self> {
201 let ToolKind::Script(script) = &decl.kind else { return None };
202 let script = script.clone();
203 Some(Self { name, plugin, decl, script, interpreter, sandbox: None, timeout: SCRIPT_TIMEOUT })
204 }
205
206 pub fn with_sandbox(mut self, sandbox: Option<Sandbox>) -> Self {
208 self.sandbox = sandbox;
209 self
210 }
211
212 pub fn with_timeout(mut self, timeout: Duration) -> Self {
213 self.timeout = timeout;
214 self
215 }
216
217 pub fn plugin(&self) -> &PluginRef {
218 &self.plugin
219 }
220
221 fn finish(&self, r: &ScriptOutcome, workspace: &Path, before: &BTreeMap<PathBuf, Option<SystemTime>>) -> ToolOutput {
222 let name = &self.plugin.name;
223 if r.stopped {
224 return ToolOutput::err("Stopped by the user.");
225 }
226 if r.timed_out {
227 return ToolOutput::err(format!("{name} took longer than {} seconds and was stopped.", self.timeout.as_secs()));
228 }
229 if r.exit_code != Some(0) {
230 let output = [r.stdout.trim(), r.stderr.trim()].iter().filter(|s| !s.is_empty()).copied().collect::<Vec<_>>().join("\n");
231 return ToolOutput::err(format!("{name} ran into a problem:\n{output}"));
232 }
233 let out = r.stdout.trim();
234 let mut text = out.to_string();
235 let mut listed: Vec<PathBuf> = Vec::new();
236 if out.starts_with('{')
237 && let Ok(Value::Object(j)) = serde_json::from_str::<Value>(out)
238 {
239 if let Some(t) = j.get("text").and_then(Value::as_str) {
240 text = t.to_string();
241 }
242 if let Some(files) = j.get("files").and_then(Value::as_array) {
243 listed = files.iter().filter_map(Value::as_str).filter(|f| path_inside(f)).map(|f| workspace.join(f)).filter(|p| p.is_file()).collect();
244 }
245 }
246 if r.truncated {
247 text.push_str("\n…(output cut: it was too long)");
248 }
249 let after = snapshot(workspace);
250 let mut changes: Vec<FileChange> = after
251 .iter()
252 .filter(|(f, m)| before.get(*f) != Some(*m))
253 .map(|(f, _)| if before.contains_key(f) { FileChange::Modified(f.clone()) } else { FileChange::Created(f.clone()) })
254 .collect();
255 for f in listed {
256 if !changes.iter().any(|c| matches!(c, FileChange::Created(p) | FileChange::Modified(p) if *p == f)) {
257 changes.push(FileChange::Modified(f));
258 }
259 }
260 if text.is_empty() {
261 text = "(no result)".into();
262 }
263 ToolOutput { text, failed: false, changes }
264 }
265}
266
267impl Tool for PythonTool {
268 fn spec(&self) -> ToolSpec {
269 ToolSpec { name: self.name.clone(), description: self.decl.description.clone(), parameters: self.decl.params.clone() }
270 }
271
272 fn title(&self, args: &Value) -> String {
273 match &self.decl.title {
274 Some(t) => fill_title(t, args),
275 None => format!("Used {}", self.plugin.name),
276 }
277 }
278
279 fn access(&self, _args: &Value) -> Access {
280 Access::Execute
281 }
282
283 fn run(&self, args: &Value, ctx: &ToolContext<'_>) -> ToolOutput {
284 if ctx.cancelled() {
285 return ToolOutput::err("Stopped by the user.");
286 }
287 let script = self.plugin.dir.join(&self.script);
288 if !path_inside(&self.script) || !script.is_file() {
289 return ToolOutput::err(format!("{}: {} is missing", self.plugin.name, self.script));
290 }
291 let args = if args.is_object() { args.clone() } else { Value::Object(serde_json::Map::new()) };
292 let chat = ctx.workspace.to_string_lossy().into_owned();
293 let dir = self.plugin.dir.to_string_lossy().into_owned();
294 let mut env = self.plugin.env.clone();
295 for (k, v) in [("HOB_CHAT_DIR", &chat), ("BYOB_CHAT_DIR", &chat), ("HOB_PLUGIN_DIR", &dir), ("BYOB_PLUGIN_DIR", &dir)] {
296 env.insert(k.into(), v.clone());
297 }
298 let run = ScriptRun {
299 program: self.interpreter.clone(),
300 script,
301 args: Vec::new(),
302 cwd: ctx.workspace.to_path_buf(),
303 stdin: serde_json::to_vec(&args).unwrap_or_else(|_| b"{}".to_vec()),
304 env,
305 timeout: self.timeout,
306 max_output: MAX_SCRIPT_OUTPUT,
307 };
308 let before = snapshot(ctx.workspace);
309 let result = match &self.sandbox {
310 Some(sb) => run_script_sandboxed(&run, sb, &self.plugin.dir, self.plugin.network, &|| ctx.cancelled()),
311 None => run_script_unsandboxed(&run, &|| ctx.cancelled()),
312 };
313 match result {
314 Ok(r) => self.finish(&r, ctx.workspace, &before),
315 Err(e) => ToolOutput::err(format!("{} couldn't run: {e}", self.plugin.name)),
316 }
317 }
318}
319
320pub fn python_tools(plugin: &Plugin, settings: &BTreeMap<String, String>, interpreter: &[OsString], names: &mut ToolNames) -> Vec<Arc<PythonTool>> {
324 python_tools_in(plugin, settings, interpreter, names, None)
325}
326
327pub fn python_tools_in(
329 plugin: &Plugin,
330 settings: &BTreeMap<String, String>,
331 interpreter: &[OsString],
332 names: &mut ToolNames,
333 sandbox: Option<&Sandbox>,
334) -> Vec<Arc<PythonTool>> {
335 let r = PluginRef::new(plugin, settings);
336 plugin
337 .manifest
338 .tools
339 .iter()
340 .filter(|t| matches!(&t.kind, ToolKind::Script(s) if plugin.dir.join(s).is_file()))
341 .filter_map(|t| {
342 let name = names.claim(&t.name, &namespaced(plugin.id(), &t.name));
343 PythonTool::new(name, r.clone(), t.clone(), interpreter.to_vec()).map(|t| Arc::new(t.with_sandbox(sandbox.cloned())))
344 })
345 .collect()
346}
347
348#[cfg(test)]
349mod tests {
350 use super::*;
351 use serde_json::json;
352
353 #[test]
354 fn titles() {
355 let a = json!({"amount": 5, "from": "EUR", "x": null});
356 assert_eq!(fill_title("Converted {amount} {from} {to} {x}", &a), "Converted 5 EUR {to} null");
357 assert_eq!(fill_title("{ not a key } {", &a), "{ not a key } {");
358 assert_eq!(fill_title(&"é".repeat(200), &a).chars().count(), 120);
359 }
360
361 #[test]
362 fn missing_interpreter_is_an_error() {
363 let run = ScriptRun {
364 program: vec![OsString::from("/definitely/not/an/interpreter")],
365 script: PathBuf::from("x"),
366 args: vec![],
367 cwd: std::env::temp_dir(),
368 stdin: vec![],
369 env: BTreeMap::new(),
370 timeout: SCRIPT_TIMEOUT,
371 max_output: 10,
372 };
373 assert!(run_script_unsandboxed(&run, &|| false).unwrap_err().starts_with("couldn't start"));
374 let empty = ScriptRun { program: vec![], ..run };
375 assert!(run_script_unsandboxed(&empty, &|| false).is_err());
376 }
377}