kindly-tools 0.11.14

use anyhow::{Context, Result};
use clap::{Args, Subcommand};
use std::io::Write;

/// Security shield commands for protecting AI interactions
#[derive(Debug, Clone, Args)]
pub struct ShieldCommand {
    #[command(subcommand)]
    pub command: ShieldSubcommand,
}

#[derive(Debug, Clone, Subcommand)]
pub enum ShieldSubcommand {
    /// Auto-wrap AI CLI commands with security scanning
    AutoWrap(AutoWrapArgs),
}

#[derive(Debug, Clone, Args)]
pub struct AutoWrapArgs {
    /// Shell type to generate for (bash, zsh)
    #[arg(short = 's', long, default_value = "bash")]
    shell: String,

    /// Output file for shell functions (defaults to stdout)
    #[arg(short = 'o', long)]
    output: Option<String>,

    /// Additional AI commands to wrap (beyond defaults)
    #[arg(short = 'c', long)]
    commands: Vec<String>,

    /// Skip default AI commands
    #[arg(long)]
    no_defaults: bool,
}

impl ShieldCommand {
    pub async fn run(self) -> Result<()> {
        match self.command {
            ShieldSubcommand::AutoWrap(args) => execute_auto_wrap(args).await,
        }
    }
}

async fn execute_auto_wrap(args: AutoWrapArgs) -> Result<()> {
    // Default AI CLI commands to wrap
    let mut ai_commands = if args.no_defaults {
        vec![]
    } else {
        vec![
            "claude".to_string(),
            "openai".to_string(),
            "gemini".to_string(),
            "gpt".to_string(),
            "chatgpt".to_string(),
            "llm".to_string(),
            "ai".to_string(),
            "ollama".to_string(),
            "anthropic".to_string(),
            "bard".to_string(),
        ]
    };

    // Add custom commands
    ai_commands.extend(args.commands);

    // Remove duplicates
    ai_commands.sort();
    ai_commands.dedup();

    // Generate shell functions
    let shell_script = match args.shell.as_str() {
        "bash" | "zsh" => generate_bash_functions(&ai_commands),
        shell => {
            return Err(anyhow::anyhow!(
                "Unsupported shell: {}. Currently only bash and zsh are supported.",
                shell
            ))
        },
    };

    // Output the script
    if let Some(output_path) = args.output {
        std::fs::write(&output_path, &shell_script)
            .with_context(|| format!("Failed to write to {}", output_path))?;
        println!("✓ Generated shield wrapper functions in: {}", output_path);
        println!("\nTo activate, add this to your shell config:");
        println!("  source {}", output_path);
    } else {
        // Write to stdout
        std::io::stdout()
            .write_all(shell_script.as_bytes())
            .with_context(|| "Failed to write to stdout")?;
    }

    Ok(())
}

fn generate_bash_functions(commands: &[String]) -> String {
    let mut script = String::new();

    // Header
    script.push_str("#!/bin/bash\n");
    script.push_str("# KindlyGuard Shield - AI CLI Security Wrappers\n");
    script.push_str("# Generated by: kindly shield auto-wrap\n");
    script.push_str("# This file wraps AI CLI commands with security scanning\n\n");

    // Check if kindly is available
    script.push_str("# Check if kindly is installed\n");
    script.push_str("if ! command -v kindly &> /dev/null; then\n");
    script.push_str(
        "    echo \"Warning: kindly command not found. Shield protection disabled.\" >&2\n",
    );
    script.push_str("    echo \"Install with: cargo install kindly-tools\" >&2\n");
    script.push_str("    KINDLY_SHIELD_DISABLED=1\n");
    script.push_str("fi\n\n");

    // Core shield function
    script.push_str("# Core shield function that wraps AI commands\n");
    script.push_str("__kindly_shield_wrap() {\n");
    script.push_str("    local cmd=\"$1\"\n");
    script.push_str("    shift\n");
    script.push_str("    \n");
    script.push_str("    # If shield is disabled, just run the command\n");
    script.push_str("    if [[ -n \"$KINDLY_SHIELD_DISABLED\" ]]; then\n");
    script.push_str("        command \"$cmd\" \"$@\"\n");
    script.push_str("        return $?\n");
    script.push_str("    fi\n");
    script.push_str("    \n");
    script.push_str("    # Check if we should scan (look for text content in args)\n");
    script.push_str("    local should_scan=0\n");
    script.push_str("    local temp_file=\"\"\n");
    script.push_str("    \n");
    script.push_str("    # Simple heuristic: scan if args contain text that looks like a prompt\n");
    script.push_str("    for arg in \"$@\"; do\n");
    script.push_str("        # Skip flags\n");
    script.push_str("        if [[ \"$arg\" =~ ^- ]]; then\n");
    script.push_str("            continue\n");
    script.push_str("        fi\n");
    script.push_str("        # If arg has more than 10 characters, it might be a prompt\n");
    script.push_str("        if [[ ${#arg} -gt 10 ]]; then\n");
    script.push_str("            should_scan=1\n");
    script.push_str("            break\n");
    script.push_str("        fi\n");
    script.push_str("    done\n");
    script.push_str("    \n");
    script.push_str("    if [[ $should_scan -eq 1 ]]; then\n");
    script.push_str("        # Create temp file with all arguments\n");
    script.push_str("        temp_file=$(mktemp /tmp/kindly-shield-XXXXXX.txt)\n");
    script.push_str("        printf '%s\\n' \"$@\" > \"$temp_file\"\n");
    script.push_str("        \n");
    script.push_str("        # Scan for threats\n");
    script.push_str("        if kindly scan \"$temp_file\" --quiet 2>/dev/null; then\n");
    script.push_str("            # No threats found, proceed\n");
    script.push_str("            rm -f \"$temp_file\"\n");
    script.push_str("            command \"$cmd\" \"$@\"\n");
    script.push_str("        else\n");
    script.push_str("            # Threats detected\n");
    script.push_str(
        "            echo \"⚠️  KindlyGuard Shield: Potential security threats detected!\" >&2\n",
    );
    script.push_str("            echo \"Run 'kindly scan $temp_file' for details\" >&2\n");
    script.push_str("            echo \"To bypass (AT YOUR OWN RISK): KINDLY_SHIELD_DISABLED=1 $cmd ...\" >&2\n");
    script.push_str("            rm -f \"$temp_file\"\n");
    script.push_str("            return 1\n");
    script.push_str("        fi\n");
    script.push_str("    else\n");
    script.push_str("        # No text content to scan, just run the command\n");
    script.push_str("        command \"$cmd\" \"$@\"\n");
    script.push_str("    fi\n");
    script.push_str("}\n\n");

    // Generate wrapper function for each command
    script.push_str("# AI command wrappers\n");
    for cmd in commands {
        // Sanitize command name for shell function
        let safe_cmd = cmd.replace('-', "_");

        script.push_str(&format!("# Wrapper for {}\n", cmd));
        script.push_str(&format!("{safe_cmd}() {{\n"));
        script.push_str(&format!("    __kindly_shield_wrap \"{}\" \"$@\"\n", cmd));
        script.push_str("}\n\n");
    }

    // Add convenience functions
    script.push_str("# Convenience functions\n");
    script.push_str("kindly-shield-status() {\n");
    script.push_str("    if [[ -n \"$KINDLY_SHIELD_DISABLED\" ]]; then\n");
    script.push_str("        echo \"🛡️  KindlyGuard Shield: DISABLED\"\n");
    script.push_str("    else\n");
    script.push_str("        echo \"🛡️  KindlyGuard Shield: ACTIVE\"\n");
    script.push_str("        echo \"Protected commands: ");
    script.push_str(&commands.join(", "));
    script.push_str("\"\n");
    script.push_str("    fi\n");
    script.push_str("}\n\n");

    script.push_str("kindly-shield-disable() {\n");
    script.push_str("    export KINDLY_SHIELD_DISABLED=1\n");
    script.push_str(
        "    echo \"⚠️  KindlyGuard Shield: DISABLED - AI commands are now unprotected!\"\n",
    );
    script.push_str("}\n\n");

    script.push_str("kindly-shield-enable() {\n");
    script.push_str("    unset KINDLY_SHIELD_DISABLED\n");
    script.push_str("    echo \"✓ KindlyGuard Shield: ENABLED - AI commands are now protected\"\n");
    script.push_str("}\n\n");

    // Show status on load
    script.push_str("# Show status when sourced\n");
    script.push_str("kindly-shield-status\n");

    script
}