tiny-trae 0.1.0

use anyhow::{Context, Result};
use serde_json::{json, Value};
use std::path::Path;
use std::process::Command;
use log::{debug, info, warn};

use crate::config::Config;
use crate::claude::{ToolDefinition, FunctionDefinition};
use super::{Tool, ToolResult};

pub struct BashTool;

#[async_trait::async_trait]
impl Tool for BashTool {
    fn get_definition(&self, name: &str) -> ToolDefinition {
        ToolDefinition {
            r#type: "function".to_string(),
            function: FunctionDefinition {
                name: name.to_string(),
                description: "Execute bash commands".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "command": {
                            "type": "string",
                            "description": "The bash command to execute"
                        },
                        "working_directory": {
                            "type": "string",
                            "description": "The working directory to run the command in (optional)"
                        },
                        "timeout": {
                            "type": "integer",
                            "description": "Timeout in seconds (optional, defaults to 30)"
                        }
                    },
                    "required": ["command"]
                }),
            },
        }
    }

    async fn execute(&self, input: Value, config: &Config) -> Result<ToolResult> {
        let command = input
            .get("command")
            .and_then(|v| v.as_str())
            .context("Missing or invalid 'command' parameter")?;

        let working_directory = input
            .get("working_directory")
            .and_then(|v| v.as_str())
            .unwrap_or(&config.workspace.root_path);

        let timeout = input
            .get("timeout")
            .and_then(|v| v.as_u64())
            .unwrap_or(30);

        // Security check - prevent dangerous commands
        if is_dangerous_command(command) {
            warn!("Dangerous command blocked: {}", command);
            return Ok(ToolResult::error("Dangerous command blocked for security reasons".to_string()));
        }

        let work_dir = Path::new(working_directory);
        if !work_dir.exists() {
            return Ok(ToolResult::error(format!("Working directory does not exist: {}", working_directory)));
        }

        debug!("Executing command: {} in directory: {}", command, working_directory);

        let mut cmd = Command::new("bash");
        cmd.arg("-c")
           .arg(command)
           .current_dir(work_dir);

        // Set timeout
        let output = tokio::time::timeout(
            std::time::Duration::from_secs(timeout),
            tokio::task::spawn_blocking(move || cmd.output())
        ).await;

        match output {
            Ok(Ok(Ok(output))) => {
                let stdout = String::from_utf8_lossy(&output.stdout);
                let stderr = String::from_utf8_lossy(&output.stderr);
                
                let mut result = String::new();
                if !stdout.is_empty() {
                    result.push_str("STDOUT:\n");
                    result.push_str(&stdout);
                }
                if !stderr.is_empty() {
                    if !result.is_empty() {
                        result.push_str("\n");
                    }
                    result.push_str("STDERR:\n");
                    result.push_str(&stderr);
                }
                
                if output.status.success() {
                    info!("Command executed successfully");
                    Ok(ToolResult::success(result))
                } else {
                    debug!("Command failed with exit code: {:?}", output.status.code());
                    Ok(ToolResult::error(format!("Command failed with exit code: {:?}\n{}", output.status.code(), result)))
                }
            }
            Ok(Ok(Err(e))) => {
                debug!("Failed to execute command: {}", e);
                Ok(ToolResult::error(format!("Failed to execute command: {}", e)))
            }
            Ok(Err(_)) => {
                debug!("Command execution was cancelled");
                Ok(ToolResult::error("Command execution was cancelled".to_string()))
            }
            Err(_) => {
                debug!("Command execution timed out");
                Ok(ToolResult::error(format!("Command execution timed out after {} seconds", timeout)))
            }
        }
    }
}

fn is_dangerous_command(command: &str) -> bool {
    let dangerous_patterns = [
        "rm -rf",
        "rm -fr",
        "rm -r",
        "rm -f",
        "rmdir",
        "del /s",
        "del /q",
        "format",
        "fdisk",
        "mkfs",
        "dd if=",
        "dd of=",
        "kill -9",
        "killall",
        "pkill",
        "shutdown",
        "reboot",
        "halt",
        "poweroff",
        "init 0",
        "init 6",
        ":/etc/passwd",
        ":/etc/shadow",
        "chmod 777",
        "chmod -R 777",
        "chown -R",
        "usermod",
        "userdel",
        "groupdel",
        "crontab -r",
        "history -c",
        "history -w",
        "shred",
        "wipe",
        "srm",
        "cat /dev/urandom",
        "cat /dev/zero",
        ":(){ :|:& };:",
        "fork bomb",
        "while true",
        "curl.*|.*bash",
        "wget.*|.*bash",
        "nc -l",
        "netcat -l",
        "python -c",
        "perl -e",
        "ruby -e",
        "eval(",
        "exec(",
        "system(",
        "shell_exec(",
        "passthru(",
        "popen(",
        "proc_open(",
    ];

    let command_lower = command.to_lowercase();
    dangerous_patterns.iter().any(|pattern| command_lower.contains(pattern))
}