redacted-cli 1.1.0

//! Claude Code PostToolUse protocol, ported from the Go 0.7 hook. The
//! updatedToolOutput replacement mirrors the tool_response shape (issue #46).

use std::collections::BTreeMap;
use std::panic::{self, AssertUnwindSafe};

use serde::Serialize;
use serde_json::Value;

use crate::scrub::ScrubResult;

/// Printed in place of the tool output when scrubbing failed.
pub(crate) const WITHHELD: &str = "[redacted] tool output withheld: the scrubber errored, so raw output was suppressed to avoid leaking secrets.";

// Typed structs rather than `json!`: serde_json sorts `json!` keys, and the
// envelope's key order is pinned byte for byte. `'a` lets fields borrow strings.
#[derive(Serialize)]
struct HookOutput<'a> {
    decision: &'a str,
    reason: &'a str,
    #[serde(rename = "hookSpecificOutput")]
    hook_specific_output: HookSpecificOutput<'a>,
}

#[derive(Serialize)]
struct HookSpecificOutput<'a> {
    #[serde(rename = "hookEventName")]
    hook_event_name: &'a str,
    // A Value, not a string: Claude Code applies the replacement only when it
    // matches the tool's output schema (an object for Bash, Read, Grep, ...).
    #[serde(rename = "updatedToolOutput", skip_serializing_if = "Value::is_null")]
    updated_tool_output: Value,
}

/// The bytes to print and the per-pattern hit counts to record.
type HookResult = (Vec<u8>, BTreeMap<String, usize>);

/// Scrubs the tool output in a PostToolUse payload. No hit returns no bytes so
/// the original passes through; a hit returns the block envelope.
fn scrub_payload(data: &[u8], scrub: &dyn Fn(&str) -> ScrubResult) -> Result<HookResult, String> {
    let payload: Value =
        serde_json::from_slice(data).map_err(|e| format!("parse hook payload: {e}"))?;
    let tool_name = string_field_or_empty(&payload, "tool_name")?;
    let response = payload.get("tool_response").unwrap_or(&Value::Null);
    if tool_name == "Bash" {
        scrub_bash_response(response, scrub)
    } else {
        scrub_tool_response(response, &tool_name, scrub)
    }
}

/// Fail closed: an error or panic withholds the output rather than letting
/// raw, unscrubbed bytes through.
pub fn scrub_payload_or_withhold(data: &[u8], scrub: &dyn Fn(&str) -> ScrubResult) -> HookResult {
    // catch_unwind turns a panic into an Err. AssertUnwindSafe vouches that nothing
    // the panic may have left half-written is used afterwards: it is all dropped.
    match panic::catch_unwind(AssertUnwindSafe(|| scrub_payload(data, scrub))) {
        Ok(Ok(hook_result)) => hook_result,
        _ => (withheld_output(data), BTreeMap::new()),
    }
}

/// The fail-closed envelope. Its replacement mirrors `tool_response` with every
/// string withheld; unparsable input or no `tool_response` gets the bare sentence.
pub fn withheld_output(data: &[u8]) -> Vec<u8> {
    let payload: Value = serde_json::from_slice(data).unwrap_or(Value::Null);
    let updated = match payload.get("tool_response") {
        None | Some(Value::Null) => Value::String(WITHHELD.into()),
        Some(response) => withhold_strings(response),
    };
    encode(&HookOutput {
        decision: "block",
        reason: WITHHELD,
        hook_specific_output: HookSpecificOutput {
            hook_event_name: "PostToolUse",
            updated_tool_output: updated,
        },
    })
    .expect("a parsed Value always serializes")
}

/// Same shape, every string leaf replaced by the WITHHELD sentence.
pub(crate) fn withhold_strings(value: &Value) -> Value {
    match value {
        Value::String(_) => Value::String(WITHHELD.into()),
        Value::Object(map) => Value::Object(
            map.iter()
                .map(|(key, child)| (key.clone(), withhold_strings(child)))
                .collect(),
        ),
        Value::Array(items) => Value::Array(items.iter().map(withhold_strings).collect()),
        other => other.clone(),
    }
}

/// Puts `line` above an envelope's reason; `updatedToolOutput` is left alone
/// because it stands in for the tool result the model reads.
pub fn prepend_reason(out: &[u8], line: &str) -> Vec<u8> {
    let Ok(mut envelope) = serde_json::from_slice::<Value>(out) else {
        return out.to_vec();
    };
    let Some(reason) = envelope["reason"].as_str() else {
        return out.to_vec();
    };
    envelope["reason"] = Value::String(format!("{line}\n{reason}"));
    encode(&envelope).unwrap_or_else(|_| out.to_vec())
}

// Go decodes into typed structs, so a wrong type is an error, null is "".
fn string_field_or_empty(value: &Value, key: &str) -> Result<String, String> {
    match value.get(key) {
        None | Some(Value::Null) => Ok(String::new()),
        Some(Value::String(s)) => Ok(s.clone()),
        Some(_) => Err(format!("parse hook payload: {key} is not a string")),
    }
}

fn scrub_bash_response(
    response: &Value,
    scrub: &dyn Fn(&str) -> ScrubResult,
) -> Result<HookResult, String> {
    if !(response.is_object() || response.is_null()) {
        return Err("parse hook payload: tool_response is not an object".into());
    }
    let stdout_result = scrub(&string_field_or_empty(response, "stdout")?);
    let stderr_result = scrub(&string_field_or_empty(response, "stderr")?);
    if !stdout_result.has_redactions() && !stderr_result.has_redactions() {
        return Ok((Vec::new(), BTreeMap::new()));
    }

    // Only fields that came in as strings are replaced: no key is invented.
    let mut updated = response.clone();
    for (key, scrubbed) in [
        ("stdout", &stdout_result.text),
        ("stderr", &stderr_result.text),
    ] {
        // `slot @ Value::String(_)` binds the matched value so it can be overwritten.
        if let Some(slot @ Value::String(_)) = updated.get_mut(key) {
            *slot = Value::String(scrubbed.clone());
        }
    }

    let mut counts = stdout_result.counts_by_pattern;
    for (pattern, count) in stderr_result.counts_by_pattern {
        // entry() finds or inserts the key; or_default() starts a new count at 0.
        *counts.entry(pattern).or_default() += count;
    }
    // Stdout lines, then stderr ones: updatedToolOutput keeps the streams apart.
    let scrubbed = format!("{}\n{}", stdout_result.text, stderr_result.text);
    let block = block_output("command", &counts, &scrubbed, updated)?;
    Ok((block, counts))
}

fn scrub_tool_response(
    response: &Value,
    tool_name: &str,
    scrub: &dyn Fn(&str) -> ScrubResult,
) -> Result<HookResult, String> {
    let text = text_to_scrub(response)?;
    let result = scrub(&text);
    if !result.has_redactions() {
        return Ok((Vec::new(), BTreeMap::new()));
    }

    // reason may summarize, the replacement may not: it stands in for the result.
    let updated = match response {
        Value::String(_) => Value::String(result.text.clone()),
        other => scrub_json(other, scrub),
    };
    let block = block_output(tool_name, &result.counts_by_pattern, &result.text, updated)?;
    Ok((block, result.counts_by_pattern))
}

/// Scrubber input: the string itself, or every string in a structured value.
pub(crate) fn text_to_scrub(response: &Value) -> Result<String, String> {
    if let Value::String(s) = response {
        return Ok(s.clone());
    }
    let mut parts = Vec::new();
    collect_strings(response, &mut parts);
    if !parts.is_empty() {
        return Ok(parts.join("\n"));
    }
    match response {
        Value::Null => Ok(String::new()),
        other => to_go_json(other),
    }
}

/// Every non-empty string leaf and object key, keys in sorted order (serde_json's
/// map is sorted). Keys count: a secret can arrive as one (e.g. Grep counts).
fn collect_strings(value: &Value, strings: &mut Vec<String>) {
    match value {
        Value::String(s) if !s.is_empty() => strings.push(s.clone()),
        Value::Object(map) => {
            for (key, child) in map {
                if !key.is_empty() {
                    strings.push(key.clone());
                }
                collect_strings(child, strings);
            }
        }
        Value::Array(items) => {
            for item in items {
                collect_strings(item, strings);
            }
        }
        _ => {}
    }
}

pub(crate) fn scrub_json(value: &Value, scrub: &dyn Fn(&str) -> ScrubResult) -> Value {
    match value {
        Value::String(s) => Value::String(scrub(s).text),
        Value::Object(map) => Value::Object(
            map.iter()
                .map(|(key, child)| (scrub(key).text, scrub_json(child, scrub)))
                .collect(),
        ),
        Value::Array(items) => {
            Value::Array(items.iter().map(|item| scrub_json(item, scrub)).collect())
        }
        other => other.clone(),
    }
}

fn marker_lines(scrubbed: &str) -> String {
    let lines: Vec<String> = scrubbed
        .split('\n')
        .filter(|line| line.contains("[REDACTED"))
        .map(|line| format!("- {line}"))
        .collect();
    if lines.is_empty() {
        return "Secret(s) removed from tool response.".to_string();
    }
    lines.join("\n")
}

/// The reason names the hits only: Claude Code shows its first line collapsed,
/// the marker lines on expand, and the model reads the output in `updated`.
fn block_output(
    source: &str,
    counts_by_pattern: &BTreeMap<String, usize>,
    scrubbed: &str,
    updated: Value,
) -> Result<Vec<u8>, String> {
    let count: usize = counts_by_pattern.values().sum();
    let mut hits: Vec<_> = counts_by_pattern.iter().collect();
    // Most hits first, then by name, as `redacted stats` orders them.
    hits.sort_by(|left, right| right.1.cmp(left.1).then(left.0.cmp(right.0)));
    let summary: Vec<String> = hits
        .iter()
        .map(|(name, hit_count)| match hit_count {
            1 => name.to_string(),
            _ => format!("{name} ×{hit_count}"),
        })
        .collect();
    let summary = summary.join(", ");
    encode(&HookOutput {
        decision: "block",
        reason: &format!(
            "[redacted] {count} secret(s) scrubbed from {source} output: {summary}\n\n{}",
            marker_lines(scrubbed)
        ),
        hook_specific_output: HookSpecificOutput {
            hook_event_name: "PostToolUse",
            updated_tool_output: updated,
        },
    })
}

/// Go's encoder (HTML escaping off) plus its trailing newline.
pub fn encode<T: Serialize>(value: &T) -> Result<Vec<u8>, String> {
    let mut out = to_go_json(value)?.into_bytes();
    out.push(b'\n');
    Ok(out)
}

// Go always escapes U+2028/U+2029; serde_json never does. Both only ever
// appear inside JSON strings, so a plain replace is safe.
fn to_go_json<T: Serialize>(value: &T) -> Result<String, String> {
    let json = serde_json::to_string(value).map_err(|e| format!("encode response: {e}"))?;
    Ok(json
        .replace('\u{2028}', "\\u2028")
        .replace('\u{2029}', "\\u2029"))
}

#[cfg(test)]
mod tests;