a3s_code_core/security/
classifier.rs1use super::config::{ClassificationRule, RedactionStrategy, SensitivityLevel};
7
8pub use a3s_common::privacy::PiiMatch;
10
11#[derive(Debug, Clone)]
13pub struct ClassificationResult {
14 pub overall_level: SensitivityLevel,
16 pub matches: Vec<PiiMatch>,
18}
19
20pub struct PrivacyClassifier {
24 inner: a3s_common::privacy::RegexClassifier,
25}
26
27impl PrivacyClassifier {
28 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 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 pub fn redact(&self, text: &str, strategy: RedactionStrategy) -> String {
46 self.inner.redact(text, strategy)
47 }
48
49 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 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 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 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}