agentzero_tools/
image_info.rs

1use agentzero_core::{Tool, ToolContext, ToolResult};
2use anyhow::{anyhow, Context};
3use async_trait::async_trait;
4use serde::Deserialize;
5use std::path::{Component, Path, PathBuf};
6
7#[derive(Debug, Deserialize)]
8struct ImageInfoInput {
9    path: String,
10}
11
12#[derive(Debug, Default, Clone, Copy)]
13pub struct ImageInfoTool;
14
15impl ImageInfoTool {
16    fn resolve_path(input_path: &str, workspace_root: &str) -> anyhow::Result<PathBuf> {
17        if input_path.trim().is_empty() {
18            return Err(anyhow!("path is required"));
19        }
20        let relative = Path::new(input_path);
21        if relative.is_absolute() {
22            return Err(anyhow!("absolute paths are not allowed"));
23        }
24        if relative
25            .components()
26            .any(|c| matches!(c, Component::ParentDir))
27        {
28            return Err(anyhow!("path traversal is not allowed"));
29        }
30        let joined = Path::new(workspace_root).join(relative);
31        let canonical_root = Path::new(workspace_root)
32            .canonicalize()
33            .context("unable to resolve workspace root")?;
34        let canonical = joined
35            .canonicalize()
36            .with_context(|| format!("file not found: {input_path}"))?;
37        if !canonical.starts_with(&canonical_root) {
38            return Err(anyhow!("path is outside workspace"));
39        }
40        Ok(canonical)
41    }
42
43    fn detect_format(header: &[u8]) -> &'static str {
44        if header.starts_with(b"\x89PNG") {
45            "PNG"
46        } else if header.starts_with(b"\xFF\xD8\xFF") {
47            "JPEG"
48        } else if header.starts_with(b"GIF8") {
49            "GIF"
50        } else if header.starts_with(b"RIFF") && header.len() >= 12 && &header[8..12] == b"WEBP" {
51            "WebP"
52        } else if header.starts_with(b"BM") {
53            "BMP"
54        } else {
55            "unknown"
56        }
57    }
58
59    fn png_dimensions(data: &[u8]) -> Option<(u32, u32)> {
60        if data.len() < 24 {
61            return None;
62        }
63        let width = u32::from_be_bytes([data[16], data[17], data[18], data[19]]);
64        let height = u32::from_be_bytes([data[20], data[21], data[22], data[23]]);
65        Some((width, height))
66    }
67
68    fn jpeg_dimensions(data: &[u8]) -> Option<(u32, u32)> {
69        let mut i = 2;
70        while i + 9 < data.len() {
71            if data[i] != 0xFF {
72                return None;
73            }
74            let marker = data[i + 1];
75            if marker == 0xC0 || marker == 0xC2 {
76                let height = u16::from_be_bytes([data[i + 5], data[i + 6]]) as u32;
77                let width = u16::from_be_bytes([data[i + 7], data[i + 8]]) as u32;
78                return Some((width, height));
79            }
80            let seg_len = u16::from_be_bytes([data[i + 2], data[i + 3]]) as usize;
81            i += 2 + seg_len;
82        }
83        None
84    }
85
86    fn gif_dimensions(data: &[u8]) -> Option<(u32, u32)> {
87        if data.len() < 10 {
88            return None;
89        }
90        let width = u16::from_le_bytes([data[6], data[7]]) as u32;
91        let height = u16::from_le_bytes([data[8], data[9]]) as u32;
92        Some((width, height))
93    }
94}
95
96#[async_trait]
97impl Tool for ImageInfoTool {
98    fn name(&self) -> &'static str {
99        "image_info"
100    }
101
102    fn description(&self) -> &'static str {
103        "Get metadata about an image file: format, dimensions, and file size."
104    }
105
106    fn input_schema(&self) -> Option<serde_json::Value> {
107        Some(serde_json::json!({
108            "type": "object",
109            "properties": {
110                "path": { "type": "string", "description": "Path to the image file" }
111            },
112            "required": ["path"]
113        }))
114    }
115
116    async fn execute(&self, input: &str, ctx: &ToolContext) -> anyhow::Result<ToolResult> {
117        let req: ImageInfoInput =
118            serde_json::from_str(input).context("image_info expects JSON: {\"path\": \"...\"}")?;
119
120        let file_path = Self::resolve_path(&req.path, &ctx.workspace_root)?;
121        let data = tokio::fs::read(&file_path)
122            .await
123            .with_context(|| format!("failed to read file: {}", req.path))?;
124
125        let format = Self::detect_format(&data);
126        let dimensions = match format {
127            "PNG" => Self::png_dimensions(&data),
128            "JPEG" => Self::jpeg_dimensions(&data),
129            "GIF" => Self::gif_dimensions(&data),
130            _ => None,
131        };
132
133        let file_size = data.len();
134        let mut output = format!("format={format}\nsize={file_size} bytes");
135        if let Some((w, h)) = dimensions {
136            output.push_str(&format!("\nwidth={w}\nheight={h}"));
137        }
138
139        Ok(ToolResult { output })
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use std::fs;
147    use std::sync::atomic::{AtomicU64, Ordering};
148    use std::time::{SystemTime, UNIX_EPOCH};
149
150    static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
151
152    fn temp_dir() -> PathBuf {
153        let nanos = SystemTime::now()
154            .duration_since(UNIX_EPOCH)
155            .expect("clock")
156            .as_nanos();
157        let seq = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
158        let dir = std::env::temp_dir().join(format!(
159            "agentzero-image-info-{}-{nanos}-{seq}",
160            std::process::id()
161        ));
162        fs::create_dir_all(&dir).expect("temp dir should be created");
163        dir
164    }
165
166    #[tokio::test]
167    async fn image_info_detects_png() {
168        let dir = temp_dir();
169        // Minimal valid PNG header (8-byte magic + 13-byte IHDR chunk with 1x1 dims)
170        let mut png = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
171        // IHDR length (13)
172        png.extend_from_slice(&[0x00, 0x00, 0x00, 0x0D]);
173        // "IHDR"
174        png.extend_from_slice(b"IHDR");
175        // width=100, height=50
176        png.extend_from_slice(&100u32.to_be_bytes());
177        png.extend_from_slice(&50u32.to_be_bytes());
178        // bit depth, color type, compression, filter, interlace
179        png.extend_from_slice(&[8, 2, 0, 0, 0]);
180        fs::write(dir.join("test.png"), &png).unwrap();
181
182        let tool = ImageInfoTool;
183        let result = tool
184            .execute(
185                r#"{"path": "test.png"}"#,
186                &ToolContext::new(dir.to_string_lossy().to_string()),
187            )
188            .await
189            .expect("should succeed");
190        assert!(result.output.contains("format=PNG"));
191        assert!(result.output.contains("width=100"));
192        assert!(result.output.contains("height=50"));
193        fs::remove_dir_all(dir).ok();
194    }
195
196    #[tokio::test]
197    async fn image_info_rejects_path_traversal() {
198        let dir = temp_dir();
199        let tool = ImageInfoTool;
200        let err = tool
201            .execute(
202                r#"{"path": "../escape.png"}"#,
203                &ToolContext::new(dir.to_string_lossy().to_string()),
204            )
205            .await
206            .expect_err("path traversal should fail");
207        assert!(err.to_string().contains("path traversal"));
208        fs::remove_dir_all(dir).ok();
209    }
210
211    #[tokio::test]
212    async fn image_info_non_image_file() {
213        let dir = temp_dir();
214        fs::write(dir.join("test.txt"), "hello world").unwrap();
215        let tool = ImageInfoTool;
216        let result = tool
217            .execute(
218                r#"{"path": "test.txt"}"#,
219                &ToolContext::new(dir.to_string_lossy().to_string()),
220            )
221            .await
222            .expect("should succeed even for non-image");
223        assert!(result.output.contains("format=unknown"));
224        fs::remove_dir_all(dir).ok();
225    }
226
227    #[test]
228    fn detect_format_from_headers() {
229        assert_eq!(ImageInfoTool::detect_format(b"\x89PNG\r\n\x1a\n"), "PNG");
230        assert_eq!(ImageInfoTool::detect_format(b"\xFF\xD8\xFF\xE0"), "JPEG");
231        assert_eq!(ImageInfoTool::detect_format(b"GIF89a"), "GIF");
232        assert_eq!(ImageInfoTool::detect_format(b"hello"), "unknown");
233    }
234}
agentzero_tools/image_info.rs

agentzero_tools/
image_info.rs
