Skip to main content

trace_weft_core/
redactor.rs

1use crate::{RedactionResult, RedactionStatus, Redactor};
2use regex::Regex;
3use std::sync::Arc;
4
5pub struct RegexRedactor {
6    patterns: Vec<Regex>,
7}
8
9impl RegexRedactor {
10    pub fn new() -> Self {
11        Self { patterns: vec![] }
12    }
13
14    pub fn with_default_patterns() -> Self {
15        let mut redactor = Self::new();
16        // Basic examples: email, simple API keys
17        if let Ok(email_re) = Regex::new(r"(?i)[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}") {
18            redactor.patterns.push(email_re);
19        }
20        if let Ok(key_re) = Regex::new(r"(?i)(sk-[a-zA-Z0-9]{32,})") {
21            redactor.patterns.push(key_re);
22        }
23        if let Ok(bearer_re) = Regex::new(r"(?i)(bearer\s+[a-zA-Z0-9\-\._~+/\\]+)") {
24            redactor.patterns.push(bearer_re);
25        }
26        redactor
27    }
28
29    pub fn add_pattern(&mut self, pattern: Regex) {
30        self.patterns.push(pattern);
31    }
32}
33
34impl Default for RegexRedactor {
35    fn default() -> Self {
36        Self::with_default_patterns()
37    }
38}
39
40impl Redactor for RegexRedactor {
41    fn redact(&self, input: &str) -> RedactionResult {
42        if self.patterns.is_empty() {
43            return RedactionResult {
44                redacted_text: input.to_string(),
45                status: RedactionStatus::Unredacted,
46            };
47        }
48
49        let mut current_text = input.to_string();
50        let mut was_redacted = false;
51
52        for re in &self.patterns {
53            let replaced = re.replace_all(&current_text, "[REDACTED]");
54            if replaced != current_text {
55                was_redacted = true;
56                current_text = replaced.to_string();
57            }
58        }
59
60        RedactionResult {
61            redacted_text: current_text,
62            status: if was_redacted {
63                RedactionStatus::Redacted
64            } else {
65                RedactionStatus::Unredacted
66            },
67        }
68    }
69}
70
71pub type ArcRedactor = Arc<dyn Redactor>;
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    fn redactor() -> RegexRedactor {
78        RegexRedactor::with_default_patterns()
79    }
80
81    #[test]
82    fn redacts_email_addresses() {
83        let result = redactor().redact("contact me at jane.doe+test@example.co.uk please");
84        assert_eq!(result.redacted_text, "contact me at [REDACTED] please");
85        assert_eq!(result.status, RedactionStatus::Redacted);
86    }
87
88    #[test]
89    fn redacts_api_keys() {
90        let key = format!("sk-{}", "a1B2".repeat(10));
91        let result = redactor().redact(&format!("key={key}"));
92        assert_eq!(result.redacted_text, "key=[REDACTED]");
93        assert_eq!(result.status, RedactionStatus::Redacted);
94    }
95
96    #[test]
97    fn keeps_short_sk_prefixed_words() {
98        let result = redactor().redact("see sk-short for details");
99        assert_eq!(result.redacted_text, "see sk-short for details");
100        assert_eq!(result.status, RedactionStatus::Unredacted);
101    }
102
103    #[test]
104    fn redacts_bearer_tokens() {
105        let result = redactor().redact("Authorization: Bearer abc.DEF-123~xyz");
106        assert_eq!(result.redacted_text, "Authorization: [REDACTED]");
107        assert_eq!(result.status, RedactionStatus::Redacted);
108    }
109
110    #[test]
111    fn redacts_multiple_findings_in_one_text() {
112        let result = redactor().redact("a@b.com and Bearer tok123");
113        assert_eq!(result.redacted_text, "[REDACTED] and [REDACTED]");
114        assert_eq!(result.status, RedactionStatus::Redacted);
115    }
116
117    #[test]
118    fn leaves_clean_text_untouched() {
119        let result = redactor().redact("the quick brown fox");
120        assert_eq!(result.redacted_text, "the quick brown fox");
121        assert_eq!(result.status, RedactionStatus::Unredacted);
122    }
123
124    #[test]
125    fn empty_redactor_passes_content_through() {
126        let result = RegexRedactor::new().redact("a@b.com");
127        assert_eq!(result.redacted_text, "a@b.com");
128        assert_eq!(result.status, RedactionStatus::Unredacted);
129    }
130
131    #[test]
132    fn supports_user_configured_patterns() {
133        let mut redactor = RegexRedactor::new();
134        redactor.add_pattern(Regex::new(r"\b\d{3}-\d{2}-\d{4}\b").unwrap());
135        let result = redactor.redact("ssn: 123-45-6789");
136        assert_eq!(result.redacted_text, "ssn: [REDACTED]");
137        assert_eq!(result.status, RedactionStatus::Redacted);
138    }
139
140    #[test]
141    fn default_uses_builtin_patterns() {
142        let result = RegexRedactor::default().redact("a@b.com");
143        assert_eq!(result.status, RedactionStatus::Redacted);
144    }
145}