use crate::security::permission::{ PermissionHandler, PermissionRequest, PermissionResult, ResourceType, }; use crate::security::{SecurityValidator, ValidationResult}; use crate::tools::Tool; use anyhow::Result; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use std::collections::HashSet; use std::process::Stdio; use std::sync::{Arc, RwLock}; use tokio::io::AsyncReadExt; use tokio::process::Command; pub struct RunCommandTool { validator: Arc, permission_handler: Arc, runtime_allowed: Arc>>, } impl RunCommandTool { pub fn new( validator: Arc, permission_handler: Arc, ) -> Self { Self { validator, permission_handler, runtime_allowed: Arc::new(RwLock::new(HashSet::new())), } } async fn execute_command(&self, params: &RunCommandParams) -> Result { let mut cmd = if cfg!(target_os = "windows") { let mut c = Command::new("cmd"); c.args(["/C", ¶ms.command]); c } else { let mut c = Command::new("sh"); c.args(["-c", ¶ms.command]); c }; if let Some(ref dir) = params.working_dir { cmd.current_dir(dir); } cmd.stdout(Stdio::piped()); cmd.stderr(Stdio::piped()); let mut child = cmd.spawn()?; let mut stdout = String::new(); let mut stderr = String::new(); if let Some(mut out) = child.stdout.take() { out.read_to_string(&mut stdout).await?; } if let Some(mut err) = child.stderr.take() { err.read_to_string(&mut stderr).await?; } let status = child.wait().await?; let output = if !stdout.is_empty() { stdout } else { stderr }; if status.success() { Ok(output) } else { anyhow::bail!("Command failed: {}", output) } } } #[derive(Debug, Serialize, Deserialize)] struct RunCommandParams { command: String, #[serde(default)] working_dir: Option, } #[async_trait] impl Tool for RunCommandTool { fn name(&self) -> &str { "run_command" } fn description(&self) -> &str { "Execute a shell command and return its output" } fn parameters(&self) -> serde_json::Value { serde_json::json!({ "type": "object", "properties": { "command": { "type": "string", "description": "The shell command to execute" }, "working_dir": { "type": "string", "description": "Optional working directory for the command" } }, "required": ["command"] }) } async fn execute(&self, params: serde_json::Value) -> Result { let params: RunCommandParams = serde_json::from_value(params)?; // Check if command was previously allowed let base_cmd = params.command.split_whitespace().next().unwrap_or(""); let is_allowed = { let allowed = self.runtime_allowed.read().unwrap(); allowed.contains(base_cmd) }; if is_allowed { return self.execute_command(¶ms).await; } // Validate command match self.validator.validate_shell_command(¶ms.command) { ValidationResult::Allowed => self.execute_command(¶ms).await, ValidationResult::Denied(reason) => { anyhow::bail!("Command denied: {}", reason) } ValidationResult::RequiresPermission(reason) => { let request = PermissionRequest { resource_type: ResourceType::ShellCommand, action: params.command.clone(), reason, }; match self.permission_handler.request_permission(&request) { PermissionResult::Allow => self.execute_command(¶ms).await, PermissionResult::Deny => { anyhow::bail!("Permission denied by user") } PermissionResult::AllowAlways(cmd) => { { let mut allowed = self.runtime_allowed.write().unwrap(); allowed.insert(cmd); } self.execute_command(¶ms).await } PermissionResult::Quit => { anyhow::bail!("User requested quit") } } } } } }