commons/src/data_retention.rs

use regex::Regex;
use std::sync::LazyLock;

/// Patterns for sensitive data that should be redacted from logs and event payloads.
static REDACTION_PATTERNS: LazyLock<Vec<(String, Regex)>> = LazyLock::new(|| {
    vec![
        (
            "api_key".to_string(),
            Regex::new(r#"(?i)(api[_-]?key|apikey)\s*[:=]\s*["']?([A-Za-z0-9\-_]{20,})["']?"#)
                .unwrap(),
        ),
        (
            "bearer_token".to_string(),
            Regex::new(
                r#"(?i)(bearer|token|authorization)\s*[:=]\s*["']?([A-Za-z0-9\-_\.]{20,})["']?"#,
            )
            .unwrap(),
        ),
        (
            "password".to_string(),
            Regex::new(r#"(?i)(password|passwd|pwd)\s*[:=]\s*["']?(\S{8,})["']?"#).unwrap(),
        ),
        (
            "private_key".to_string(),
            Regex::new(r#"-----BEGIN\s+(RSA\s+)?PRIVATE\s+KEY-----"#).unwrap(),
        ),
        (
            "connection_string".to_string(),
            Regex::new(r#"(?i)(mongodb|postgres|mysql|redis|amqp)://\S+"#).unwrap(),
        ),
        (
            "jwt_token".to_string(),
            Regex::new(r#"eyJ[A-Za-z0-9_-]{10,}\.eyJ[A-Za-z0-9_-]{10,}"#).unwrap(),
        ),
    ]
});

/// Redact sensitive information from a string.
///
/// Replaces matched patterns with `<REDACTED:{label}>` to prevent accidental
/// leakage of secrets through logs, event payloads, or error messages.
pub fn redact_secrets(input: &str) -> String {
    let mut result = input.to_string();
    for (label, pattern) in REDACTION_PATTERNS.iter() {
        result = pattern
            .replace_all(&result, format!("<REDACTED:{}>", label))
            .to_string();
    }
    result
}

/// Truncate a string to a maximum byte length, appending "..." if truncated.
pub fn truncate_bytes(input: &str, max_bytes: usize) -> String {
    let bytes = input.as_bytes();
    if bytes.len() <= max_bytes {
        input.to_string()
    } else {
        let truncated = String::from_utf8_lossy(&bytes[..max_bytes.saturating_sub(3)]);
        format!("{}...", truncated)
    }
}

/// Data classification for storage rules.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DataClass {
    /// Source code files — retained for the project lifetime.
    SourceCode,
    /// LLM prompts and system instructions — confidential, short retention.
    Prompt,
    /// Model-generated output — confidential, short retention.
    ModelOutput,
    /// Tool call arguments — confidential, may contain secrets.
    ToolArguments,
    /// Tool call results — confidential, may contain file content.
    ToolResult,
    /// User approval decisions — audit trail, long retention.
    ApprovalData,
    /// API keys, tokens, credentials — secret, never persisted in logs.
    Credentials,
    /// File mutation content (diffs, snapshots) — internal, medium retention.
    MutationContent,
    /// Operational traces (heartbeats, health checks) — internal, short retention.
    OperationalTrace,
    /// Domain events — internal, medium retention.
    DomainEvent,
    /// UI events (Lookout replay) — internal, short retention.
    UiEvent,
    /// Execution logs — confidential, short retention.
    ExecutionLog,
    /// Agent run records — internal, medium retention.
    AgentRun,
    /// Documentation output — source code, project lifetime.
    Documentation,
}

/// Retention configuration per data class.
#[derive(Debug, Clone)]
pub struct RetentionPolicy {
    /// Maximum number of days to retain processed events.
    pub event_retention_days: u32,
    /// Maximum number of days to retain completed agent runs.
    pub run_retention_days: u32,
    /// Maximum number of days to retain audit logs.
    pub audit_retention_days: u32,
    /// Maximum size of event payloads in bytes (larger payloads are truncated).
    pub max_event_payload_bytes: usize,
    /// Maximum number of days to retain execution logs.
    pub execution_log_retention_days: u32,
    /// Maximum number of days to retain UI replay events.
    pub ui_event_retention_days: u32,
    /// Maximum number of days to retain documentation.
    pub documentation_retention_days: u32,
}

impl Default for RetentionPolicy {
    fn default() -> Self {
        Self {
            event_retention_days: 30,
            run_retention_days: 90,
            audit_retention_days: 60,
            max_event_payload_bytes: 64 * 1024, // 64 KB
            execution_log_retention_days: 14,
            ui_event_retention_days: 7,
            documentation_retention_days: 365,
        }
    }
}

impl RetentionPolicy {
    /// Create a retention policy from environment variables with sensible defaults.
    pub fn from_env() -> Self {
        Self {
            event_retention_days: read_env_u32("SHOAL_EVENT_RETENTION_DAYS", 30),
            run_retention_days: read_env_u32("SHOAL_RUN_RETENTION_DAYS", 90),
            audit_retention_days: read_env_u32("SHOAL_AUDIT_RETENTION_DAYS", 60),
            max_event_payload_bytes: read_env_usize("SHOAL_MAX_EVENT_PAYLOAD_BYTES", 64 * 1024),
            execution_log_retention_days: read_env_u32("SHOAL_EXEC_LOG_RETENTION_DAYS", 14),
            ui_event_retention_days: read_env_u32("SHOAL_UI_EVENT_RETENTION_DAYS", 7),
            documentation_retention_days: read_env_u32("SHOAL_DOCS_RETENTION_DAYS", 365),
        }
    }

    /// Returns the retention period in days for a given data class.
    pub fn retention_days(&self, class: DataClass) -> u32 {
        match class {
            DataClass::SourceCode | DataClass::Documentation => {
                self.documentation_retention_days
            }
            DataClass::Prompt | DataClass::ModelOutput => 7,
            DataClass::ToolArguments | DataClass::ToolResult => {
                self.execution_log_retention_days
            }
            DataClass::ApprovalData | DataClass::AgentRun => self.run_retention_days,
            DataClass::Credentials => 0, // never retained
            DataClass::MutationContent => self.run_retention_days,
            DataClass::OperationalTrace => 3,
            DataClass::DomainEvent => self.event_retention_days,
            DataClass::UiEvent => self.ui_event_retention_days,
            DataClass::ExecutionLog => self.execution_log_retention_days,
        }
    }

    /// Whether the data class should be redacted before persistence.
    pub fn should_redact(&self, class: DataClass) -> bool {
        matches!(
            class,
            DataClass::ToolArguments
                | DataClass::ToolResult
                | DataClass::Credentials
                | DataClass::Prompt
                | DataClass::ModelOutput
                | DataClass::ExecutionLog
        )
    }
}

/// Applied retention policy service. Wraps a `RetentionPolicy` and provides
/// methods to enforce redaction and truncation on data before it reaches any
/// store (domain events, UI events, execution logs, agent runs, etc.).
pub struct DataRetentionService {
    policy: RetentionPolicy,
}

impl DataRetentionService {
    pub fn new(policy: RetentionPolicy) -> Self {
        Self { policy }
    }

    pub fn from_env() -> Self {
        Self::new(RetentionPolicy::from_env())
    }

    /// Returns a reference to the underlying policy.
    pub fn policy(&self) -> &RetentionPolicy {
        &self.policy
    }

    /// Apply redaction and truncation to a payload string based on its data class.
    pub fn prepare_for_storage(&self, data: &str, class: DataClass) -> String {
        let result = if self.policy.should_redact(class) {
            redact_secrets(data)
        } else {
            data.to_string()
        };
        let max_bytes = match class {
            DataClass::DomainEvent | DataClass::UiEvent => self.policy.max_event_payload_bytes,
            DataClass::ExecutionLog => 32 * 1024, // 32 KB
            DataClass::AgentRun => 128 * 1024,    // 128 KB
            _ => usize::MAX,
        };
        truncate_bytes(&result, max_bytes)
    }

    /// Prepare a JSON value for storage, redacting sensitive fields.
    pub fn prepare_json_for_storage(
        &self,
        value: &serde_json::Value,
        class: DataClass,
    ) -> serde_json::Value {
        if !self.policy.should_redact(class) {
            return value.clone();
        }
        redact_json_secrets(value)
    }

    /// Returns true if the given data class should be retained at all.
    /// Credentials, for example, should never be persisted.
    pub fn should_retain(&self, class: DataClass) -> bool {
        self.policy.retention_days(class) > 0
    }
}

/// Recursively redact sensitive fields in a JSON value.
fn redact_json_secrets(value: &serde_json::Value) -> serde_json::Value {
    match value {
        serde_json::Value::String(s) => serde_json::Value::String(redact_secrets(s)),
        serde_json::Value::Object(map) => {
            let redacted: serde_json::Map<String, serde_json::Value> = map
                .iter()
                .map(|(k, v)| {
                    let new_v = if is_sensitive_json_field(k) {
                        serde_json::Value::String("<REDACTED>".to_string())
                    } else {
                        redact_json_secrets(v)
                    };
                    (k.clone(), new_v)
                })
                .collect();
            serde_json::Value::Object(redacted)
        }
        serde_json::Value::Array(arr) => {
            serde_json::Value::Array(arr.iter().map(redact_json_secrets).collect())
        }
        other => other.clone(),
    }
}

fn is_sensitive_json_field(name: &str) -> bool {
    matches!(
        name.to_lowercase().as_str(),
        "password" | "secret" | "api_key" | "apikey" | "token" | "authorization"
            | "private_key" | "credentials" | "env" | "environment" | "bearer"
    )
}

fn read_env_u32(key: &str, default: u32) -> u32 {
    std::env::var(key)
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(default)
}

fn read_env_usize(key: &str, default: usize) -> usize {
    std::env::var(key)
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(default)
}

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

    #[test]
    fn test_redact_api_key() {
        let input = r#"{"api_key": "sk-abcdefghijklmnopqrstuvwxyz123456"}"#;
        let redacted = redact_secrets(input);
        assert!(redacted.contains("<REDACTED:api_key>"));
        assert!(!redacted.contains("sk-abcdefghijklmnopqrstuvwxyz123456"));
    }

    #[test]
    fn test_redact_bearer_token() {
        let input = "Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9";
        let redacted = redact_secrets(input);
        assert!(redacted.contains("<REDACTED:bearer_token>"));
    }

    #[test]
    fn test_redact_password() {
        let input = r#"password = "supersecret123""#;
        let redacted = redact_secrets(input);
        assert!(redacted.contains("<REDACTED:password>"));
        assert!(!redacted.contains("supersecret123"));
    }

    #[test]
    fn test_redact_private_key() {
        let input = "-----BEGIN RSA PRIVATE KEY-----";
        let redacted = redact_secrets(input);
        assert!(redacted.contains("<REDACTED:private_key>"));
    }

    #[test]
    fn test_no_redaction_for_clean_text() {
        let input = "This is a normal log message with no secrets.";
        let redacted = redact_secrets(input);
        assert_eq!(input, redacted);
    }

    #[test]
    fn test_truncate_bytes() {
        assert_eq!(truncate_bytes("hello", 10), "hello");
        assert_eq!(truncate_bytes("hello world", 8), "hello...");
        assert_eq!(truncate_bytes("", 5), "");
    }

    #[test]
    fn test_retention_policy_defaults() {
        let policy = RetentionPolicy::default();
        assert_eq!(policy.event_retention_days, 30);
        assert_eq!(policy.run_retention_days, 90);
        assert_eq!(policy.audit_retention_days, 60);
        assert_eq!(policy.max_event_payload_bytes, 64 * 1024);
        assert_eq!(policy.execution_log_retention_days, 14);
        assert_eq!(policy.ui_event_retention_days, 7);
        assert_eq!(policy.documentation_retention_days, 365);
    }

    #[test]
    fn test_data_retention_service_redacts_tool_args() {
        let svc = DataRetentionService::from_env();
        let input = r#"{"api_key": "sk-abcdefghijklmnopqrstuvwxyz123456"}"#;
        let result = svc.prepare_for_storage(input, DataClass::ToolArguments);
        assert!(result.contains("<REDACTED:api_key>"));
    }

    #[test]
    fn test_data_retention_service_no_redact_for_source_code() {
        let svc = DataRetentionService::from_env();
        let input = "fn main() { let api_key = \"test\"; }";
        let result = svc.prepare_for_storage(input, DataClass::SourceCode);
        assert_eq!(input, result);
    }

    #[test]
    fn test_data_retention_service_credentials_never_retained() {
        let svc = DataRetentionService::from_env();
        assert!(!svc.should_retain(DataClass::Credentials));
    }

    #[test]
    fn test_redact_json_secrets() {
        let value = serde_json::json!({
            "name": "test",
            "api_key": "sk-abcdefghijklmnopqrstuvwxyz123456",
            "nested": {
                "password": "hunter2"
            }
        });
        let redacted = redact_json_secrets(&value);
        assert_eq!(redacted["name"], "test");
        assert_eq!(redacted["api_key"], "<REDACTED>");
        assert_eq!(redacted["nested"]["password"], "<REDACTED>");
    }

    #[test]
    fn test_redact_connection_string() {
        let input = "Connecting to postgres://user:pass@localhost/db";
        let redacted = redact_secrets(input);
        assert!(redacted.contains("<REDACTED:connection_string>"));
    }

    #[test]
    fn test_retention_days_by_class() {
        let policy = RetentionPolicy::default();
        assert_eq!(policy.retention_days(DataClass::Credentials), 0);
        assert_eq!(policy.retention_days(DataClass::OperationalTrace), 3);
        assert_eq!(policy.retention_days(DataClass::Prompt), 7);
        assert_eq!(policy.retention_days(DataClass::DomainEvent), 30);
        assert_eq!(policy.retention_days(DataClass::AgentRun), 90);
    }
}