limit_cli/tools/
bash.rs

1use async_trait::async_trait;
2use limit_agent::error::AgentError;
3use limit_agent::Tool;
4use serde_json::Value;
5use std::path::Path;
6use std::time::Duration;
7use tokio::process::Command;
8use tokio::time::timeout;
9
10pub struct BashTool;
11
12impl BashTool {
13    pub fn new() -> Self {
14        BashTool
15    }
16
17    const DEFAULT_TIMEOUT_SECS: u64 = 60;
18
19    /// Check if a command is dangerous and should be blocked
20    fn is_dangerous_command(command: &str) -> bool {
21        let lower_cmd = command.to_lowercase();
22
23        // Block specific dangerous patterns
24        let dangerous_patterns = [
25            "rm -rf /",        // Remove root directory
26            "rm -rf /*",       // Remove all files
27            ":(){ :|:& };:",   // Fork bomb
28            "dd if=/dev/zero", // Disk wipe
29            "mkfs.",           // Format filesystem
30            "mv / /dev/null",  // Move root to null
31            "chmod -R 777 /",  // Recursive chmod on root
32            "chown -R",        // Recursive chown on root
33            "killall -9",      // Kill all processes (with no filter)
34            "kill -9 -1",      // Kill all processes
35            "shred",           // Secure delete
36            "> /dev/sda",      // Direct write to disk
37            "wget http://",    // Random downloads
38            "curl http://",    // Random downloads
39        ];
40
41        dangerous_patterns
42            .iter()
43            .any(|pattern| lower_cmd.contains(pattern))
44    }
45}
46
47impl Default for BashTool {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53#[async_trait]
54impl Tool for BashTool {
55    fn name(&self) -> &str {
56        "bash"
57    }
58
59    async fn execute(&self, args: Value) -> Result<Value, AgentError> {
60        // Extract command from arguments
61        let command = args
62            .get("command")
63            .and_then(|v| v.as_str())
64            .ok_or_else(|| AgentError::ToolError("Missing 'command' argument".to_string()))?;
65
66        // Check for dangerous commands
67        if Self::is_dangerous_command(command) {
68            return Err(AgentError::ToolError(format!(
69                "Dangerous command blocked: {}",
70                command
71            )));
72        }
73
74        // Get working directory from current directory or args
75        let workdir = args.get("workdir").and_then(|v| v.as_str()).unwrap_or(".");
76
77        // Validate working directory exists
78        if !Path::new(workdir).exists() {
79            return Err(AgentError::ToolError(format!(
80                "Working directory does not exist: {}",
81                workdir
82            )));
83        }
84
85        // Get timeout from args or use default
86        let timeout_secs = args
87            .get("timeout")
88            .and_then(|v| v.as_u64())
89            .unwrap_or(Self::DEFAULT_TIMEOUT_SECS);
90
91        // Execute command with timeout
92        let result = timeout(
93            Duration::from_secs(timeout_secs),
94            Command::new("sh")
95                .args(["-c", command])
96                .current_dir(workdir)
97                .output(),
98        )
99        .await;
100
101        match result {
102            Ok(output) => {
103                let output = output.map_err(|e| {
104                    AgentError::ToolError(format!("Failed to execute command: {}", e))
105                })?;
106
107                let stdout = String::from_utf8_lossy(&output.stdout).to_string();
108                let stderr = String::from_utf8_lossy(&output.stderr).to_string();
109                let exit_code = output.status.code().unwrap_or(-1);
110
111                Ok(serde_json::json!({
112                    "stdout": stdout,
113                    "stderr": stderr,
114                    "exit_code": exit_code
115                }))
116            }
117            Err(_) => Err(AgentError::ToolError(format!(
118                "Command timed out after {} seconds",
119                timeout_secs
120            ))),
121        }
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[tokio::test]
130    async fn test_bash_tool_name() {
131        let tool = BashTool::new();
132        assert_eq!(tool.name(), "bash");
133    }
134
135    #[tokio::test]
136    async fn test_bash_tool_default() {
137        let tool = BashTool;
138        assert_eq!(tool.name(), "bash");
139    }
140
141    #[tokio::test]
142    async fn test_bash_tool_execute_simple() {
143        let tool = BashTool::new();
144        let args = serde_json::json!({
145            "command": "echo 'hello world'"
146        });
147
148        let result = tool.execute(args).await.unwrap();
149
150        assert_eq!(result["stdout"], "hello world\n");
151        assert_eq!(result["exit_code"], 0);
152        assert!(result["stderr"].as_str().unwrap().is_empty());
153    }
154
155    #[tokio::test]
156    async fn test_bash_tool_execute_with_stderr() {
157        let tool = BashTool::new();
158        let args = serde_json::json!({
159            "command": "echo 'error' >&2; exit 1"
160        });
161
162        let result = tool.execute(args).await.unwrap();
163
164        assert_eq!(result["stderr"], "error\n");
165        assert_eq!(result["exit_code"], 1);
166    }
167
168    #[tokio::test]
169    async fn test_bash_tool_missing_command() {
170        let tool = BashTool::new();
171        let args = serde_json::json!({});
172
173        let result = tool.execute(args).await;
174
175        assert!(result.is_err());
176        let err = result.unwrap_err();
177        assert!(err.to_string().contains("Missing 'command'"));
178    }
179
180    #[tokio::test]
181    async fn test_bash_tool_dangerous_command_blocked() {
182        let tool = BashTool::new();
183        let args = serde_json::json!({
184            "command": "rm -rf /"
185        });
186
187        let result = tool.execute(args).await;
188
189        assert!(result.is_err());
190        let err = result.unwrap_err();
191        assert!(err.to_string().contains("Dangerous command blocked"));
192    }
193
194    #[tokio::test]
195    async fn test_bash_tool_fork_bomb_blocked() {
196        let tool = BashTool::new();
197        let args = serde_json::json!({
198            "command": ":(){ :|:& };:"
199        });
200
201        let result = tool.execute(args).await;
202
203        assert!(result.is_err());
204        let err = result.unwrap_err();
205        assert!(err.to_string().contains("Dangerous command blocked"));
206    }
207
208    #[tokio::test]
209    async fn test_bash_tool_timeout() {
210        let tool = BashTool::new();
211        let args = serde_json::json!({
212            "command": "sleep 10",
213            "timeout": 1
214        });
215
216        let result = tool.execute(args).await;
217
218        assert!(result.is_err());
219        let err = result.unwrap_err();
220        assert!(err.to_string().contains("timed out"));
221    }
222
223    #[tokio::test]
224    async fn test_bash_tool_invalid_workdir() {
225        let tool = BashTool::new();
226        let args = serde_json::json!({
227            "command": "echo test",
228            "workdir": "/nonexistent/directory"
229        });
230
231        let result = tool.execute(args).await;
232
233        assert!(result.is_err());
234        let err = result.unwrap_err();
235        assert!(err.to_string().contains("Working directory does not exist"));
236    }
237
238    #[tokio::test]
239    async fn test_bash_tool_current_dir() {
240        let tool = BashTool::new();
241        let args = serde_json::json!({
242            "command": "pwd"
243        });
244
245        let result = tool.execute(args).await.unwrap();
246
247        // Just check we got some output
248        assert!(!result["stdout"].as_str().unwrap().is_empty());
249        assert_eq!(result["exit_code"], 0);
250    }
251}
limit_cli/tools/bash.rs

limit_cli/tools/
bash.rs
