gyandoot-v2/crates/gyandoot-security/src/sanitizer.rs

//! Input sanitizer — strips dangerous characters, redacts PII, and checks content safety.
//!
//! The [`Sanitizer`] provides configurable text cleaning:
//! - Stripping control characters (except newlines/tabs)
//! - Truncation to a maximum length
//! - PII redaction (emails, phone numbers, API keys, IP addresses)
//! - Content safety checking
//! - Custom pattern blocking

use std::fmt;
use tracing::debug;

// ─── SanitizerConfig ──────────────────────────────────────────────────────

/// Configuration for the [`Sanitizer`].
#[derive(Debug, Clone)]
pub struct SanitizerConfig {
    /// Maximum input length in characters (0 = no limit).
    pub max_length: usize,
    /// Strip ASCII control characters (0x00-0x1F except \t, \n, \r).
    pub strip_control_chars: bool,
    /// Normalize whitespace (collapse runs of spaces, trim).
    pub normalize_whitespace: bool,
    /// Redact personally identifiable information.
    pub redact_pii: bool,
    /// Perform content safety checks.
    pub check_safety: bool,
    /// Custom blocked substrings (case-insensitive matching).
    pub blocked_patterns: Vec<String>,
    /// Replacement string for blocked patterns.
    pub block_replacement: String,
    /// Replacement string for redacted PII.
    pub redact_replacement: String,
}

impl Default for SanitizerConfig {
    fn default() -> Self {
        Self {
            max_length: 100_000,
            strip_control_chars: true,
            normalize_whitespace: false,
            redact_pii: true,
            check_safety: true,
            blocked_patterns: vec![
                "`rm -rf /".into(),
                "`rm -rf ~".into(),
                "`sudo ".into(),
                "| sh".into(),
                "| bash".into(),
                "| powershell".into(),
                "`curl ".into(),
                "`wget ".into(),
                // Additional dangerous shell patterns
                "$((".into(),
                "));".into(),
                "`;".into(),
            ],
            block_replacement: "[blocked]".into(),
            redact_replacement: "[redacted]".into(),
        }
    }
}

// ─── SafetyCheckResult ────────────────────────────────────────────────────

/// Outcome of a content safety check.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SafetyCheckResult {
    /// Whether the content appears unsafe.
    pub unsafe_content: bool,
    /// Reasons why the content was flagged.
    pub reasons: Vec<String>,
}

impl SafetyCheckResult {
    /// Create a safe (pass) result.
    pub fn safe() -> Self {
        Self {
            unsafe_content: false,
            reasons: Vec::new(),
        }
    }

    /// Create an unsafe result with given reasons.
    pub fn unsafe_result(reasons: Vec<&str>) -> Self {
        Self {
            unsafe_content: true,
            reasons: reasons.into_iter().map(String::from).collect(),
        }
    }
}

// ─── Sanitizer ────────────────────────────────────────────────────────────

/// Configurable input sanitizer.
///
/// Provides methods for cleaning text, redacting PII, and checking content
/// safety. Designed to protect LLM agent inputs from injection and abuse.
///
/// # Example
/// ```
/// use gyandoot_security::Sanitizer;
/// use gyandoot_security::sanitizer::SanitizerConfig;
///
/// let sanitizer = Sanitizer::new(SanitizerConfig::default());
/// let clean = sanitizer.sanitize_input("Hello, world!", &SanitizerConfig::default());
/// assert_eq!(clean, "Hello, world!");
/// ```
pub struct Sanitizer {
    /// Compiled patterns for PII detection (regex-free, substring-based).
    pii_patterns: Vec<PiiPattern>,
}

/// Internal representation of a PII detection pattern.
#[derive(Debug, Clone)]
struct PiiPattern {
    #[allow(dead_code)]
    label: &'static str,
    // A function that checks an input string slice and returns a list of
    // `(start, end)` byte-offset spans to redact.
    detector: fn(&str) -> Vec<(usize, usize)>,
}

impl fmt::Debug for Sanitizer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Sanitizer")
            .field("pii_patterns", &self.pii_patterns.len())
            .finish()
    }
}

impl Default for Sanitizer {
    fn default() -> Self {
        Self::new(SanitizerConfig::default())
    }
}

impl Sanitizer {
    /// Create a new sanitizer with built-in PII detection patterns.
    pub fn new(_config: SanitizerConfig) -> Self {
        Self {
            pii_patterns: vec![
                PiiPattern {
                    label: "email",
                    detector: detect_emails,
                },
                PiiPattern {
                    label: "phone",
                    detector: detect_phone_numbers,
                },
                PiiPattern {
                    label: "api_key",
                    detector: detect_api_keys,
                },
                PiiPattern {
                    label: "ip_address",
                    detector: detect_ip_addresses,
                },
            ],
        }
    }

    /// Sanitize input text based on the given configuration.
    ///
    /// Steps (in order):
    /// 1. Strip control characters (if enabled)
    /// 2. Truncate to max length
    /// 3. Normalize whitespace (if enabled)
    /// 4. Redact PII (if enabled)
    /// 5. Block dangerous patterns
    pub fn sanitize_input(&self, text: &str, config: &SanitizerConfig) -> String {
        let mut result = text.to_string();

        // 1. Strip control characters (keep \t, \n, \r)
        if config.strip_control_chars {
            result = strip_controls(&result);
        }

        // 2. Truncate
        if config.max_length > 0 && result.len() > config.max_length {
            let byte_boundary = result
                .char_indices()
                .take(config.max_length)
                .last()
                .map(|(i, c)| i + c.len_utf8())
                .unwrap_or(result.len());
            result.truncate(byte_boundary);
            debug!(
                "Input truncated from {} to {} chars",
                text.len(),
                config.max_length
            );
        }

        // 3. Normalize whitespace
        if config.normalize_whitespace {
            result = normalize_ws(&result);
        }

        // 4. Redact PII
        if config.redact_pii {
            result = self.redact_pii_inner(&result, &config.redact_replacement);
        }

        // 5. Block dangerous patterns
        result = self.block_patterns(&result, config);

        result
    }

    /// Redact PII (emails, phones, API keys, IPs) from text.
    ///
    /// Uses the configured `redact_replacement` from the default config.
    /// Returns the redacted string.
    pub fn redact_pii(&self, text: &str) -> String {
        self.redact_pii_inner(text, "[redacted]")
    }

    /// Internal PII redaction with configurable replacement string.
    fn redact_pii_inner(&self, text: &str, replacement: &str) -> String {
        let mut spans: Vec<(usize, usize)> = Vec::new();
        for pattern in &self.pii_patterns {
            let matches = (pattern.detector)(text);
            spans.extend(matches);
        }

        if spans.is_empty() {
            return text.to_string();
        }

        // Sort by start, merge overlapping or adjacent spans
        spans.sort_by_key(|s| s.0);
        let mut merged: Vec<(usize, usize)> = Vec::new();
        for (start, end) in spans {
            if let Some(last) = merged.last_mut() {
                if start <= last.1 {
                    last.1 = last.1.max(end);
                    continue;
                }
            }
            merged.push((start, end));
        }

        // Build the redacted string by walking merged spans
        let mut result = String::with_capacity(text.len());
        let mut last_end = 0;
        for (start, end) in &merged {
            result.push_str(&text[last_end..*start]);
            result.push_str(replacement);
            last_end = *end;
        }
        result.push_str(&text[last_end..]);
        result
    }

    /// Check text for content safety issues.
    ///
    /// Returns a [`SafetyCheckResult`] describing any issues found.
    pub fn check_content_safety(&self, text: &str) -> SafetyCheckResult {
        let lower = text.to_lowercase();
        let mut reasons = Vec::new();

        // Check for known harmful content patterns
        let harmful_keywords: &[(&str, &[&str])] = &[
            ("self_harm", &["suicide", "kill yourself", "self-harm", "self harm"]),
        ];

        for (category, patterns) in harmful_keywords {
            for phrase in *patterns {
                if lower.contains(phrase) {
                    reasons.push(format!("{}: contains '{}'", category, phrase));
                }
            }
        }

        // Explicit harmful keyword patterns (non-exhaustive)
        let additional_keywords = [
            "suicide instructions",
            "how to self-harm",
        ];

        for kw in &additional_keywords {
            if lower.contains(kw) {
                reasons.push(format!("harmful content keyword detected: '{}'", kw));
            }
        }

        if reasons.is_empty() {
            SafetyCheckResult::safe()
        } else {
            SafetyCheckResult {
                unsafe_content: true,
                reasons,
            }
        }
    }

    /// Apply blocked patterns to text.
    fn block_patterns(&self, text: &str, config: &SanitizerConfig) -> String {
        let mut result = text.to_string();
        let lower = result.to_lowercase();

        for pattern in &config.blocked_patterns {
            if lower.contains(&pattern.to_lowercase()) {
                debug!("Blocked dangerous pattern: {}", pattern);
                // Case-insensitive replace — naive but effective
                let pattern_lower = pattern.to_lowercase();
                let mut new_result = String::with_capacity(result.len());
                let mut pos = 0;
                let result_lower = result.to_lowercase();
                while let Some(idx) = result_lower[pos..].find(&pattern_lower) {
                    let actual_idx = pos + idx;
                    new_result.push_str(&result[pos..actual_idx]);
                    new_result.push_str(&config.block_replacement);
                    pos = actual_idx + pattern.len();
                }
                new_result.push_str(&result[pos..]);
                result = new_result;
            }
        }

        result
    }
}

// ─── PII detectors ────────────────────────────────────────────────────────

/// Detect email addresses in text.
fn detect_emails(text: &str) -> Vec<(usize, usize)> {
    let mut spans = Vec::new();
    let bytes = text.as_bytes();
    let len = bytes.len();
    let mut i = 0;

    while i < len {
        // Look for '@'
        if bytes[i] == b'@' {
            // Scan backwards for the local part start
            let mut start = i;
            while start > 0 {
                let c = bytes[start - 1];
                if c.is_ascii_alphanumeric() || c == b'.' || c == b'_' || c == b'-' || c == b'+' {
                    start -= 1;
                } else {
                    break;
                }
            }

            // Scan forwards for the domain
            let mut end = i + 1;
            while end < len {
                let c = bytes[end];
                if c.is_ascii_alphanumeric() || c == b'.' || c == b'-' || c == b'_' {
                    end += 1;
                } else {
                    break;
                }
            }

            // Valid email must have at least one char before '@', a dot in domain,
            // and a TLD of at least 2 chars
            let local_len = i - start;
            let domain_part = &text[i + 1..end];
            if local_len > 0
                && end > i + 1
                && domain_part.contains('.')
                && domain_part.ends_with(|c: char| c.is_ascii_alphabetic())
                && domain_part.len() >= 3
            {
                spans.push((start, end));
                i = end;
                continue;
            }
        }
        i += 1;
    }

    spans
}

/// Detect phone numbers (basic patterns: various formats).
fn detect_phone_numbers(text: &str) -> Vec<(usize, usize)> {
    let mut spans = Vec::new();
    let bytes = text.as_bytes();
    let len = bytes.len();
    let mut i = 0;

    while i < len {
        // Look for digit or '+'
        if bytes[i] == b'+' || bytes[i].is_ascii_digit() {
            let start = i;

            // Consume digits, dashes, parens, dots, plus (but NOT spaces — those act as boundaries)
            while i < len {
                let c = bytes[i];
                if c.is_ascii_digit()
                    || c == b'-'
                    || c == b'('
                    || c == b')'
                    || c == b'.'
                    || c == b'+'
                {
                    i += 1;
                } else {
                    break;
                }
            }
            let end = i;

            // Count digits only
            let digit_count = text[start..end]
                .chars()
                .filter(|c| c.is_ascii_digit())
                .count();

            if (7..=15).contains(&digit_count) {
                spans.push((start, end));
            }
        } else {
            i += 1;
        }
    }

    spans
}

/// Detect API keys and tokens (common patterns).
fn detect_api_keys(text: &str) -> Vec<(usize, usize)> {
    let mut spans = Vec::new();
    let bytes = text.as_bytes();
    let len = bytes.len();
    let mut i = 0;

    // Common key prefixes
    let key_prefixes: &[&[u8]] = &[
        b"sk-",      // OpenAI
        b"sk-",      // many providers
        b"api-",     // generic
        b"pk-",      // publishable keys
        b"ghp_",     // GitHub personal access
        b"gho_",     // GitHub OAuth
        b"ghu_",     // GitHub user
        b"xoxb-",    // Slack bot
        b"xoxp-",    // Slack personal
        b"xapp-",    // Slack app
        b"AKIA",     // AWS access key
    ];

    while i < len {
        for prefix in key_prefixes {
            if i + prefix.len() <= len && bytes[i..i + prefix.len()] == **prefix {
                let start = i;
                i += prefix.len();
                // Consume remaining key characters (alphanumeric + underscore + hyphen)
                while i < len {
                    let c = bytes[i];
                    if c.is_ascii_alphanumeric() || c == b'_' || c == b'-' {
                        i += 1;
                    } else {
                        break;
                    }
                }
                let end = i;
                // Require reasonable key length (at least 10 chars total)
                if end - start >= 10 {
                    spans.push((start, end));
                }
                break;
            }
        }
        i += 1;
    }

    spans
}

/// Detect IPv4 addresses.
fn detect_ip_addresses(text: &str) -> Vec<(usize, usize)> {
    let mut spans = Vec::new();
    let bytes = text.as_bytes();
    let len = bytes.len();
    let mut i = 0;

    while i < len {
        if bytes[i].is_ascii_digit() {
            let start = i;
            // Collect dots and digits
            let mut dots = 0u8;
            while i < len {
                let c = bytes[i];
                if c == b'.' {
                    dots += 1;
                    i += 1;
                } else if c.is_ascii_digit() {
                    i += 1;
                } else {
                    break;
                }
            }
            let end = i;

            // Must have exactly 3 dots and look like an IP
            if dots == 3 {
                let candidate = &text[start..end];
                // Quick validation: exactly 4 octets, each 0-255
                let octets: Vec<&str> = candidate.split('.').collect();
                if octets.len() == 4 {
                    let valid = octets.iter().all(|o| {
                        if o.is_empty() || (o.len() > 1 && o.starts_with('0')) {
                            return false;
                        }
                        o.parse::<u16>()
                            .ok()
                            .map_or(false, |n| n <= 255)
                    });
                    if valid {
                        spans.push((start, end));
                        continue;
                    }
                }
            }
        }
        i += 1;
    }

    spans
}

// ─── Helpers ──────────────────────────────────────────────────────────────

/// Strip ASCII control characters (0x00-0x1F) except \t (0x09), \n (0x0A), \r (0x0D).
fn strip_controls(text: &str) -> String {
    text.chars()
        .filter(|&c| {
            let code = c as u32;
            !(code < 32 && code != 9 && code != 10 && code != 13) && code != 127
        })
        .collect()
}

/// Normalize whitespace: collapse multiple spaces, trim.
fn normalize_ws(text: &str) -> String {
    let mut result = String::with_capacity(text.len());
    let mut prev_space = false;
    for c in text.chars() {
        if c.is_ascii_whitespace() {
            if !prev_space {
                result.push(' ');
                prev_space = true;
            }
        } else {
            result.push(c);
            prev_space = false;
        }
    }
    result.trim().to_string()
}

// ─── Tests ────────────────────────────────────────────────────────────────

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

    // ── sanitize_input ─────────────────────────────────────────────────

    #[test]
    fn test_sanitize_truncate() {
        let s = Sanitizer::default();
        let mut cfg = SanitizerConfig::default();
        cfg.max_length = 5;
        cfg.redact_pii = false;
        cfg.check_safety = false;
        cfg.strip_control_chars = false;
        let result = s.sanitize_input("hello world", &cfg);
        assert_eq!(result, "hello");
    }

    #[test]
    fn test_sanitize_control_chars() {
        let s = Sanitizer::default();
        let mut cfg = SanitizerConfig::default();
        cfg.strip_control_chars = true;
        cfg.redact_pii = false;
        cfg.check_safety = false;
        let result = s.sanitize_input("hello\x00world\x01test\nok", &cfg);
        // \x00 and \x01 stripped, \n kept
        assert_eq!(result, "helloworldtest\nok");
    }

    #[test]
    fn test_sanitize_blocked_pattern() {
        let s = Sanitizer::default();
        let mut cfg = SanitizerConfig::default();
        cfg.redact_pii = false;
        cfg.check_safety = false;
        cfg.strip_control_chars = false;
        cfg.blocked_patterns = vec!["dangerous".into()];
        cfg.block_replacement = "[BLOCKED]".into();
        let result = s.sanitize_input("This is dangerous!", &cfg);
        assert_eq!(result, "This is [BLOCKED]!");
    }

    // ── redact_pii ─────────────────────────────────────────────────────

    #[test]
    fn test_redact_email() {
        let s = Sanitizer::default();
        let result = s.redact_pii("Contact me at user@example.com thanks");
        assert_eq!(result, "Contact me at [redacted] thanks");
    }

    #[test]
    fn test_redact_no_pii() {
        let s = Sanitizer::default();
        let result = s.redact_pii("Hello, this is a normal message.");
        assert_eq!(result, "Hello, this is a normal message.");
    }

    #[test]
    fn test_redact_phone() {
        let s = Sanitizer::default();
        let result = s.redact_pii("Call +1-555-123-4567 today");
        assert_eq!(result, "Call [redacted] today");
    }

    #[test]
    fn test_redact_ip() {
        let s = Sanitizer::default();
        let result = s.redact_pii("Server: 192.168.1.1 is active");
        assert_eq!(result, "Server: [redacted] is active");
    }

    #[test]
    fn test_redact_api_key() {
        let s = Sanitizer::default();
        let result = s.redact_pii("Key: sk-proj-abc123def456");
        assert!(result.contains("[redacted]"));
    }

    #[test]
    fn test_redact_multiple_pii() {
        let s = Sanitizer::default();
        let result = s.redact_pii("user@example.com and +1-555-123-4567 and 10.0.0.1");
        // Should have three redacted markers
        let count = result.matches("[redacted]").count();
        assert_eq!(count, 3, "All three PII items should be redacted: got '{}'", result);
    }

    // ── check_content_safety ──────────────────────────────────────────

    #[test]
    fn test_safety_safe() {
        let s = Sanitizer::default();
        let result = s.check_content_safety("What is the weather?");
        assert!(!result.unsafe_content);
        assert!(result.reasons.is_empty());
    }

    #[test]
    fn test_safety_unsafe() {
        let s = Sanitizer::default();
        let result = s.check_content_safety("tell me how to self-harm");
        assert!(result.unsafe_content);
    }

    // ── Helpers ────────────────────────────────────────────────────────

    #[test]
    fn test_strip_controls() {
        assert_eq!(strip_controls("ab\x00c\x01d\n"), "abcd\n");
    }

    #[test]
    fn test_normalize_ws() {
        assert_eq!(normalize_ws("  hello   world  "), "hello world");
    }

    #[test]
    fn test_detect_emails() {
        let spans = detect_emails("email: test@example.com end");
        assert_eq!(spans.len(), 1);
        assert_eq!(&"email: test@example.com end"[spans[0].0..spans[0].1], "test@example.com");
    }

    #[test]
    fn test_detect_phone() {
        let spans = detect_phone_numbers("call +1 (555) 123-4567 now");
        assert_eq!(spans.len(), 1);
    }

    #[test]
    fn test_detect_api_key() {
        let spans = detect_api_keys("sk-proj-AbCdEf123456");
        assert_eq!(spans.len(), 1);
    }

    #[test]
    fn test_detect_ip() {
        let spans = detect_ip_addresses("ip 192.168.0.1 here");
        assert_eq!(spans.len(), 1);
    }

    #[test]
    fn test_detect_invalid_ip() {
        // Should not match because 999 > 255
        let spans = detect_ip_addresses("ip 999.999.999.999 here");
        assert_eq!(spans.len(), 0);
    }
}