aaai_core/masking/
engine.rs

1//! Secret masking engine.
2//!
3//! Applies regex-based patterns to strings and replaces matched secrets with
4//! `***MASKED***`.  Used by the report generator and CLI output when
5//! `--mask-secrets` is active.
6
7use regex::Regex;
8
9use super::patterns::{BUILTIN_PATTERNS, SecretPattern};
10
11const MASK: &str = "***MASKED***";
12
13/// A compiled set of masking rules.
14pub struct MaskingEngine {
15    rules: Vec<(Regex, Option<usize>)>,
16}
17
18impl MaskingEngine {
19    /// Build an engine from the built-in patterns only.
20    pub fn builtin() -> Self {
21        Self::from_patterns(BUILTIN_PATTERNS, &[])
22    }
23
24    /// Build an engine from built-in patterns plus custom regex strings.
25    pub fn with_custom(custom: &[String]) -> Self {
26        Self::from_patterns(BUILTIN_PATTERNS, custom)
27    }
28
29    fn from_patterns(builtin: &[SecretPattern], custom: &[String]) -> Self {
30        let mut rules = Vec::new();
31        for sp in builtin {
32            match Regex::new(sp.pattern) {
33                Ok(re) => rules.push((re, sp.value_group)),
34                Err(e) => log::warn!("Built-in mask pattern {:?} failed to compile: {e}", sp.name),
35            }
36        }
37        for pat in custom {
38            match Regex::new(pat) {
39                Ok(re) => rules.push((re, None)),
40                Err(e) => log::warn!("Custom mask pattern {:?} failed to compile: {e}", pat),
41            }
42        }
43        Self { rules }
44    }
45
46    /// Apply all masking rules to `text`, returning the masked version.
47    pub fn mask(&self, text: &str) -> String {
48        let mut result = text.to_string();
49        for (re, group) in &self.rules {
50            result = mask_with_regex(&result, re, *group);
51        }
52        result
53    }
54
55    /// Mask only if secrets are present; return `None` if no change.
56    pub fn mask_if_needed(&self, text: &str) -> Option<String> {
57        let masked = self.mask(text);
58        if masked == text { None } else { Some(masked) }
59    }
60}
61
62fn mask_with_regex(text: &str, re: &Regex, group: Option<usize>) -> String {
63    match group {
64        None => re.replace_all(text, MASK).to_string(),
65        Some(g) => {
66            let mut result = text.to_string();
67            // Process matches in reverse order so offsets remain valid.
68            let captures: Vec<_> = re.captures_iter(text).collect();
69            for cap in captures.iter().rev() {
70                if let Some(m) = cap.get(g) {
71                    result.replace_range(m.start()..m.end(), MASK);
72                }
73            }
74            result
75        }
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    fn engine() -> MaskingEngine { MaskingEngine::builtin() }
84
85    #[test]
86    fn masks_api_key_assignment() {
87        let text = r#"api_key = "sk-abcdefghijklmnop1234567890""#;
88        let masked = engine().mask(text);
89        assert!(masked.contains(MASK), "expected mask in: {masked}");
90        assert!(!masked.contains("sk-abcdefghijklmnop"), "key should be masked");
91    }
92
93    #[test]
94    fn masks_aws_access_key() {
95        let text = "AKIAIOSFODNN7EXAMPLE";
96        let masked = engine().mask(text);
97        assert!(masked.contains(MASK));
98    }
99
100    #[test]
101    fn safe_text_unchanged() {
102        let text = "port = 8080";
103        let masked = engine().mask(text);
104        assert_eq!(masked, text);
105    }
106
107    #[test]
108    fn mask_if_needed_returns_none_for_safe_text() {
109        assert!(engine().mask_if_needed("ordinary text").is_none());
110    }
111
112    #[test]
113    fn custom_pattern_applied() {
114        let engine = MaskingEngine::with_custom(&["SECRET_VALUE".to_string()]);
115        let masked = engine.mask("here is SECRET_VALUE exposed");
116        assert!(masked.contains(MASK));
117    }
118
119    #[test]
120    fn private_key_header_masked() {
121        let text = "-----BEGIN RSA PRIVATE KEY-----\nABC123\n-----END RSA PRIVATE KEY-----";
122        let masked = engine().mask(text);
123        assert!(masked.contains(MASK));
124    }
125
126    #[test]
127    fn github_token_masked() {
128        let text = "token = ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZabcdef123456";
129        let masked = engine().mask(text);
130        assert!(masked.contains(MASK));
131    }
132}
133// (tests already included in the module above)
aaai_core/masking/engine.rs

aaai_core/masking/
engine.rs
