Skip to main content

kindly_tools/commands/
shield.rs

1use anyhow::{Context, Result};
2use clap::{Args, Subcommand};
3use std::io::Write;
4
5/// Security shield commands for protecting AI interactions
6#[derive(Debug, Clone, Args)]
7pub struct ShieldCommand {
8    #[command(subcommand)]
9    pub command: ShieldSubcommand,
10}
11
12#[derive(Debug, Clone, Subcommand)]
13pub enum ShieldSubcommand {
14    /// Auto-wrap AI CLI commands with security scanning
15    AutoWrap(AutoWrapArgs),
16}
17
18#[derive(Debug, Clone, Args)]
19pub struct AutoWrapArgs {
20    /// Shell type to generate for (bash, zsh)
21    #[arg(short = 's', long, default_value = "bash")]
22    shell: String,
23
24    /// Output file for shell functions (defaults to stdout)
25    #[arg(short = 'o', long)]
26    output: Option<String>,
27
28    /// Additional AI commands to wrap (beyond defaults)
29    #[arg(short = 'c', long)]
30    commands: Vec<String>,
31
32    /// Skip default AI commands
33    #[arg(long)]
34    no_defaults: bool,
35}
36
37impl ShieldCommand {
38    pub async fn run(self) -> Result<()> {
39        match self.command {
40            ShieldSubcommand::AutoWrap(args) => execute_auto_wrap(args).await,
41        }
42    }
43}
44
45async fn execute_auto_wrap(args: AutoWrapArgs) -> Result<()> {
46    // Default AI CLI commands to wrap
47    let mut ai_commands = if args.no_defaults {
48        vec![]
49    } else {
50        vec![
51            "claude".to_string(),
52            "openai".to_string(),
53            "gemini".to_string(),
54            "gpt".to_string(),
55            "chatgpt".to_string(),
56            "llm".to_string(),
57            "ai".to_string(),
58            "ollama".to_string(),
59            "anthropic".to_string(),
60            "bard".to_string(),
61        ]
62    };
63
64    // Add custom commands
65    ai_commands.extend(args.commands);
66
67    // Remove duplicates
68    ai_commands.sort();
69    ai_commands.dedup();
70
71    // Generate shell functions
72    let shell_script = match args.shell.as_str() {
73        "bash" | "zsh" => generate_bash_functions(&ai_commands),
74        shell => {
75            return Err(anyhow::anyhow!(
76                "Unsupported shell: {}. Currently only bash and zsh are supported.",
77                shell
78            ))
79        },
80    };
81
82    // Output the script
83    if let Some(output_path) = args.output {
84        std::fs::write(&output_path, &shell_script)
85            .with_context(|| format!("Failed to write to {}", output_path))?;
86        println!("✓ Generated shield wrapper functions in: {}", output_path);
87        println!("\nTo activate, add this to your shell config:");
88        println!("  source {}", output_path);
89    } else {
90        // Write to stdout
91        std::io::stdout()
92            .write_all(shell_script.as_bytes())
93            .with_context(|| "Failed to write to stdout")?;
94    }
95
96    Ok(())
97}
98
99fn generate_bash_functions(commands: &[String]) -> String {
100    let mut script = String::new();
101
102    // Header
103    script.push_str("#!/bin/bash\n");
104    script.push_str("# KindlyGuard Shield - AI CLI Security Wrappers\n");
105    script.push_str("# Generated by: kindly shield auto-wrap\n");
106    script.push_str("# This file wraps AI CLI commands with security scanning\n\n");
107
108    // Check if kindly is available
109    script.push_str("# Check if kindly is installed\n");
110    script.push_str("if ! command -v kindly &> /dev/null; then\n");
111    script.push_str(
112        "    echo \"Warning: kindly command not found. Shield protection disabled.\" >&2\n",
113    );
114    script.push_str("    echo \"Install with: cargo install kindly-tools\" >&2\n");
115    script.push_str("    KINDLY_SHIELD_DISABLED=1\n");
116    script.push_str("fi\n\n");
117
118    // Core shield function
119    script.push_str("# Core shield function that wraps AI commands\n");
120    script.push_str("__kindly_shield_wrap() {\n");
121    script.push_str("    local cmd=\"$1\"\n");
122    script.push_str("    shift\n");
123    script.push_str("    \n");
124    script.push_str("    # If shield is disabled, just run the command\n");
125    script.push_str("    if [[ -n \"$KINDLY_SHIELD_DISABLED\" ]]; then\n");
126    script.push_str("        command \"$cmd\" \"$@\"\n");
127    script.push_str("        return $?\n");
128    script.push_str("    fi\n");
129    script.push_str("    \n");
130    script.push_str("    # Check if we should scan (look for text content in args)\n");
131    script.push_str("    local should_scan=0\n");
132    script.push_str("    local temp_file=\"\"\n");
133    script.push_str("    \n");
134    script.push_str("    # Simple heuristic: scan if args contain text that looks like a prompt\n");
135    script.push_str("    for arg in \"$@\"; do\n");
136    script.push_str("        # Skip flags\n");
137    script.push_str("        if [[ \"$arg\" =~ ^- ]]; then\n");
138    script.push_str("            continue\n");
139    script.push_str("        fi\n");
140    script.push_str("        # If arg has more than 10 characters, it might be a prompt\n");
141    script.push_str("        if [[ ${#arg} -gt 10 ]]; then\n");
142    script.push_str("            should_scan=1\n");
143    script.push_str("            break\n");
144    script.push_str("        fi\n");
145    script.push_str("    done\n");
146    script.push_str("    \n");
147    script.push_str("    if [[ $should_scan -eq 1 ]]; then\n");
148    script.push_str("        # Create temp file with all arguments\n");
149    script.push_str("        temp_file=$(mktemp /tmp/kindly-shield-XXXXXX.txt)\n");
150    script.push_str("        printf '%s\\n' \"$@\" > \"$temp_file\"\n");
151    script.push_str("        \n");
152    script.push_str("        # Scan for threats\n");
153    script.push_str("        if kindly scan \"$temp_file\" --quiet 2>/dev/null; then\n");
154    script.push_str("            # No threats found, proceed\n");
155    script.push_str("            rm -f \"$temp_file\"\n");
156    script.push_str("            command \"$cmd\" \"$@\"\n");
157    script.push_str("        else\n");
158    script.push_str("            # Threats detected\n");
159    script.push_str(
160        "            echo \"⚠️  KindlyGuard Shield: Potential security threats detected!\" >&2\n",
161    );
162    script.push_str("            echo \"Run 'kindly scan $temp_file' for details\" >&2\n");
163    script.push_str("            echo \"To bypass (AT YOUR OWN RISK): KINDLY_SHIELD_DISABLED=1 $cmd ...\" >&2\n");
164    script.push_str("            rm -f \"$temp_file\"\n");
165    script.push_str("            return 1\n");
166    script.push_str("        fi\n");
167    script.push_str("    else\n");
168    script.push_str("        # No text content to scan, just run the command\n");
169    script.push_str("        command \"$cmd\" \"$@\"\n");
170    script.push_str("    fi\n");
171    script.push_str("}\n\n");
172
173    // Generate wrapper function for each command
174    script.push_str("# AI command wrappers\n");
175    for cmd in commands {
176        // Sanitize command name for shell function
177        let safe_cmd = cmd.replace('-', "_");
178
179        script.push_str(&format!("# Wrapper for {}\n", cmd));
180        script.push_str(&format!("{safe_cmd}() {{\n"));
181        script.push_str(&format!("    __kindly_shield_wrap \"{}\" \"$@\"\n", cmd));
182        script.push_str("}\n\n");
183    }
184
185    // Add convenience functions
186    script.push_str("# Convenience functions\n");
187    script.push_str("kindly-shield-status() {\n");
188    script.push_str("    if [[ -n \"$KINDLY_SHIELD_DISABLED\" ]]; then\n");
189    script.push_str("        echo \"🛡️  KindlyGuard Shield: DISABLED\"\n");
190    script.push_str("    else\n");
191    script.push_str("        echo \"🛡️  KindlyGuard Shield: ACTIVE\"\n");
192    script.push_str("        echo \"Protected commands: ");
193    script.push_str(&commands.join(", "));
194    script.push_str("\"\n");
195    script.push_str("    fi\n");
196    script.push_str("}\n\n");
197
198    script.push_str("kindly-shield-disable() {\n");
199    script.push_str("    export KINDLY_SHIELD_DISABLED=1\n");
200    script.push_str(
201        "    echo \"⚠️  KindlyGuard Shield: DISABLED - AI commands are now unprotected!\"\n",
202    );
203    script.push_str("}\n\n");
204
205    script.push_str("kindly-shield-enable() {\n");
206    script.push_str("    unset KINDLY_SHIELD_DISABLED\n");
207    script.push_str("    echo \"✓ KindlyGuard Shield: ENABLED - AI commands are now protected\"\n");
208    script.push_str("}\n\n");
209
210    // Show status on load
211    script.push_str("# Show status when sourced\n");
212    script.push_str("kindly-shield-status\n");
213
214    script
215}