aperion_shield/shims/

check_cmd.rs

//! `aperion-shield --check-cmd -- <command> [args...]`
//!
//! This is the runtime entry point the installed shims invoke. The
//! argv (minus the `--` separator) is reassembled into a single shell-
//! style command string, classified as a `shell` tool call, and
//! evaluated against the active engine. The exit code drives whether
//! the shim then execs the real binary.
//!
//! ## Why we go through the `shell` scope of the engine
//!
//! Every rule that catches destructive shell behaviour today (`rm -rf
//! /`, `aws s3 rm --recursive`, `kubectl delete namespace`, ...) is
//! already authored against the canonical
//! `{"name": "shell", "arguments": {"command": "<...>"}}` shape that
//! both `--check-staged` and the MCP `tools/call` path use. We
//! deliberately reuse that codepath rather than inventing a new rule
//! scope: zero new YAML, zero per-command branching here, and any rule
//! the operator adds for one surface automatically applies to all
//! three (MCP, pre-commit, shim).
//!
//! ## Exit code policy
//!
//! Mirrors `--check-staged` exactly so operators only have to remember
//! one table:
//!
//! | Code | Meaning                                                    |
//! |------|------------------------------------------------------------|
//! | 0    | Engine returned `Allow` (or shadow). Shim execs the real   |
//! |      | binary.                                                    |
//! | 1    | `Block` decision. Shim refuses; banner already printed.    |
//! | 2    | `Approval` / `IdentityVerification`. Can't prompt at shim  |
//! |      | invocation time (no inbox loop), so we surface as a refuse |
//! |      | with a note pointing the user at MCP-mediated invocation.  |
//! | 3    | Operational error (couldn't load shieldset, argv empty).   |

use anyhow::{anyhow, Result};
use serde_json::json;

use crate::engine::Engine;
use crate::{decide, Adjustments, BurstDetector, Decision, WorkspaceContext};

/// Wire-shape result a caller (CLI dispatcher) uses to drive process
/// exit + banner printing. Kept distinct from `Decision` so we can
/// add shim-specific fields later (e.g. captured argv for the audit
/// log) without churning the engine types.
#[derive(Debug, Clone)]
pub struct CheckCmdReport {
    /// Reconstructed command line, useful for the banner and the
    /// JSON-Lines audit record.
    pub command_line: String,
    /// The decision the engine returned for the canonical shell call.
    pub decision: Decision,
    /// The single most-severe rule that fired, if any (for the banner).
    pub primary: Option<PrimaryFinding>,
}

#[derive(Debug, Clone)]
pub struct PrimaryFinding {
    pub rule_id: String,
    pub severity: String,
    pub reason: String,
    pub safer_alternative: Option<String>,
}

impl CheckCmdReport {
    /// See module docstring for the exit-code table.
    pub fn exit_code(&self) -> u8 {
        match &self.decision {
            Decision::Allow => 0,
            Decision::Warn { .. } => 0,
            Decision::Block { .. } => 1,
            Decision::Approval { .. } | Decision::IdentityVerification { .. } => 2,
        }
    }
}

/// Top-level entry point: take the argv that the shim passed through,
/// build a canonical `shell` tool call, and run it through `engine`.
///
/// `argv[0]` is the command name (e.g. "aws"); `argv[1..]` are the
/// per-invocation arguments. We deliberately don't ship a fancy
/// command-line parser here: the engine inspects the reassembled
/// string with the same predicates it already uses for MCP shell
/// calls, which is good enough for v0.8's precision target. v0.9 may
/// add per-CLI argument-tree parsers (proper `aws`, `kubectl`, `gcloud`
/// grammars) on top of this scaffold.
pub fn run(engine: &Engine, argv: &[String]) -> Result<CheckCmdReport> {
    if argv.is_empty() {
        return Err(anyhow!(
            "--check-cmd requires at least the command name after `--`"
        ));
    }

    let cmd_line = reassemble_command_line(argv);

    // Adaptive layer: same setup as the hooks path uses. Decision
    // memory is intentionally absent (each shim invocation is a one-
    // shot in a fresh process; the memory state on disk would be a
    // false signal here -- the operator's last MCP approval doesn't
    // mean their next shell `aws s3 rm` should slide through).
    let cwd = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
    let workspace = WorkspaceContext::probe_at(&engine.policy, &cwd);
    let burst = BurstDetector::new(engine.policy.burst_detector.clone());
    let adj = Adjustments {
        workspace_is_prod: workspace.is_prod,
        burst_in_progress: burst.in_burst(),
        ..Default::default()
    };

    let canonical = json!({"name": "shell", "arguments": {"command": cmd_line}});
    let eval = engine.evaluate("shell", &canonical, adj);
    let decision = decide(&eval);

    let primary = eval
        .matches
        .iter()
        .max_by(|a, b| a.severity.cmp(&b.severity).then(a.points.cmp(&b.points)))
        .map(|m| PrimaryFinding {
            rule_id: m.rule_id.clone(),
            severity: format!("{:?}", m.severity),
            reason: m.reason.clone(),
            safer_alternative: m.safer_alternative.clone(),
        });

    Ok(CheckCmdReport {
        command_line: cmd_line,
        decision,
        primary,
    })
}

/// Reassemble argv into a single shell command string, quoting args
/// that contain shell metacharacters so the engine's existing regex /
/// predicate set sees something close to what the user actually typed.
///
/// This is a presentation step, not a security boundary -- the engine
/// is matching patterns, not re-parsing for execution. We just need
/// `aws s3 rm --recursive s3://prod-bucket` to look like that string
/// rather than `aws s3 rm --recursive s3://prod-bucket` with the
/// fragments concatenated without spaces.
fn reassemble_command_line(argv: &[String]) -> String {
    let mut out = String::new();
    for (i, arg) in argv.iter().enumerate() {
        if i > 0 {
            out.push(' ');
        }
        if needs_quoting(arg) {
            out.push('\'');
            out.push_str(&arg.replace('\'', "'\\''"));
            out.push('\'');
        } else {
            out.push_str(arg);
        }
    }
    out
}

fn needs_quoting(arg: &str) -> bool {
    if arg.is_empty() {
        return true;
    }
    arg.chars().any(|c| {
        c.is_whitespace()
            || c == '"'
            || c == '\''
            || c == '`'
            || c == '$'
            || c == '\\'
            || c == ';'
            || c == '|'
            || c == '&'
            || c == '<'
            || c == '>'
            || c == '('
            || c == ')'
    })
}

/// Render the stderr banner Shield prints when the shim refuses an
/// invocation. Kept in this module (instead of CLI-side) so the format
/// is testable without going through `process::exit`.
pub fn refusal_banner(report: &CheckCmdReport) -> String {
    let mut out = String::new();
    let decision_label = match &report.decision {
        Decision::Block { .. } => "BLOCKED",
        Decision::Approval { .. } => "APPROVAL-REQUIRED",
        Decision::IdentityVerification { .. } => "IDENTITY-REQUIRED",
        Decision::Warn { .. } => "WARN",
        Decision::Allow => "ALLOW",
    };

    out.push_str(&format!("[aperion-shield/check-cmd] {} -- ", decision_label));
    out.push_str(&format!("`{}`\n", short_command(&report.command_line)));

    if let Some(p) = &report.primary {
        out.push_str(&format!(
            "  rule    : {}  (severity={})\n",
            p.rule_id, p.severity
        ));
        out.push_str(&format!("  reason  : {}\n", p.reason));
        if let Some(sa) = &p.safer_alternative {
            out.push_str(&format!("  suggest : {}\n", sa));
        }
    }

    match &report.decision {
        Decision::Approval { .. } | Decision::IdentityVerification { .. } => {
            out.push_str(
                "  note    : approvals require an MCP-mediated invocation (this shim cannot prompt)\n",
            );
        }
        _ => {}
    }

    out.push_str("\nbypass options for a single invocation:\n");
    out.push_str("  SHIELD_SHIMS_DISABLE=1 <command> ...   (env override, one-shot)\n");
    out.push_str("  aperion-shield --uninstall-shims        (remove all shims)\n");
    out
}

/// Trim a very long command line to the first 200 chars so the banner
/// stays readable. Engine match still ran against the full string.
fn short_command(cmd: &str) -> String {
    const CAP: usize = 200;
    if cmd.len() <= CAP {
        cmd.to_string()
    } else {
        let mut s = cmd.chars().take(CAP).collect::<String>();
        s.push_str(" …");
        s
    }
}

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

    #[test]
    fn empty_argv_is_an_operational_error() {
        let engine = Engine::builtin_default();
        let err = run(&engine, &[]).expect_err("empty argv should error");
        assert!(err.to_string().contains("at least the command name"));
    }

    #[test]
    fn reassembly_preserves_simple_argv() {
        let v = vec!["aws".to_string(), "s3".to_string(), "ls".to_string()];
        assert_eq!(reassemble_command_line(&v), "aws s3 ls");
    }

    #[test]
    fn reassembly_quotes_args_with_spaces_and_metacharacters() {
        let v = vec![
            "psql".to_string(),
            "-c".to_string(),
            "DROP TABLE users;".to_string(),
        ];
        let line = reassemble_command_line(&v);
        // The SQL fragment has a space and a `;`, so it should be quoted.
        assert!(line.contains("'DROP TABLE users;'"), "got: {}", line);
    }

    #[test]
    fn reassembly_escapes_embedded_single_quotes() {
        let v = vec!["sh".to_string(), "-c".to_string(), "echo 'hi'".to_string()];
        let line = reassemble_command_line(&v);
        assert!(line.contains("'echo '\\''hi'\\'''"), "got: {}", line);
    }

    #[test]
    fn needs_quoting_picks_up_shell_metacharacters() {
        assert!(needs_quoting("a b"));
        assert!(needs_quoting("a;b"));
        assert!(needs_quoting("a|b"));
        assert!(needs_quoting("a>b"));
        assert!(needs_quoting("`whoami`"));
        assert!(!needs_quoting("aws"));
        assert!(!needs_quoting("--recursive"));
        assert!(!needs_quoting("s3://prod-bucket"));
    }

    #[test]
    fn exit_code_for_allow_is_zero() {
        let report = CheckCmdReport {
            command_line: "aws s3 ls".into(),
            decision: Decision::Allow,
            primary: None,
        };
        assert_eq!(report.exit_code(), 0);
    }

    #[test]
    fn run_with_innocuous_command_returns_allow() {
        let engine = Engine::builtin_default();
        let report = run(&engine, &["aws".to_string(), "s3".to_string(), "ls".to_string()])
            .expect("run");
        assert!(matches!(
            report.decision,
            Decision::Allow | Decision::Warn { .. }
        ));
        assert_eq!(report.exit_code(), 0);
    }

    #[test]
    fn refusal_banner_includes_command_rule_and_bypass_note() {
        use crate::Severity;
        let report = CheckCmdReport {
            command_line: "rm -rf /".into(),
            decision: Decision::Block {
                rule_id: "fs.rm_root".into(),
                severity: Severity::Critical,
                reason: "rm -rf / is non-recoverable".into(),
                safer_alternative: Some("rm -rf <specific-path>".into()),
                contributing_rules: vec!["fs.rm_root".into()],
            },
            primary: Some(PrimaryFinding {
                rule_id: "fs.rm_root".into(),
                severity: "Critical".into(),
                reason: "rm -rf / is non-recoverable".into(),
                safer_alternative: Some("rm -rf <specific-path>".into()),
            }),
        };
        let banner = refusal_banner(&report);
        assert!(banner.contains("BLOCKED"));
        assert!(banner.contains("rm -rf /"));
        assert!(banner.contains("fs.rm_root"));
        assert!(banner.contains("SHIELD_SHIMS_DISABLE"));
        assert!(banner.contains("aperion-shield --uninstall-shims"));
    }

    #[test]
    fn short_command_truncates_long_lines() {
        let long = "a".repeat(500);
        let s = short_command(&long);
        assert!(s.len() < long.len());
        assert!(s.ends_with(" …"));
    }
}