Skip to main content

atman_runtime/
redact.rs

1use std::collections::HashSet;
2
3use regex::{Regex, RegexSet};
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
7pub enum RedactMode {
8    #[default]
9    Full,
10    Partial,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct RedactHit {
15    pub kind: String,
16    pub start: usize,
17    pub end: usize,
18}
19
20pub struct Redactor {
21    kinds: Vec<String>,
22    regexes: Vec<Regex>,
23    set: RegexSet,
24    mode: RedactMode,
25    allowlist: HashSet<String>,
26}
27
28impl Redactor {
29    pub fn builtin() -> Self {
30        Self::from_pairs(BUILTIN_PATTERNS, RedactMode::default())
31    }
32
33    pub fn from_pairs(pairs: &[(&str, &str)], mode: RedactMode) -> Self {
34        let mut kinds = Vec::with_capacity(pairs.len());
35        let mut regexes = Vec::with_capacity(pairs.len());
36        let mut src = Vec::with_capacity(pairs.len());
37        for (kind, pattern) in pairs {
38            let compiled = Regex::new(pattern)
39                .unwrap_or_else(|e| panic!("invalid redact regex `{kind}`: {e}"));
40            kinds.push((*kind).to_string());
41            regexes.push(compiled);
42            src.push(*pattern);
43        }
44        let set = RegexSet::new(&src).expect("build RegexSet");
45        Self {
46            kinds,
47            regexes,
48            set,
49            mode,
50            allowlist: HashSet::new(),
51        }
52    }
53
54    pub fn with_mode(mut self, mode: RedactMode) -> Self {
55        self.mode = mode;
56        self
57    }
58
59    pub fn with_allowlist(mut self, items: impl IntoIterator<Item = String>) -> Self {
60        self.allowlist = items.into_iter().collect();
61        self
62    }
63
64    pub fn is_enabled(&self) -> bool {
65        !self.regexes.is_empty()
66    }
67
68    pub fn scan(&self, text: &str) -> Vec<RedactHit> {
69        if !self.set.is_match(text) {
70            return Vec::new();
71        }
72        let mut hits = Vec::new();
73        for &idx in self.set.matches(text).iter().collect::<Vec<_>>().iter() {
74            let regex = &self.regexes[idx];
75            let kind = &self.kinds[idx];
76            for m in regex.find_iter(text) {
77                if self.allowlist.contains(m.as_str()) {
78                    continue;
79                }
80                hits.push(RedactHit {
81                    kind: kind.clone(),
82                    start: m.start(),
83                    end: m.end(),
84                });
85            }
86        }
87        hits.sort_by(|a, b| {
88            a.start
89                .cmp(&b.start)
90                .then_with(|| (b.end - b.start).cmp(&(a.end - a.start)))
91        });
92        dedup_overlaps(hits)
93    }
94
95    pub fn redact(&self, text: &str) -> (String, Vec<RedactHit>) {
96        let hits = self.scan(text);
97        if hits.is_empty() {
98            return (text.to_string(), hits);
99        }
100        let mut out = String::with_capacity(text.len());
101        let mut cursor = 0;
102        for hit in &hits {
103            out.push_str(&text[cursor..hit.start]);
104            let matched = &text[hit.start..hit.end];
105            out.push_str(&render_replacement(matched, &hit.kind, self.mode));
106            cursor = hit.end;
107        }
108        out.push_str(&text[cursor..]);
109        (out, hits)
110    }
111
112    pub fn redact_json(&self, value: &mut serde_json::Value) -> Vec<RedactHit> {
113        let mut hits = Vec::new();
114        walk_json(value, self, &mut hits);
115        hits
116    }
117}
118
119fn dedup_overlaps(mut hits: Vec<RedactHit>) -> Vec<RedactHit> {
120    let mut out: Vec<RedactHit> = Vec::with_capacity(hits.len());
121    for hit in hits.drain(..) {
122        if let Some(last) = out.last_mut()
123            && hit.start < last.end
124        {
125            if hit.end > last.end {
126                last.end = hit.end;
127            }
128            continue;
129        }
130        out.push(hit);
131    }
132    out
133}
134
135fn render_replacement(matched: &str, kind: &str, mode: RedactMode) -> String {
136    match mode {
137        RedactMode::Full => format!("<REDACTED:{kind}>"),
138        RedactMode::Partial => {
139            if matched.len() <= 8 {
140                format!("<REDACTED:{kind}>")
141            } else {
142                let head: String = matched.chars().take(3).collect();
143                let tail: String = matched
144                    .chars()
145                    .rev()
146                    .take(3)
147                    .collect::<String>()
148                    .chars()
149                    .rev()
150                    .collect();
151                format!("{head}***{tail} <REDACTED:{kind}>")
152            }
153        }
154    }
155}
156
157fn walk_json(value: &mut serde_json::Value, r: &Redactor, hits: &mut Vec<RedactHit>) {
158    match value {
159        serde_json::Value::String(s) => {
160            let (redacted, mut new_hits) = r.redact(s);
161            if !new_hits.is_empty() {
162                *s = redacted;
163                hits.append(&mut new_hits);
164            }
165        }
166        serde_json::Value::Array(arr) => {
167            for v in arr {
168                walk_json(v, r, hits);
169            }
170        }
171        serde_json::Value::Object(map) => {
172            for (_, v) in map {
173                walk_json(v, r, hits);
174            }
175        }
176        _ => {}
177    }
178}
179
180pub const BUILTIN_PATTERNS: &[(&str, &str)] = &[
181    ("anthropic_api_key", r"sk-ant-[A-Za-z0-9_-]{20,}"),
182    ("openai_api_key", r"sk-[A-Za-z0-9_-]{20,}"),
183    ("github_token", r"gh[pousr]_[A-Za-z0-9]{20,}"),
184    ("google_api_key", r"AIza[0-9A-Za-z_-]{35}"),
185    ("aws_access_key", r"AKIA[0-9A-Z]{16}"),
186    ("bearer_token", r"Bearer\s+[A-Za-z0-9\-_.]{20,}"),
187    ("private_key_header", r"-----BEGIN [A-Z ]+PRIVATE KEY-----"),
188    (
189        "jwt",
190        r"eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}",
191    ),
192    (
193        "email",
194        r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b",
195    ),
196    ("credit_card", r"\b(?:\d[ -]*?){13,16}\b"),
197];
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[test]
204    fn openai_key_is_replaced_with_full_marker() {
205        let r = Redactor::builtin();
206        let (out, hits) = r.redact("token=sk-abcdefghijklmnop1234567890 rest");
207        assert_eq!(hits.len(), 1);
208        assert_eq!(hits[0].kind, "openai_api_key");
209        assert!(out.contains("<REDACTED:openai_api_key>"), "out: {out}");
210        assert!(!out.contains("sk-abcdef"), "sensitive prefix leaked: {out}");
211    }
212
213    #[test]
214    fn anthropic_key_wins_over_openai_prefix() {
215        let r = Redactor::builtin();
216        let (out, hits) = r.redact("key=sk-ant-abcdefghijklmn12345678 end");
217        let kinds: Vec<_> = hits.iter().map(|h| h.kind.as_str()).collect();
218        assert!(kinds.contains(&"anthropic_api_key"), "kinds: {kinds:?}");
219        assert!(out.contains("<REDACTED:"), "out: {out}");
220    }
221
222    #[test]
223    fn multiple_patterns_in_one_string_are_all_replaced() {
224        let r = Redactor::builtin();
225        let input = "gh=ghp_abcdefghij1234567890xyzXYZ11 email=alice@example.com";
226        let (out, hits) = r.redact(input);
227        assert_eq!(hits.len(), 2, "hits: {hits:?}");
228        assert!(out.contains("<REDACTED:github_token>"));
229        assert!(out.contains("<REDACTED:email>"));
230    }
231
232    #[test]
233    fn partial_mode_shows_prefix_and_suffix() {
234        let r = Redactor::builtin().with_mode(RedactMode::Partial);
235        let (out, _) = r.redact("token=sk-abcdefghijklmnop1234567890 rest");
236        assert!(out.contains("sk-***"), "partial marker missing: {out}");
237        assert!(out.contains("<REDACTED:openai_api_key>"));
238        assert!(!out.contains("sk-abcdefghijklmnop"), "full leak: {out}");
239    }
240
241    #[test]
242    fn allowlist_lets_specific_matches_pass_through() {
243        let r =
244            Redactor::builtin().with_allowlist(["sk-test-fixture-value-1234567890".to_string()]);
245        let (out, hits) = r.redact("cfg=sk-test-fixture-value-1234567890");
246        assert!(hits.is_empty(), "allowlist should suppress: {hits:?}");
247        assert_eq!(out, "cfg=sk-test-fixture-value-1234567890");
248    }
249
250    #[test]
251    fn clean_input_is_left_alone() {
252        let r = Redactor::builtin();
253        let (out, hits) = r.redact("no secrets here, just prose");
254        assert!(hits.is_empty());
255        assert_eq!(out, "no secrets here, just prose");
256    }
257
258    #[test]
259    fn overlapping_matches_are_deduplicated() {
260        let r = Redactor::from_pairs(
261            &[("kind_a", r"foo\d+"), ("kind_b", r"foo123bar")],
262            RedactMode::Full,
263        );
264        let (out, hits) = r.redact("foo123bar tail");
265        assert_eq!(hits.len(), 1, "overlapping match should collapse: {hits:?}");
266        assert!(out.starts_with("<REDACTED:"), "out: {out}");
267    }
268
269    #[test]
270    fn redact_json_walks_strings_arrays_and_objects() {
271        let r = Redactor::builtin();
272        let mut v = serde_json::json!({
273            "token": "sk-abcdefghijklmn1234567890",
274            "notes": [
275                "safe text",
276                "email me at bob@example.com",
277                {"nested": "gh=ghp_abcdefghij1234567890xyzXYZ11"}
278            ]
279        });
280        let hits = r.redact_json(&mut v);
281        assert_eq!(hits.len(), 3, "hits: {hits:?}");
282        let flat = v.to_string();
283        assert!(flat.contains("<REDACTED:openai_api_key>"));
284        assert!(flat.contains("<REDACTED:email>"));
285        assert!(flat.contains("<REDACTED:github_token>"));
286        assert!(!flat.contains("bob@example.com"));
287    }
288}