1use crate::masker::deduplicate;
4use crate::{
5 apply_mask, CloakError, Confidence, EntityType, Locale, MaskStrategy, PiiEntity, Recognizer,
6 RecognizerRegistry, Result,
7};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::time::Instant;
11
12pub struct ScannerBuilder {
23 registry: RecognizerRegistry,
24 locale: Locale,
25 strategy: Option<MaskStrategy>,
26 min_confidence: Confidence,
27 allow_list: Vec<String>,
28 deny_list: Vec<String>,
29}
30
31impl Default for ScannerBuilder {
32 fn default() -> Self {
33 Self {
34 registry: RecognizerRegistry::new(),
35 locale: Locale::Universal,
36 strategy: Some(MaskStrategy::default()),
37 min_confidence: Confidence::ZERO,
38 allow_list: Vec::new(),
39 deny_list: Vec::new(),
40 }
41 }
42}
43
44impl ScannerBuilder {
45 #[must_use]
47 pub fn new() -> Self {
48 Self::default()
49 }
50
51 #[must_use]
53 pub fn from_registry(registry: RecognizerRegistry) -> Self {
54 Self {
55 registry,
56 ..Self::default()
57 }
58 }
59
60 #[must_use]
62 pub fn locale(mut self, locale: Locale) -> Self {
63 self.locale = locale;
64 self
65 }
66
67 #[must_use]
69 pub fn strategy(mut self, strategy: MaskStrategy) -> Self {
70 self.strategy = Some(strategy);
71 self
72 }
73
74 #[must_use]
76 pub fn without_masking(mut self) -> Self {
77 self.strategy = None;
78 self
79 }
80
81 #[must_use]
83 pub fn recognizer<R>(mut self, recognizer: R) -> Self
84 where
85 R: Recognizer + 'static,
86 {
87 self.registry.register(recognizer);
88 self
89 }
90
91 #[must_use]
93 pub fn boxed_recognizer(mut self, recognizer: Box<dyn Recognizer>) -> Self {
94 self.registry.register_boxed(recognizer);
95 self
96 }
97
98 pub fn min_confidence(mut self, confidence: f64) -> Result<Self> {
100 self.min_confidence = Confidence::new(confidence)?;
101 Ok(self)
102 }
103
104 #[must_use]
106 pub fn allow_list<I, S>(mut self, values: I) -> Self
107 where
108 I: IntoIterator<Item = S>,
109 S: Into<String>,
110 {
111 self.allow_list.extend(
112 values
113 .into_iter()
114 .map(Into::into)
115 .filter(|value| !value.is_empty()),
116 );
117 self
118 }
119
120 #[must_use]
122 pub fn deny_list<I, S>(mut self, values: I) -> Self
123 where
124 I: IntoIterator<Item = S>,
125 S: Into<String>,
126 {
127 self.deny_list.extend(
128 values
129 .into_iter()
130 .map(Into::into)
131 .filter(|value| !value.is_empty()),
132 );
133 self
134 }
135
136 pub fn build(self) -> Result<Scanner> {
138 if self.registry.is_empty() {
139 return Err(CloakError::NoRecognizers);
140 }
141
142 Ok(Scanner {
143 registry: self.registry,
144 locale: self.locale,
145 strategy: self.strategy,
146 min_confidence: self.min_confidence,
147 allow_list: self.allow_list,
148 deny_list: self.deny_list,
149 })
150 }
151}
152
153pub struct Scanner {
155 registry: RecognizerRegistry,
156 locale: Locale,
157 strategy: Option<MaskStrategy>,
158 min_confidence: Confidence,
159 allow_list: Vec<String>,
160 deny_list: Vec<String>,
161}
162
163impl Scanner {
164 #[must_use]
166 pub fn builder() -> ScannerBuilder {
167 ScannerBuilder::new()
168 }
169
170 pub fn scan(&self, text: &str) -> Result<ScanResult> {
172 let started = Instant::now();
173 let mut findings = self.registry.scan_locale(text, &self.locale);
174 findings.extend(deny_list_findings(text, &self.deny_list));
175 let allow_spans = literal_spans(text, &self.allow_list);
176 findings.retain(|finding| !allow_spans.iter().any(|span| span.overlaps(finding.span)));
177 findings.retain(|finding| finding.confidence >= self.min_confidence);
178 findings = deduplicate_for_reporting(&findings);
179 findings.sort_by_key(|finding| finding.span.start);
180
181 let masked_text = self
182 .strategy
183 .as_ref()
184 .map(|strategy| apply_mask(text, &findings, strategy))
185 .transpose()?;
186
187 let stats = ScanStats::from_findings(&findings, started.elapsed().as_millis(), text.len());
188
189 Ok(ScanResult {
190 findings,
191 masked_text,
192 stats,
193 })
194 }
195}
196
197fn deny_list_findings(text: &str, deny_list: &[String]) -> Vec<PiiEntity> {
198 literal_spans(text, deny_list)
199 .into_iter()
200 .map(|span| PiiEntity {
201 entity_type: EntityType::Custom("DenyList".to_string()),
202 span,
203 text: text[span.start..span.end].to_string(),
204 confidence: Confidence::ONE,
205 recognizer_id: "deny_list_v1".to_string(),
206 })
207 .collect()
208}
209
210fn literal_spans(text: &str, values: &[String]) -> Vec<crate::Span> {
211 let mut spans = Vec::new();
212 for value in values {
213 if value.is_empty() {
214 continue;
215 }
216 let mut search_start = 0;
217 while let Some(offset) = text[search_start..].find(value) {
218 let start = search_start + offset;
219 let end = start + value.len();
220 spans.push(crate::Span::new(start, end));
221 search_start = end;
222 }
223 }
224 spans
225}
226
227fn deduplicate_for_reporting(findings: &[PiiEntity]) -> Vec<PiiEntity> {
228 let mut sorted = findings.to_vec();
229 sorted.sort_by_key(|finding| (finding.span.start, std::cmp::Reverse(finding.span.end)));
230
231 let mut keep: Vec<PiiEntity> = Vec::with_capacity(sorted.len());
232 for finding in sorted {
233 if keep
234 .iter()
235 .any(|kept| should_preserve_nested_url_query(kept, &finding))
236 {
237 keep.push(finding);
238 continue;
239 }
240
241 if let Some(overlap_index) = keep
242 .iter()
243 .rposition(|kept| finding.span.overlaps(kept.span))
244 {
245 if should_keep_existing_url_query(&keep[overlap_index], &finding) {
246 continue;
247 }
248 if should_replace_with_url_query(&keep[overlap_index], &finding) {
249 keep[overlap_index] = finding;
250 continue;
251 }
252
253 let merged = deduplicate(&[keep[overlap_index].clone(), finding])
254 .into_iter()
255 .next()
256 .unwrap_or_else(|| keep[overlap_index].clone());
257 keep[overlap_index] = merged;
258 continue;
259 }
260 keep.push(finding);
261 }
262
263 keep
264}
265
266fn should_preserve_nested_url_query(outer: &PiiEntity, inner: &PiiEntity) -> bool {
267 outer.entity_type == EntityType::Url
268 && inner.recognizer_id.starts_with("url_query_")
269 && inner.span.start >= outer.span.start
270 && inner.span.end <= outer.span.end
271}
272
273fn should_keep_existing_url_query(existing: &PiiEntity, incoming: &PiiEntity) -> bool {
274 existing.span == incoming.span
275 && existing.entity_type == incoming.entity_type
276 && existing.recognizer_id.starts_with("url_query_")
277 && !incoming.recognizer_id.starts_with("url_query_")
278}
279
280fn should_replace_with_url_query(existing: &PiiEntity, incoming: &PiiEntity) -> bool {
281 existing.span == incoming.span
282 && existing.entity_type == incoming.entity_type
283 && !existing.recognizer_id.starts_with("url_query_")
284 && incoming.recognizer_id.starts_with("url_query_")
285}
286
287#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
289pub struct ScanResult {
290 pub findings: Vec<PiiEntity>,
292 pub masked_text: Option<String>,
294 pub stats: ScanStats,
296}
297
298#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
300pub struct ScanStats {
301 pub total_findings: usize,
303 pub findings_by_type: HashMap<EntityType, usize>,
305 pub scan_duration_ms: u64,
307 pub bytes_scanned: usize,
309}
310
311impl ScanStats {
312 fn from_findings(findings: &[PiiEntity], duration_ms: u128, bytes_scanned: usize) -> Self {
313 let mut findings_by_type = HashMap::new();
314 for finding in findings {
315 *findings_by_type
316 .entry(finding.entity_type.clone())
317 .or_insert(0) += 1;
318 }
319
320 Self {
321 total_findings: findings.len(),
322 findings_by_type,
323 scan_duration_ms: duration_ms.try_into().unwrap_or(u64::MAX),
324 bytes_scanned,
325 }
326 }
327}
328
329#[cfg(test)]
330mod tests {
331 use super::*;
332 use crate::Span;
333
334 struct EmailRecognizer;
335
336 impl Recognizer for EmailRecognizer {
337 fn id(&self) -> &str {
338 "email_test_v1"
339 }
340
341 fn entity_type(&self) -> EntityType {
342 EntityType::Email
343 }
344
345 fn supported_locales(&self) -> &[Locale] {
346 &[]
347 }
348
349 fn scan(&self, text: &str) -> Vec<PiiEntity> {
350 let Some(start) = text
351 .find('@')
352 .and_then(|at| text[..at].rfind(' ').map(|space| space + 1).or(Some(0)))
353 else {
354 return Vec::new();
355 };
356 let end = text[start..]
357 .find(' ')
358 .map_or(text.len(), |offset| start + offset);
359 vec![PiiEntity {
360 entity_type: EntityType::Email,
361 span: Span::new(start, end),
362 text: text[start..end].to_string(),
363 confidence: Confidence::new(0.95).unwrap(),
364 recognizer_id: self.id().to_string(),
365 }]
366 }
367 }
368
369 #[test]
370 fn test_scanner_builder_without_recognizers_errors() {
371 assert!(Scanner::builder().build().is_err());
372 }
373
374 #[test]
375 fn test_scanner_scan_returns_findings_and_masked_text() {
376 let scanner = Scanner::builder()
377 .recognizer(EmailRecognizer)
378 .build()
379 .unwrap();
380 let result = scanner.scan("Contact user@example.com").unwrap();
381 assert_eq!(result.findings.len(), 1);
382 assert_eq!(result.masked_text.as_deref(), Some("Contact [EMAIL]"));
383 }
384
385 #[test]
386 fn test_scanner_without_masking_returns_no_masked_text() {
387 let scanner = Scanner::builder()
388 .recognizer(EmailRecognizer)
389 .without_masking()
390 .build()
391 .unwrap();
392 let result = scanner.scan("Contact user@example.com").unwrap();
393 assert!(result.masked_text.is_none());
394 }
395
396 #[test]
397 fn test_scanner_min_confidence_filters_low_confidence_findings() {
398 let scanner = Scanner::builder()
399 .recognizer(EmailRecognizer)
400 .min_confidence(1.0)
401 .unwrap()
402 .build()
403 .unwrap();
404 let result = scanner.scan("Contact user@example.com").unwrap();
405 assert!(result.findings.is_empty());
406 }
407
408 #[test]
409 fn test_scanner_allow_list_suppresses_overlapping_findings() {
410 let scanner = Scanner::builder()
411 .recognizer(EmailRecognizer)
412 .allow_list(["user@example.com"])
413 .build()
414 .unwrap();
415 let result = scanner.scan("Contact user@example.com").unwrap();
416 assert!(result.findings.is_empty());
417 assert_eq!(
418 result.masked_text.as_deref(),
419 Some("Contact user@example.com")
420 );
421 }
422
423 #[test]
424 fn test_scanner_deny_list_adds_custom_finding() {
425 let scanner = Scanner::builder()
426 .recognizer(EmailRecognizer)
427 .deny_list(["PRJ-12345"])
428 .build()
429 .unwrap();
430 let result = scanner.scan("ticket PRJ-12345").unwrap();
431 assert_eq!(result.findings.len(), 1);
432 assert_eq!(
433 result.findings[0].entity_type,
434 EntityType::Custom("DenyList".to_string())
435 );
436 assert_eq!(result.masked_text.as_deref(), Some("ticket [DENYLIST]"));
437 }
438}