agentzero_tools/
screenshot.rs

1use agentzero_core::{Tool, ToolContext, ToolResult};
2use anyhow::{anyhow, Context};
3use async_trait::async_trait;
4use serde::Deserialize;
5use std::path::PathBuf;
6use std::process::Stdio;
7use tokio::process::Command;
8
9#[derive(Debug, Deserialize)]
10struct ScreenshotInput {
11    #[serde(default = "default_filename")]
12    filename: String,
13}
14
15fn default_filename() -> String {
16    "screenshot.png".to_string()
17}
18
19#[derive(Debug, Default, Clone, Copy)]
20pub struct ScreenshotTool;
21
22impl ScreenshotTool {
23    fn screenshot_command() -> &'static str {
24        if cfg!(target_os = "macos") {
25            "screencapture"
26        } else {
27            "import"
28        }
29    }
30
31    fn screenshot_args(output_path: &str) -> Vec<String> {
32        if cfg!(target_os = "macos") {
33            vec!["-x".to_string(), output_path.to_string()]
34        } else {
35            // ImageMagick's import tool
36            vec![
37                "-window".to_string(),
38                "root".to_string(),
39                output_path.to_string(),
40            ]
41        }
42    }
43}
44
45#[async_trait]
46impl Tool for ScreenshotTool {
47    fn name(&self) -> &'static str {
48        "screenshot"
49    }
50
51    fn description(&self) -> &'static str {
52        "Capture a screenshot of the current display and save it to a file."
53    }
54
55    fn input_schema(&self) -> Option<serde_json::Value> {
56        Some(serde_json::json!({
57            "type": "object",
58            "properties": {
59                "filename": { "type": "string", "description": "Output filename for the screenshot" }
60            },
61            "required": ["filename"]
62        }))
63    }
64
65    async fn execute(&self, input: &str, ctx: &ToolContext) -> anyhow::Result<ToolResult> {
66        let req: ScreenshotInput = serde_json::from_str(input)
67            .context("screenshot expects JSON: {\"filename\": \"...\"}")?;
68
69        if req.filename.contains("..") || req.filename.starts_with('/') {
70            return Err(anyhow!("invalid filename"));
71        }
72
73        let output_path = PathBuf::from(&ctx.workspace_root).join(&req.filename);
74        let output_str = output_path.to_string_lossy().to_string();
75
76        let cmd = Self::screenshot_command();
77        let args = Self::screenshot_args(&output_str);
78
79        let status = Command::new(cmd)
80            .args(&args)
81            .stdout(Stdio::null())
82            .stderr(Stdio::null())
83            .status()
84            .await
85            .with_context(|| format!("failed to run {cmd}"))?;
86
87        if status.success() {
88            Ok(ToolResult {
89                output: format!("screenshot saved to {}", req.filename),
90            })
91        } else {
92            Err(anyhow!(
93                "screenshot failed with exit code {}",
94                status.code().unwrap_or(-1)
95            ))
96        }
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[tokio::test]
105    async fn screenshot_rejects_path_traversal() {
106        let tool = ScreenshotTool;
107        let err = tool
108            .execute(
109                r#"{"filename": "../escape.png"}"#,
110                &ToolContext::new(".".to_string()),
111            )
112            .await
113            .expect_err("path traversal should fail");
114        assert!(err.to_string().contains("invalid filename"));
115    }
116
117    #[tokio::test]
118    async fn screenshot_rejects_absolute_path() {
119        let tool = ScreenshotTool;
120        let err = tool
121            .execute(
122                r#"{"filename": "/tmp/screenshot.png"}"#,
123                &ToolContext::new(".".to_string()),
124            )
125            .await
126            .expect_err("absolute path should fail");
127        assert!(err.to_string().contains("invalid filename"));
128    }
129
130    #[test]
131    fn screenshot_command_returns_platform_binary() {
132        let cmd = ScreenshotTool::screenshot_command();
133        assert!(!cmd.is_empty());
134    }
135}
agentzero_tools/screenshot.rs

agentzero_tools/
screenshot.rs
