1use crate::registry::{Tool, ToolContext};
10use aegis_security::check_path;
11use anyhow::Result;
12use async_trait::async_trait;
13use serde_json::{json, Value};
14use std::path::Path;
15
16pub struct DiskUsageTool;
22
23impl DiskUsageTool {
24 pub fn new() -> Self {
26 DiskUsageTool
27 }
28}
29
30impl Default for DiskUsageTool {
31 fn default() -> Self {
32 Self::new()
33 }
34}
35
36#[async_trait]
37impl Tool for DiskUsageTool {
38 fn name(&self) -> &str {
39 "disk_usage"
40 }
41
42 fn description(&self) -> &str {
43 "Analyze disk usage of a directory (du-style): total size plus the largest top-level children. Read-only, within the working directory."
44 }
45
46 fn parameters(&self) -> Value {
47 json!({
48 "type": "object",
49 "properties": {
50 "path": { "type": "string", "description": "Directory to analyze (default: working directory)" },
51 "top": { "type": "integer", "description": "How many largest children to list (default 20)" }
52 }
53 })
54 }
55
56 async fn execute(&self, args: Value, ctx: &ToolContext<'_>) -> Result<String> {
57 let path_arg = args["path"].as_str().unwrap_or(".");
58 let top = args["top"].as_u64().unwrap_or(20) as usize;
59 let dir = check_path(path_arg, &ctx.cwd)?;
60 if !dir.exists() {
61 anyhow::bail!("Path not found: {path_arg}");
62 }
63 if !dir.is_dir() {
64 let size = std::fs::metadata(&dir).map(|m| m.len()).unwrap_or(0);
66 return Ok(format!("{} {}", human_bytes(size), path_arg));
67 }
68
69 let dir_owned = dir.clone();
70 let (total, mut children) = tokio::task::spawn_blocking(move || {
71 let mut children: Vec<(String, u64)> = Vec::new();
72 let mut total = 0u64;
73 if let Ok(rd) = std::fs::read_dir(&dir_owned) {
74 for entry in rd.flatten() {
75 let name = entry.file_name().to_string_lossy().to_string();
76 let sz = dir_size(&entry.path());
77 total += sz;
78 children.push((name, sz));
79 }
80 }
81 (total, children)
82 })
83 .await
84 .map_err(|e| anyhow::anyhow!("disk_usage task failed: {e}"))?;
85
86 children.sort_by(|a, b| b.1.cmp(&a.1));
87 children.truncate(top);
88
89 let mut out = format!("Total: {} ({})\n", human_bytes(total), path_arg);
90 for (name, sz) in children {
91 out.push_str(&format!("{:>10} {}\n", human_bytes(sz), name));
92 }
93 Ok(out.trim_end().to_string())
94 }
95}
96
97fn dir_size(path: &Path) -> u64 {
99 let meta = match std::fs::symlink_metadata(path) {
100 Ok(m) => m,
101 Err(_) => return 0,
102 };
103 if meta.file_type().is_symlink() {
104 return 0;
105 }
106 if meta.is_file() {
107 return meta.len();
108 }
109 if meta.is_dir() {
110 let mut total = 0u64;
111 if let Ok(rd) = std::fs::read_dir(path) {
112 for entry in rd.flatten() {
113 total += dir_size(&entry.path());
114 }
115 }
116 return total;
117 }
118 0
119}
120
121fn human_bytes(bytes: u64) -> String {
122 const UNITS: &[&str] = &["B", "KiB", "MiB", "GiB", "TiB"];
123 let mut v = bytes as f64;
124 let mut i = 0;
125 while v >= 1024.0 && i < UNITS.len() - 1 {
126 v /= 1024.0;
127 i += 1;
128 }
129 format!("{v:.1} {}", UNITS[i])
130}
131
132pub struct ListeningPortsTool;
138
139impl ListeningPortsTool {
140 pub fn new() -> Self {
142 ListeningPortsTool
143 }
144}
145
146impl Default for ListeningPortsTool {
147 fn default() -> Self {
148 Self::new()
149 }
150}
151
152#[async_trait]
153impl Tool for ListeningPortsTool {
154 fn name(&self) -> &str {
155 "listening_ports"
156 }
157
158 fn description(&self) -> &str {
159 "List TCP sockets in the LISTEN state (what's listening on this host) by reading /proc/net/tcp and /proc/net/tcp6. Linux only, read-only."
160 }
161
162 fn parameters(&self) -> Value {
163 json!({
164 "type": "object",
165 "properties": {}
166 })
167 }
168
169 async fn execute(&self, _args: Value, _ctx: &ToolContext<'_>) -> Result<String> {
170 let mut listeners: Vec<String> = Vec::new();
171 let mut any_file = false;
172
173 for (proc_path, is_v6) in [("/proc/net/tcp", false), ("/proc/net/tcp6", true)] {
174 if let Ok(content) = std::fs::read_to_string(proc_path) {
175 any_file = true;
176 for line in content.lines().skip(1) {
177 if let Some(addr) = parse_listen_line(line, is_v6) {
178 listeners.push(addr);
179 }
180 }
181 }
182 }
183
184 if !any_file {
185 return Ok("listening_ports is only supported on Linux (could not read /proc/net/tcp).".to_string());
186 }
187 if listeners.is_empty() {
188 return Ok("No listening TCP sockets found.".to_string());
189 }
190 listeners.sort();
191 listeners.dedup();
192 Ok(format!("Listening TCP sockets:\n{}", listeners.join("\n")))
193 }
194}
195
196fn parse_listen_line(line: &str, is_v6: bool) -> Option<String> {
199 let fields: Vec<&str> = line.split_whitespace().collect();
200 if fields.len() < 4 {
202 return None;
203 }
204 let st = fields[3];
205 if !st.eq_ignore_ascii_case("0A") {
206 return None; }
208 let (addr_hex, port_hex) = fields[1].split_once(':')?;
209 let port = u16::from_str_radix(port_hex, 16).ok()?;
210 let ip = if is_v6 {
211 parse_ipv6_hex(addr_hex)?
212 } else {
213 parse_ipv4_hex(addr_hex)?
214 };
215 if is_v6 {
216 Some(format!("[{ip}]:{port}"))
217 } else {
218 Some(format!("{ip}:{port}"))
219 }
220}
221
222fn parse_ipv4_hex(hex: &str) -> Option<String> {
224 if hex.len() != 8 {
225 return None;
226 }
227 let n = u32::from_str_radix(hex, 16).ok()?;
228 let b = n.to_le_bytes(); Some(std::net::Ipv4Addr::new(b[0], b[1], b[2], b[3]).to_string())
230}
231
232fn parse_ipv6_hex(hex: &str) -> Option<String> {
234 if hex.len() != 32 {
235 return None;
236 }
237 let mut bytes = [0u8; 16];
238 for word in 0..4 {
239 let chunk = &hex[word * 8..word * 8 + 8];
240 let n = u32::from_str_radix(chunk, 16).ok()?;
241 let le = n.to_le_bytes();
242 bytes[word * 4..word * 4 + 4].copy_from_slice(&le);
243 }
244 Some(std::net::Ipv6Addr::from(bytes).to_string())
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250
251 #[test]
252 fn ipv4_hex_is_little_endian() {
253 assert_eq!(parse_ipv4_hex("0100007F").as_deref(), Some("127.0.0.1"));
254 assert_eq!(parse_ipv4_hex("00000000").as_deref(), Some("0.0.0.0"));
255 assert_eq!(parse_ipv4_hex("bad"), None);
256 }
257
258 #[test]
259 fn listen_line_filters_state() {
260 let listen = " 0: 00000000:0016 00000000:0000 0A 00000000:00000000 00:00000000 00000000 0";
262 assert_eq!(parse_listen_line(listen, false).as_deref(), Some("0.0.0.0:22"));
263 let est = " 1: 0100007F:1F90 0100007F:C000 01 00000000:00000000 00:00000000 00000000 0";
265 assert_eq!(parse_listen_line(est, false), None);
266 }
267
268 #[test]
269 fn human_bytes_ok() {
270 assert_eq!(human_bytes(1536), "1.5 KiB");
271 }
272}