mrapids 0.1.31

//! Response redaction for sensitive data
//!
//! Automatically redacts sensitive information from API responses before
//! displaying to users or returning to AI agents.

use regex::Regex;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use std::collections::HashSet;
use std::sync::LazyLock;

/// Common patterns for sensitive field names
static SENSITIVE_KEY_PATTERNS: LazyLock<Vec<Regex>> = LazyLock::new(|| {
    vec![
        Regex::new(r"(?i)(api[_-]?key|apikey)").unwrap(),
        Regex::new(r"(?i)(password|passwd|pwd)").unwrap(),
        Regex::new(r"(?i)(secret|private[_-]?key)").unwrap(),
        Regex::new(r"(?i)(token|access[_-]?token|refresh[_-]?token)").unwrap(),
        Regex::new(r"(?i)(auth|authorization)").unwrap(),
        Regex::new(r"(?i)(credential|cred)").unwrap(),
        Regex::new(r"(?i)(ssn|social[_-]?security)").unwrap(),
        Regex::new(r"(?i)(credit[_-]?card|card[_-]?number|cvv|cvc)").unwrap(),
    ]
});

/// Patterns for sensitive values
// Matches SSN with dashes (123-45-6789) or without (123456789)
static SSN_PATTERN: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"\b(\d{3})-?(\d{2})-?(\d{4})\b").unwrap());

static CARD_PATTERN: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"\b(\d{4})[\s-]?(\d{4})[\s-]?(\d{4})[\s-]?(\d{4})\b").unwrap());

static JWT_PATTERN: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"^[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+$").unwrap());

static EMAIL_PATTERN: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b").unwrap());

static PHONE_PATTERN: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"\b(\d{3})[-.]?(\d{3})[-.]?(\d{4})\b").unwrap());

/// Known token prefixes that indicate sensitive values
static TOKEN_PREFIXES: [&str; 7] = ["Bearer ", "Basic ", "Token ", "AWS", "sk_", "pk_", "ghp_"];

/// Controls how sensitive values are masked
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum MaskingMode {
    /// Replace entirely with `[REDACTED]` — legacy behavior
    Full,
    /// Keep last 4 chars visible (or type-specific partial mask)
    Partial,
}

/// Redactor for sensitive data in API responses
pub struct ResponseRedactor {
    custom_patterns: Vec<Regex>,
    sensitive_keys: HashSet<String>,
    enabled: bool,
    masking_mode: MaskingMode,
}

impl ResponseRedactor {
    /// Create a new redactor with optional custom patterns (defaults to Partial masking)
    pub fn new(custom_patterns: Vec<String>, enabled: bool) -> Self {
        Self::new_with_mode(custom_patterns, enabled, MaskingMode::Partial)
    }

    /// Create a new redactor with explicit masking mode
    pub fn new_with_mode(custom_patterns: Vec<String>, enabled: bool, mode: MaskingMode) -> Self {
        let mut sensitive_keys = HashSet::new();

        // Default sensitive key names
        for key in &[
            "password",
            "secret",
            "token",
            "api_key",
            "apikey",
            "auth",
            "authorization",
            "credential",
            "private_key",
            "access_token",
            "refresh_token",
            "ssn",
            "credit_card",
            "cvv",
        ] {
            sensitive_keys.insert(key.to_string());
        }

        let compiled_patterns = custom_patterns
            .into_iter()
            .filter_map(|pattern| match Regex::new(&pattern) {
                Ok(regex) => Some(regex),
                Err(e) => {
                    eprintln!("Warning: Invalid redact pattern '{}': {}", pattern, e);
                    None
                }
            })
            .collect();

        Self {
            custom_patterns: compiled_patterns,
            sensitive_keys,
            enabled,
            masking_mode: mode,
        }
    }

    /// Create a default redactor (enabled, Partial masking)
    #[allow(dead_code)]
    pub fn default_enabled() -> Self {
        Self::new(vec![], true)
    }

    /// Redact sensitive data from a JSON value
    pub fn redact(&self, value: &Value) -> Value {
        if !self.enabled {
            return value.clone();
        }
        self.redact_value(value)
    }

    /// Produce a masked representation of `value` based on `data_type`.
    ///
    /// `data_type` hints: `"ssn"`, `"card"`, `"email"`, `"phone"`, `"jwt"`,
    /// `"api_key"` / `"token"`, or `"generic"`.
    fn mask_value(&self, value: &str, data_type: &str) -> String {
        if self.masking_mode == MaskingMode::Full {
            return "[REDACTED]".to_string();
        }

        match data_type {
            "ssn" => {
                // 123-45-6789 → ***-**-6789
                if let Some(caps) = SSN_PATTERN.captures(value) {
                    if let Some(last) = caps.get(3) {
                        return format!("***-**-{}", last.as_str());
                    }
                }
                mask_generic(value)
            }
            "card" => {
                // 4111-1111-1111-1234 → ****-****-****-1234
                if let Some(caps) = CARD_PATTERN.captures(value) {
                    if let Some(last) = caps.get(4) {
                        return format!("****-****-****-{}", last.as_str());
                    }
                }
                mask_generic(value)
            }
            "email" => {
                // alice@example.com → a***@***.com
                if let Some(at_pos) = value.find('@') {
                    let local = &value[..at_pos];
                    let domain = &value[at_pos + 1..];
                    let first_char = local.chars().next().unwrap_or('*');
                    let tld = domain.rsplit('.').next().unwrap_or("com");
                    return format!("{}***@***.{}", first_char, tld);
                }
                mask_generic(value)
            }
            "phone" => {
                // 555-123-7890 → ***-***-7890
                if let Some(caps) = PHONE_PATTERN.captures(value) {
                    if let Some(last) = caps.get(3) {
                        return format!("***-***-{}", last.as_str());
                    }
                }
                mask_generic(value)
            }
            "jwt" => {
                // Compute a short fingerprint
                let mut hasher = Sha256::new();
                hasher.update(value.as_bytes());
                let hash = hasher.finalize();
                let hex = hex::encode(hash);
                format!("[JWT:{}]", &hex[..12])
            }
            "api_key" | "token" => {
                // sk_test_abc123 → sk_t****bc23
                let len = value.len();
                if len <= 8 {
                    return mask_generic(value);
                }
                let prefix: String = value.chars().take(4).collect();
                let suffix: String = value.chars().skip(len - 4).collect();
                format!("{}****{}", prefix, suffix)
            }
            _ => mask_generic(value),
        }
    }

    /// Determine the data_type hint for a key-based redaction
    fn data_type_for_key(key: &str) -> &'static str {
        let k = key.to_lowercase();
        if k.contains("ssn") || k.contains("social_security") || k.contains("social-security") {
            "ssn"
        } else if k.contains("card") || k.contains("cvv") || k.contains("cvc") {
            "card"
        } else if k.contains("email") {
            "email"
        } else if k.contains("phone") || k.contains("mobile") || k.contains("tel") {
            "phone"
        } else if k.contains("token")
            || k.contains("api_key")
            || k.contains("apikey")
            || k.contains("api-key")
            || k.contains("secret")
            || k.contains("private_key")
            || k.contains("credential")
        {
            "api_key"
        } else {
            "generic"
        }
    }

    fn redact_value(&self, value: &Value) -> Value {
        match value {
            Value::Object(map) => {
                let mut redacted_map = Map::new();

                for (key, val) in map {
                    let key_lower = key.to_lowercase();

                    if self.is_sensitive_key(&key_lower) {
                        let masked =
                            match val {
                                Value::String(s) => {
                                    let dt = Self::data_type_for_key(&key_lower);
                                    Value::String(self.mask_value(s, dt))
                                }
                                _ => Value::String(self.mask_value(
                                    &val.to_string(),
                                    Self::data_type_for_key(&key_lower),
                                )),
                            };
                        redacted_map.insert(key.clone(), masked);
                    } else {
                        let redacted_val = match val {
                            Value::String(s) => {
                                if self.is_sensitive_value(s) {
                                    let dt = self.detect_value_type(s);
                                    Value::String(self.mask_value(s, dt))
                                } else {
                                    Value::String(self.redact_patterns_in_string(s))
                                }
                            }
                            _ => self.redact_value(val),
                        };
                        redacted_map.insert(key.clone(), redacted_val);
                    }
                }

                Value::Object(redacted_map)
            }
            Value::Array(arr) => Value::Array(arr.iter().map(|v| self.redact_value(v)).collect()),
            Value::String(s) => {
                if self.is_sensitive_value(s) {
                    let dt = self.detect_value_type(s);
                    Value::String(self.mask_value(s, dt))
                } else {
                    Value::String(self.redact_patterns_in_string(s))
                }
            }
            _ => value.clone(),
        }
    }

    /// Detect the data type of a sensitive value by inspecting its content
    fn detect_value_type(&self, value: &str) -> &'static str {
        if JWT_PATTERN.is_match(value) {
            return "jwt";
        }
        for prefix in &TOKEN_PREFIXES {
            if value.starts_with(prefix) {
                return "api_key";
            }
        }
        "generic"
    }

    pub(crate) fn is_sensitive_key(&self, key: &str) -> bool {
        // Check exact matches
        if self.sensitive_keys.contains(key) {
            return true;
        }

        // Check built-in patterns
        for pattern in SENSITIVE_KEY_PATTERNS.iter() {
            if pattern.is_match(key) {
                return true;
            }
        }

        // Check custom patterns
        for pattern in &self.custom_patterns {
            if pattern.is_match(key) {
                return true;
            }
        }

        false
    }

    pub(crate) fn is_sensitive_value(&self, value: &str) -> bool {
        // Skip short values
        if value.len() < 8 {
            return false;
        }

        // Check for JWT
        if JWT_PATTERN.is_match(value) {
            return true;
        }

        // Check for known token prefixes
        for prefix in &TOKEN_PREFIXES {
            if value.starts_with(prefix) {
                return true;
            }
        }

        // Check custom patterns against value
        for pattern in &self.custom_patterns {
            if pattern.is_match(value) {
                return true;
            }
        }

        false
    }

    fn redact_patterns_in_string(&self, value: &str) -> String {
        let mut result = value.to_string();

        // Redact SSN patterns (123-45-6789)
        if SSN_PATTERN.is_match(&result) {
            if self.masking_mode == MaskingMode::Partial {
                result = SSN_PATTERN
                    .replace_all(&result, |caps: &regex::Captures| {
                        format!("***-**-{}", &caps[3])
                    })
                    .to_string();
            } else {
                result = SSN_PATTERN
                    .replace_all(&result, "[SSN-REDACTED]")
                    .to_string();
            }
        }

        // Redact credit card patterns
        if CARD_PATTERN.is_match(&result) {
            if self.masking_mode == MaskingMode::Partial {
                result = CARD_PATTERN
                    .replace_all(&result, |caps: &regex::Captures| {
                        format!("****-****-****-{}", &caps[4])
                    })
                    .to_string();
            } else {
                result = CARD_PATTERN
                    .replace_all(&result, "[CARD-REDACTED]")
                    .to_string();
            }
        }

        // Redact email patterns
        if EMAIL_PATTERN.is_match(&result) {
            if self.masking_mode == MaskingMode::Partial {
                result = EMAIL_PATTERN
                    .replace_all(&result, |caps: &regex::Captures| {
                        let full = &caps[0];
                        self.mask_value(full, "email")
                    })
                    .to_string();
            } else {
                result = EMAIL_PATTERN
                    .replace_all(&result, "[EMAIL-REDACTED]")
                    .to_string();
            }
        }

        // Redact phone patterns
        if PHONE_PATTERN.is_match(&result) {
            // Avoid matching fragments already masked (e.g. card last-4)
            if self.masking_mode == MaskingMode::Partial {
                result = PHONE_PATTERN
                    .replace_all(&result, |caps: &regex::Captures| {
                        format!("***-***-{}", &caps[3])
                    })
                    .to_string();
            } else {
                result = PHONE_PATTERN
                    .replace_all(&result, "[PHONE-REDACTED]")
                    .to_string();
            }
        }

        result
    }
}

/// Produce a generic partial mask: `****` + last 4 chars, or full mask for short values
fn mask_generic(value: &str) -> String {
    let len = value.len();
    if len <= 4 {
        return "****".to_string();
    }
    let suffix: String = value.chars().skip(len - 4).collect();
    format!("****{}", suffix)
}

impl Default for ResponseRedactor {
    fn default() -> Self {
        Self::new(vec![], false) // Disabled by default
    }
}

/// Convenience function to redact a response (defaults to Partial masking)
pub fn redact_response(response: &Value, enabled: bool) -> Value {
    let redactor = ResponseRedactor::new(vec![], enabled);
    redactor.redact(response)
}

/// Sanitize headers JSON for database storage.
/// Always enabled — never persist secrets in cleartext.
/// Preserves header names, redacts sensitive values.
/// Uses Full redaction for storage safety.
pub fn sanitize_headers_for_storage(headers: &Value) -> Value {
    let redactor = ResponseRedactor::new_with_mode(vec![], true, MaskingMode::Full);
    redactor.redact(headers)
}

/// Compute SHA-256 fingerprints for sensitive header values.
/// Returns a JSON object mapping header names to truncated hashes
/// for correlation (e.g., "did I use the same token across runs?").
/// Only includes entries for headers that contain sensitive values.
pub fn compute_secret_fingerprints(headers: &Value) -> Value {
    let redactor = ResponseRedactor::new_with_mode(vec![], true, MaskingMode::Full);
    let mut fingerprints = Map::new();

    if let Value::Object(map) = headers {
        for (key, val) in map {
            let key_lower = key.to_lowercase();
            let is_secret = redactor.is_sensitive_key(&key_lower)
                || matches!(val, Value::String(s) if redactor.is_sensitive_value(s));

            if is_secret {
                if let Value::String(raw) = val {
                    let mut hasher = Sha256::new();
                    hasher.update(raw.as_bytes());
                    let hash = hasher.finalize();
                    let hex = hex::encode(hash);
                    fingerprints.insert(key.clone(), Value::String(hex[..12].to_string()));
                }
            }
        }
    }

    Value::Object(fingerprints)
}

/// Sanitize a JSON response body for database storage.
/// Always enabled — never persist secrets in cleartext.
/// Uses Full redaction for storage safety.
pub fn sanitize_body_for_storage(body: &str) -> String {
    match serde_json::from_str::<Value>(body) {
        Ok(json_val) => {
            let redactor = ResponseRedactor::new_with_mode(vec![], true, MaskingMode::Full);
            let redacted = redactor.redact(&json_val);
            serde_json::to_string(&redacted).unwrap_or_else(|_| body.to_string())
        }
        Err(_) => body.to_string(), // non-JSON body, store as-is
    }
}

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

    #[test]
    fn test_redact_sensitive_keys() {
        let redactor = ResponseRedactor::default_enabled();

        let data = json!({
            "username": "john_doe",
            "password": "super_secret_123",
            "api_key": "sk_test_123456789",
            "data": {
                "token": "ghp_abcdef123456",
                "safe_field": "normal_value"
            }
        });

        let redacted = redactor.redact(&data);

        assert_eq!(redacted["username"], "john_doe");
        // Partial masking: password → generic mask (last 4)
        assert_ne!(redacted["password"].as_str().unwrap(), "super_secret_123");
        assert!(redacted["password"].as_str().unwrap().contains("****"));
        // api_key → token-style mask (first 4 + **** + last 4)
        let api_key_masked = redacted["api_key"].as_str().unwrap();
        assert!(api_key_masked.starts_with("sk_t"));
        assert!(api_key_masked.contains("****"));
        // token → token-style mask
        let token_masked = redacted["data"]["token"].as_str().unwrap();
        assert!(token_masked.contains("****"));
        assert_eq!(redacted["data"]["safe_field"], "normal_value");
    }

    #[test]
    fn test_redact_ssn_in_string() {
        let redactor = ResponseRedactor::default_enabled();

        // Dashed format
        let data = json!({
            "message": "Customer SSN is 123-45-6789 on file"
        });
        let redacted = redactor.redact(&data);
        let msg = redacted["message"].as_str().unwrap();
        assert!(msg.contains("***-**-6789"));
        assert!(!msg.contains("123-45"));

        // Undashed format (also matched)
        let data2 = json!({
            "message": "SSN: 123456789"
        });
        let redacted2 = redactor.redact(&data2);
        let msg2 = redacted2["message"].as_str().unwrap();
        assert!(
            msg2.contains("***-**-6789"),
            "Undashed SSN should be masked: {}",
            msg2
        );
        assert!(!msg2.contains("12345"));
    }

    #[test]
    fn test_redact_credit_card() {
        let redactor = ResponseRedactor::default_enabled();

        let data = json!({
            "card": "4111-1111-1111-1111"
        });

        let redacted = redactor.redact(&data);
        let card = redacted["card"].as_str().unwrap();
        // Partial masking: last 4 digits visible
        assert!(card.contains("****-****-****-1111"));
        assert!(!card.contains("4111-1111-1111-1111"));
    }

    #[test]
    fn test_redact_jwt() {
        let redactor = ResponseRedactor::default_enabled();

        let data = json!({
            "jwt": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U"
        });

        let redacted = redactor.redact(&data);
        let jwt_val = redacted["jwt"].as_str().unwrap();
        // Partial mode: JWT fingerprint
        assert!(jwt_val.starts_with("[JWT:"));
        assert!(jwt_val.ends_with(']'));
        assert_eq!(jwt_val.len(), 18); // [JWT: + 12 hex + ]
    }

    #[test]
    fn test_redact_disabled() {
        let redactor = ResponseRedactor::new(vec![], false);

        let data = json!({
            "password": "secret123"
        });

        let redacted = redactor.redact(&data);
        assert_eq!(redacted["password"], "secret123"); // Not redacted
    }

    #[test]
    fn test_redact_array() {
        let redactor = ResponseRedactor::default_enabled();

        let data = json!({
            "users": [
                {"name": "alice", "password": "pass1234"},
                {"name": "bob", "password": "pass5678"}
            ]
        });

        let redacted = redactor.redact(&data);
        assert_eq!(redacted["users"][0]["name"], "alice");
        assert!(redacted["users"][0]["password"]
            .as_str()
            .unwrap()
            .contains("****"));
        assert!(redacted["users"][1]["password"]
            .as_str()
            .unwrap()
            .contains("****"));
    }

    #[test]
    fn test_sanitize_headers_for_storage() {
        let headers = json!({
            "Authorization": "Bearer sk_test_abc123456789",
            "Accept": "application/json",
            "X-API-Key": "secret_key_value_12345",
            "Content-Type": "application/json",
            "User-Agent": "mrapids/1.0"
        });

        let sanitized = sanitize_headers_for_storage(&headers);
        // Storage always uses Full redaction
        assert_eq!(sanitized["Authorization"], "[REDACTED]");
        assert_eq!(sanitized["X-API-Key"], "[REDACTED]");
        assert_eq!(sanitized["Accept"], "application/json");
        assert_eq!(sanitized["Content-Type"], "application/json");
        assert_eq!(sanitized["User-Agent"], "mrapids/1.0");
    }

    #[test]
    fn test_compute_secret_fingerprints() {
        let headers = json!({
            "Authorization": "Bearer sk_test_abc123456789",
            "Accept": "application/json",
            "X-API-Key": "secret_key_value_12345"
        });

        let fingerprints = compute_secret_fingerprints(&headers);
        let fp_map = fingerprints.as_object().unwrap();

        // Should have fingerprints for sensitive headers only
        assert!(fp_map.contains_key("Authorization"));
        assert!(fp_map.contains_key("X-API-Key"));
        assert!(!fp_map.contains_key("Accept"));

        // Fingerprints should be 12 hex chars
        let auth_fp = fp_map["Authorization"].as_str().unwrap();
        assert_eq!(auth_fp.len(), 12);
        assert!(auth_fp.chars().all(|c| c.is_ascii_hexdigit()));
    }

    #[test]
    fn test_fingerprint_correlation() {
        let token = "Bearer sk_test_same_token_12345";
        let headers1 = json!({"Authorization": token});
        let headers2 = json!({"Authorization": token});
        let headers3 = json!({"Authorization": "Bearer sk_test_different_token"});

        let fp1 = compute_secret_fingerprints(&headers1);
        let fp2 = compute_secret_fingerprints(&headers2);
        let fp3 = compute_secret_fingerprints(&headers3);

        // Same token → same fingerprint
        assert_eq!(
            fp1["Authorization"].as_str().unwrap(),
            fp2["Authorization"].as_str().unwrap()
        );
        // Different token → different fingerprint
        assert_ne!(
            fp1["Authorization"].as_str().unwrap(),
            fp3["Authorization"].as_str().unwrap()
        );
    }

    #[test]
    fn test_sanitize_body_for_storage() {
        let body = r#"{"user": "alice", "password": "secret123", "email": "alice@test.com"}"#;
        let sanitized = sanitize_body_for_storage(body);
        let parsed: Value = serde_json::from_str(&sanitized).unwrap();

        assert_eq!(parsed["user"], "alice");
        // Storage uses Full mode
        assert_eq!(parsed["password"], "[REDACTED]");
    }

    #[test]
    fn test_sanitize_body_non_json() {
        let body = "plain text response body";
        let sanitized = sanitize_body_for_storage(body);
        assert_eq!(sanitized, body);
    }

    // --- New tests for partial masking ---

    #[test]
    fn test_partial_mask_email() {
        let redactor = ResponseRedactor::default_enabled();

        let data = json!({
            "message": "Contact alice.smith@example.com for details"
        });

        let redacted = redactor.redact(&data);
        let msg = redacted["message"].as_str().unwrap();
        // Should partially mask: first char + *** @ ***.tld
        assert!(msg.contains("a***@***.com"), "got: {}", msg);
        assert!(!msg.contains("alice.smith@example.com"));
    }

    #[test]
    fn test_partial_mask_phone() {
        let redactor = ResponseRedactor::default_enabled();

        let data = json!({
            "info": "Call 555-123-7890 for support"
        });

        let redacted = redactor.redact(&data);
        let info = redacted["info"].as_str().unwrap();
        assert!(info.contains("***-***-7890"), "got: {}", info);
        assert!(!info.contains("555-123"));
    }

    #[test]
    fn test_partial_mask_api_key() {
        let redactor = ResponseRedactor::default_enabled();

        let data = json!({
            "api_key": "sk_test_abc123xyz"
        });

        let redacted = redactor.redact(&data);
        let masked = redacted["api_key"].as_str().unwrap();
        // token-style: first 4 + **** + last 4
        assert!(masked.starts_with("sk_t"), "got: {}", masked);
        assert!(masked.contains("****"), "got: {}", masked);
        assert!(masked.ends_with("3xyz"), "got: {}", masked);
    }

    #[test]
    fn test_full_redaction_mode() {
        let redactor = ResponseRedactor::new_with_mode(vec![], true, MaskingMode::Full);

        let data = json!({
            "password": "super_secret_123",
            "api_key": "sk_test_123456789",
            "message": "SSN is 123-45-6789",
            "card": "4111-1111-1111-1111"
        });

        let redacted = redactor.redact(&data);
        assert_eq!(redacted["password"], "[REDACTED]");
        assert_eq!(redacted["api_key"], "[REDACTED]");
        assert!(redacted["message"]
            .as_str()
            .unwrap()
            .contains("[SSN-REDACTED]"));
        assert!(redacted["card"]
            .as_str()
            .unwrap()
            .contains("[CARD-REDACTED]"));
    }
}