Skip to main content

a3s_code_core/security/
classifier.rs

1//! Security Privacy Classifier
2//!
3//! Thin wrapper around `a3s_common::privacy::RegexClassifier` that preserves the
4//! existing a3s-code API while delegating to the shared implementation.
5
6use super::config::{ClassificationRule, RedactionStrategy, SensitivityLevel};
7
8// Re-export PiiMatch from shared crate for consumers that need it
9pub use a3s_common::privacy::PiiMatch;
10
11/// Result of classifying a piece of text
12#[derive(Debug, Clone)]
13pub struct ClassificationResult {
14    /// Overall highest sensitivity level found
15    pub overall_level: SensitivityLevel,
16    /// All matches found
17    pub matches: Vec<PiiMatch>,
18}
19
20/// Privacy classifier with pre-compiled regex rules.
21///
22/// Wraps `a3s_common::privacy::RegexClassifier` with the a3s-code-specific API.
23pub struct PrivacyClassifier {
24    inner: a3s_common::privacy::RegexClassifier,
25}
26
27impl PrivacyClassifier {
28    /// Create a new classifier from classification rules
29    pub fn new(rules: &[ClassificationRule]) -> Self {
30        let inner = a3s_common::privacy::RegexClassifier::new(rules, SensitivityLevel::Public)
31            .expect("default rules should always compile");
32        Self { inner }
33    }
34
35    /// Classify text and return all matches
36    pub fn classify(&self, text: &str) -> ClassificationResult {
37        let result = self.inner.classify(text);
38        ClassificationResult {
39            overall_level: result.overall_level,
40            matches: result.matches,
41        }
42    }
43
44    /// Redact all matches in text using the given strategy
45    pub fn redact(&self, text: &str, strategy: RedactionStrategy) -> String {
46        self.inner.redact(text, strategy)
47    }
48
49    /// Quick check: does the text contain any sensitive data?
50    pub fn contains_sensitive(&self, text: &str) -> bool {
51        self.inner.contains_sensitive(text)
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use crate::security::config::default_classification_rules;
59
60    fn make_classifier() -> PrivacyClassifier {
61        PrivacyClassifier::new(&default_classification_rules())
62    }
63
64    #[test]
65    fn test_detect_credit_card() {
66        let classifier = make_classifier();
67        let result = classifier.classify("My card is 4111-1111-1111-1111");
68        assert!(!result.matches.is_empty());
69        assert!(result.overall_level >= SensitivityLevel::HighlySensitive);
70    }
71
72    #[test]
73    fn test_detect_ssn() {
74        let classifier = make_classifier();
75        let result = classifier.classify("SSN: 123-45-6789");
76        assert!(!result.matches.is_empty());
77        let ssn_match = result.matches.iter().find(|m| m.rule_name == "ssn");
78        assert!(ssn_match.is_some());
79    }
80
81    #[test]
82    fn test_detect_email() {
83        let classifier = make_classifier();
84        let result = classifier.classify("Contact me at user@example.com");
85        assert!(!result.matches.is_empty());
86        let email_match = result.matches.iter().find(|m| m.rule_name == "email");
87        assert!(email_match.is_some());
88    }
89
90    #[test]
91    fn test_detect_phone() {
92        let classifier = make_classifier();
93        // The phone regex expects digit-only prefix: \b\d{3}[-.\s]?\d{3}[-.\s]?\d{4}\b
94        let result = classifier.classify("Call me at 555-123-4567");
95        assert!(!result.matches.is_empty());
96    }
97
98    #[test]
99    fn test_detect_api_key() {
100        let classifier = make_classifier();
101        // The api_key regex requires 32+ alphanumeric/dash/underscore chars
102        let result = classifier.classify("Use key sk_test_0123456789abcdefghijklmnop");
103        assert!(!result.matches.is_empty());
104        assert!(result.overall_level >= SensitivityLevel::HighlySensitive);
105    }
106
107    #[test]
108    fn test_clean_text_no_matches() {
109        let classifier = make_classifier();
110        let result = classifier.classify("Hello, this is a normal message.");
111        assert!(result.matches.is_empty());
112        assert_eq!(result.overall_level, SensitivityLevel::Public);
113    }
114
115    #[test]
116    fn test_redact_remove() {
117        let classifier = make_classifier();
118        let redacted = classifier.redact("SSN: 123-45-6789", RedactionStrategy::Remove);
119        assert!(!redacted.contains("123-45-6789"));
120        // Remove strategy replaces with empty string
121        assert!(!redacted.contains("[REDACTED]"));
122    }
123
124    #[test]
125    fn test_redact_mask() {
126        let classifier = make_classifier();
127        let redacted = classifier.redact("SSN: 123-45-6789", RedactionStrategy::Mask);
128        assert!(!redacted.contains("123-45-6789"));
129        assert!(redacted.contains("***"));
130    }
131
132    #[test]
133    fn test_redact_hash() {
134        let classifier = make_classifier();
135        let redacted = classifier.redact("SSN: 123-45-6789", RedactionStrategy::Hash);
136        assert!(!redacted.contains("123-45-6789"));
137        assert!(redacted.contains("[HASH:"));
138    }
139
140    #[test]
141    fn test_contains_sensitive() {
142        let classifier = make_classifier();
143        assert!(classifier.contains_sensitive("SSN: 123-45-6789"));
144        assert!(!classifier.contains_sensitive("Hello world"));
145    }
146
147    #[test]
148    fn test_multiple_matches() {
149        let classifier = make_classifier();
150        let result = classifier.classify("SSN: 123-45-6789, email: test@example.com");
151        assert!(result.matches.len() >= 2);
152        assert_eq!(result.overall_level, SensitivityLevel::HighlySensitive);
153    }
154}