1use std::collections::HashMap;
8
9use serde::{Deserialize, Serialize};
10
11use crate::{DataMasker, MaskingRule};
12
13#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct MaskingCondition {
22 field: String,
24 expected_value: String,
26}
27
28impl MaskingCondition {
29 pub fn new(field: &str, expected_value: &str) -> Self {
31 Self {
32 field: field.to_string(),
33 expected_value: expected_value.to_string(),
34 }
35 }
36
37 pub fn field(&self) -> &str {
39 &self.field
40 }
41
42 pub fn expected_value(&self) -> &str {
44 &self.expected_value
45 }
46
47 pub fn is_satisfied(&self, data: &HashMap<String, String>) -> bool {
49 data.get(&self.field)
50 .map(|v| v == &self.expected_value)
51 .unwrap_or(false)
52 }
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct FieldMaskingRule {
65 field: String,
66 rule: MaskingRule,
67 condition: Option<MaskingCondition>,
68 priority: u32,
69}
70
71impl FieldMaskingRule {
72 pub fn new(field: &str, rule: MaskingRule) -> Self {
74 Self {
75 field: field.to_string(),
76 rule,
77 condition: None,
78 priority: 100,
79 }
80 }
81
82 pub fn with_condition(mut self, cond: MaskingCondition) -> Self {
84 self.condition = Some(cond);
85 self
86 }
87
88 pub fn with_priority(mut self, priority: u32) -> Self {
90 self.priority = priority;
91 self
92 }
93
94 pub fn field(&self) -> &str {
96 &self.field
97 }
98
99 pub fn rule(&self) -> &MaskingRule {
101 &self.rule
102 }
103
104 pub fn condition(&self) -> Option<&MaskingCondition> {
106 self.condition.as_ref()
107 }
108
109 pub fn priority(&self) -> u32 {
111 self.priority
112 }
113
114 pub fn should_apply(&self, data: &HashMap<String, String>) -> bool {
116 match &self.condition {
117 Some(cond) => cond.is_satisfied(data),
118 None => true,
119 }
120 }
121}
122
123#[derive(Debug, Clone, Default)]
132pub struct MaskingStrategyEngine {
133 rules: Vec<FieldMaskingRule>,
134}
135
136impl MaskingStrategyEngine {
137 pub fn new() -> Self {
139 Self::default()
140 }
141
142 pub fn add_rule(mut self, rule: FieldMaskingRule) -> Self {
144 self.rules.push(rule);
145 self
146 }
147
148 pub fn add_rules(mut self, rules: Vec<FieldMaskingRule>) -> Self {
150 self.rules.extend(rules);
151 self
152 }
153
154 pub fn rule_count(&self) -> usize {
156 self.rules.len()
157 }
158
159 pub fn clear(&mut self) {
161 self.rules.clear();
162 }
163
164 pub fn effective_rule(
166 &self,
167 field: &str,
168 data: &HashMap<String, String>,
169 ) -> Option<&MaskingRule> {
170 let mut candidates: Vec<&FieldMaskingRule> = self
171 .rules
172 .iter()
173 .filter(|r| r.field() == field && r.should_apply(data))
174 .collect();
175 candidates.sort_by_key(|r| r.priority());
176 candidates.first().map(|r| r.rule())
177 }
178
179 pub fn apply_to_map(&self, data: &HashMap<String, String>) -> HashMap<String, String> {
181 data.iter()
182 .map(|(k, v)| match self.effective_rule(k, data) {
183 Some(rule) => (k.clone(), DataMasker::apply(rule, v)),
184 None => (k.clone(), v.clone()),
185 })
186 .collect()
187 }
188
189 pub fn apply_to_json(&self, json: &str) -> String {
191 let Ok(mut value) = serde_json::from_str::<serde_json::Value>(json) else {
192 return json.to_string();
193 };
194 let Some(obj) = value.as_object_mut() else {
195 return json.to_string();
196 };
197 let snapshot: HashMap<String, String> = obj
199 .iter()
200 .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
201 .collect();
202 let keys: Vec<String> = obj.keys().cloned().collect();
203 for key in keys {
204 if let Some(rule) = self.effective_rule(&key, &snapshot) {
205 if let Some(serde_json::Value::String(s)) = obj.get_mut(&key) {
206 *s = DataMasker::apply(rule, s);
207 }
208 }
209 }
210 serde_json::to_string(&value).unwrap_or_else(|_| json.to_string())
211 }
212
213 pub fn remove_rules_for_field(&mut self, field: &str) -> usize {
215 let before = self.rules.len();
216 self.rules.retain(|r| r.field() != field);
217 before - self.rules.len()
218 }
219}
220
221#[derive(Debug, Clone)]
227pub struct PipelineStage {
228 name: String,
229 rules: HashMap<String, MaskingRule>,
230}
231
232impl PipelineStage {
233 pub fn new(name: &str) -> Self {
235 Self {
236 name: name.to_string(),
237 rules: HashMap::new(),
238 }
239 }
240
241 pub fn name(&self) -> &str {
243 &self.name
244 }
245
246 pub fn with_rule(mut self, field: &str, rule: MaskingRule) -> Self {
248 self.rules.insert(field.to_string(), rule);
249 self
250 }
251
252 pub fn rule_count(&self) -> usize {
254 self.rules.len()
255 }
256
257 pub fn apply(&self, data: &HashMap<String, String>) -> HashMap<String, String> {
259 DataMasker::mask_map(&self.rules, data)
260 }
261
262 pub fn apply_to_json(&self, json: &str) -> String {
264 DataMasker::mask_json(&self.rules, json)
265 }
266}
267
268#[derive(Debug, Clone, Default)]
273pub struct MaskingPipeline {
274 stages: Vec<PipelineStage>,
275}
276
277impl MaskingPipeline {
278 pub fn new() -> Self {
280 Self::default()
281 }
282
283 pub fn add_stage(mut self, stage: PipelineStage) -> Self {
285 self.stages.push(stage);
286 self
287 }
288
289 pub fn stage_count(&self) -> usize {
291 self.stages.len()
292 }
293
294 pub fn stage_names(&self) -> Vec<&str> {
296 self.stages.iter().map(|s| s.name()).collect()
297 }
298
299 pub fn apply_to_map(&self, data: &HashMap<String, String>) -> HashMap<String, String> {
301 let mut current = data.clone();
302 for stage in &self.stages {
303 current = stage.apply(¤t);
304 }
305 current
306 }
307
308 pub fn apply_to_json(&self, json: &str) -> String {
310 let mut current = json.to_string();
311 for stage in &self.stages {
312 current = stage.apply_to_json(¤t);
313 }
314 current
315 }
316
317 pub fn clear(&mut self) {
319 self.stages.clear();
320 }
321}
322
323#[derive(Debug, Clone, Default)]
329pub struct MaskingValidator;
330
331impl MaskingValidator {
332 pub fn new() -> Self {
334 Self
335 }
336
337 pub fn is_masked(original: &str, masked: &str) -> bool {
339 original != masked || original.is_empty()
340 }
341
342 pub fn contains_mask_char(masked: &str) -> bool {
344 masked.contains('*')
345 }
346
347 pub fn no_sensitive_substring(masked: &str, sensitive: &str) -> bool {
349 if sensitive.len() <= 2 {
350 return true;
351 }
352 !masked.contains(sensitive)
353 }
354
355 pub fn validate_map(
357 data: &HashMap<String, String>,
358 masked: &HashMap<String, String>,
359 ) -> Vec<String> {
360 data.iter()
361 .filter(|(k, v)| masked.get(*k).map(|m| m == *v).unwrap_or(true) && !v.is_empty())
362 .map(|(k, _)| k.clone())
363 .collect()
364 }
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370
371 #[test]
374 fn condition_new() {
375 let c = MaskingCondition::new("is_vip", "false");
376 assert_eq!(c.field(), "is_vip");
377 assert_eq!(c.expected_value(), "false");
378 }
379
380 #[test]
381 fn condition_satisfied() {
382 let c = MaskingCondition::new("is_vip", "false");
383 let mut data = HashMap::new();
384 data.insert("is_vip".to_string(), "false".to_string());
385 assert!(c.is_satisfied(&data));
386 }
387
388 #[test]
389 fn condition_not_satisfied() {
390 let c = MaskingCondition::new("is_vip", "false");
391 let mut data = HashMap::new();
392 data.insert("is_vip".to_string(), "true".to_string());
393 assert!(!c.is_satisfied(&data));
394 }
395
396 #[test]
397 fn condition_field_missing() {
398 let c = MaskingCondition::new("is_vip", "false");
399 let data = HashMap::new();
400 assert!(!c.is_satisfied(&data));
401 }
402
403 #[test]
406 fn field_rule_new() {
407 let r = FieldMaskingRule::new("phone", MaskingRule::Phone);
408 assert_eq!(r.field(), "phone");
409 assert_eq!(r.rule(), &MaskingRule::Phone);
410 assert!(r.condition().is_none());
411 assert_eq!(r.priority(), 100);
412 }
413
414 #[test]
415 fn field_rule_with_condition() {
416 let r = FieldMaskingRule::new("phone", MaskingRule::Phone)
417 .with_condition(MaskingCondition::new("is_vip", "false"));
418 assert!(r.condition().is_some());
419 }
420
421 #[test]
422 fn field_rule_with_priority() {
423 let r = FieldMaskingRule::new("phone", MaskingRule::Phone).with_priority(10);
424 assert_eq!(r.priority(), 10);
425 }
426
427 #[test]
428 fn field_rule_should_apply_no_condition() {
429 let r = FieldMaskingRule::new("phone", MaskingRule::Phone);
430 let data = HashMap::new();
431 assert!(r.should_apply(&data));
432 }
433
434 #[test]
435 fn field_rule_should_apply_condition_met() {
436 let r = FieldMaskingRule::new("phone", MaskingRule::Phone)
437 .with_condition(MaskingCondition::new("is_vip", "false"));
438 let mut data = HashMap::new();
439 data.insert("is_vip".to_string(), "false".to_string());
440 assert!(r.should_apply(&data));
441 }
442
443 #[test]
444 fn field_rule_should_apply_condition_not_met() {
445 let r = FieldMaskingRule::new("phone", MaskingRule::Phone)
446 .with_condition(MaskingCondition::new("is_vip", "false"));
447 let mut data = HashMap::new();
448 data.insert("is_vip".to_string(), "true".to_string());
449 assert!(!r.should_apply(&data));
450 }
451
452 #[test]
455 fn strategy_engine_default_empty() {
456 let e = MaskingStrategyEngine::new();
457 assert_eq!(e.rule_count(), 0);
458 }
459
460 #[test]
461 fn strategy_engine_add_rule() {
462 let e = MaskingStrategyEngine::new()
463 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Phone));
464 assert_eq!(e.rule_count(), 1);
465 }
466
467 #[test]
468 fn strategy_engine_add_rules_batch() {
469 let rules = vec![
470 FieldMaskingRule::new("phone", MaskingRule::Phone),
471 FieldMaskingRule::new("email", MaskingRule::Email),
472 ];
473 let e = MaskingStrategyEngine::new().add_rules(rules);
474 assert_eq!(e.rule_count(), 2);
475 }
476
477 #[test]
478 fn strategy_engine_apply_to_map() {
479 let e = MaskingStrategyEngine::new()
480 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Phone));
481 let mut data = HashMap::new();
482 data.insert("phone".to_string(), "13812345678".to_string());
483 data.insert("name".to_string(), "Alice".to_string());
484 let result = e.apply_to_map(&data);
485 assert_eq!(result["phone"], "138****5678");
486 assert_eq!(result["name"], "Alice");
487 }
488
489 #[test]
490 fn strategy_engine_apply_to_json() {
491 let e = MaskingStrategyEngine::new()
492 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Phone));
493 let json = r#"{"phone":"13812345678","name":"Alice"}"#;
494 let result = e.apply_to_json(json);
495 assert!(result.contains("138****5678"));
496 assert!(result.contains("Alice"));
497 }
498
499 #[test]
500 fn strategy_engine_apply_to_json_invalid() {
501 let e = MaskingStrategyEngine::new()
502 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Phone));
503 assert_eq!(e.apply_to_json("not json"), "not json");
504 }
505
506 #[test]
507 fn strategy_engine_conditional_masking() {
508 let e = MaskingStrategyEngine::new().add_rule(
509 FieldMaskingRule::new("phone", MaskingRule::Phone)
510 .with_condition(MaskingCondition::new("is_vip", "false")),
511 );
512 let mut data = HashMap::new();
513 data.insert("phone".to_string(), "13812345678".to_string());
514 data.insert("is_vip".to_string(), "true".to_string());
515 let result = e.apply_to_map(&data);
516 assert_eq!(result["phone"], "13812345678");
518 }
519
520 #[test]
521 fn strategy_engine_conditional_masking_applied() {
522 let e = MaskingStrategyEngine::new().add_rule(
523 FieldMaskingRule::new("phone", MaskingRule::Phone)
524 .with_condition(MaskingCondition::new("is_vip", "false")),
525 );
526 let mut data = HashMap::new();
527 data.insert("phone".to_string(), "13812345678".to_string());
528 data.insert("is_vip".to_string(), "false".to_string());
529 let result = e.apply_to_map(&data);
530 assert_eq!(result["phone"], "138****5678");
531 }
532
533 #[test]
534 fn strategy_engine_priority_resolution() {
535 let e = MaskingStrategyEngine::new()
536 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Password).with_priority(50))
537 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Phone).with_priority(10));
538 let mut data = HashMap::new();
539 data.insert("phone".to_string(), "13812345678".to_string());
540 let result = e.apply_to_map(&data);
541 assert_eq!(result["phone"], "138****5678");
543 }
544
545 #[test]
546 fn strategy_engine_remove_rules_for_field() {
547 let mut e = MaskingStrategyEngine::new()
548 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Phone))
549 .add_rule(FieldMaskingRule::new("email", MaskingRule::Email));
550 let removed = e.remove_rules_for_field("phone");
551 assert_eq!(removed, 1);
552 assert_eq!(e.rule_count(), 1);
553 }
554
555 #[test]
556 fn strategy_engine_clear() {
557 let mut e = MaskingStrategyEngine::new()
558 .add_rule(FieldMaskingRule::new("phone", MaskingRule::Phone));
559 e.clear();
560 assert_eq!(e.rule_count(), 0);
561 }
562
563 #[test]
564 fn strategy_engine_effective_rule_none() {
565 let e = MaskingStrategyEngine::new();
566 let data = HashMap::new();
567 assert!(e.effective_rule("phone", &data).is_none());
568 }
569
570 #[test]
573 fn pipeline_stage_new() {
574 let s = PipelineStage::new("stage1");
575 assert_eq!(s.name(), "stage1");
576 assert_eq!(s.rule_count(), 0);
577 }
578
579 #[test]
580 fn pipeline_stage_with_rule() {
581 let s = PipelineStage::new("stage1").with_rule("phone", MaskingRule::Phone);
582 assert_eq!(s.rule_count(), 1);
583 }
584
585 #[test]
586 fn pipeline_stage_apply() {
587 let s = PipelineStage::new("stage1").with_rule("phone", MaskingRule::Phone);
588 let mut data = HashMap::new();
589 data.insert("phone".to_string(), "13812345678".to_string());
590 let result = s.apply(&data);
591 assert_eq!(result["phone"], "138****5678");
592 }
593
594 #[test]
595 fn pipeline_stage_apply_to_json() {
596 let s = PipelineStage::new("stage1").with_rule("phone", MaskingRule::Phone);
597 let json = r#"{"phone":"13812345678"}"#;
598 let result = s.apply_to_json(json);
599 assert!(result.contains("138****5678"));
600 }
601
602 #[test]
605 fn pipeline_default_empty() {
606 let p = MaskingPipeline::new();
607 assert_eq!(p.stage_count(), 0);
608 }
609
610 #[test]
611 fn pipeline_add_stage() {
612 let p = MaskingPipeline::new().add_stage(PipelineStage::new("s1"));
613 assert_eq!(p.stage_count(), 1);
614 assert_eq!(p.stage_names(), vec!["s1"]);
615 }
616
617 #[test]
618 fn pipeline_apply_single_stage() {
619 let p = MaskingPipeline::new()
620 .add_stage(PipelineStage::new("mask").with_rule("phone", MaskingRule::Phone));
621 let mut data = HashMap::new();
622 data.insert("phone".to_string(), "13812345678".to_string());
623 let result = p.apply_to_map(&data);
624 assert_eq!(result["phone"], "138****5678");
625 }
626
627 #[test]
628 fn pipeline_apply_multi_stage() {
629 let p = MaskingPipeline::new()
630 .add_stage(
631 PipelineStage::new("type_mask")
632 .with_rule("phone", MaskingRule::Phone)
633 .with_rule("email", MaskingRule::Email),
634 )
635 .add_stage(PipelineStage::new("name_mask").with_rule("name", MaskingRule::Name));
636 let mut data = HashMap::new();
637 data.insert("phone".to_string(), "13812345678".to_string());
638 data.insert("email".to_string(), "test@example.com".to_string());
639 data.insert("name".to_string(), "Alice".to_string());
640 let result = p.apply_to_map(&data);
641 assert_eq!(result["phone"], "138****5678");
642 assert_eq!(result["email"], "t***@example.com");
643 assert_eq!(result["name"], "A****");
644 }
645
646 #[test]
647 fn pipeline_apply_to_json() {
648 let p = MaskingPipeline::new()
649 .add_stage(PipelineStage::new("mask").with_rule("phone", MaskingRule::Phone));
650 let json = r#"{"phone":"13812345678"}"#;
651 let result = p.apply_to_json(json);
652 assert!(result.contains("138****5678"));
653 }
654
655 #[test]
656 fn pipeline_apply_to_json_invalid() {
657 let p = MaskingPipeline::new().add_stage(PipelineStage::new("mask"));
658 assert_eq!(p.apply_to_json("not json"), "not json");
659 }
660
661 #[test]
662 fn pipeline_clear() {
663 let mut p = MaskingPipeline::new().add_stage(PipelineStage::new("s1"));
664 p.clear();
665 assert_eq!(p.stage_count(), 0);
666 }
667
668 #[test]
669 fn pipeline_empty_passthrough() {
670 let p = MaskingPipeline::new();
671 let mut data = HashMap::new();
672 data.insert("phone".to_string(), "13812345678".to_string());
673 let result = p.apply_to_map(&data);
674 assert_eq!(result["phone"], "13812345678");
675 }
676
677 #[test]
680 fn validator_is_masked() {
681 assert!(MaskingValidator::is_masked("13812345678", "138****5678"));
682 assert!(!MaskingValidator::is_masked("same", "same"));
683 assert!(MaskingValidator::is_masked("", ""));
684 }
685
686 #[test]
687 fn validator_contains_mask_char() {
688 assert!(MaskingValidator::contains_mask_char("138****5678"));
689 assert!(!MaskingValidator::contains_mask_char("13812345678"));
690 }
691
692 #[test]
693 fn validator_no_sensitive_substring() {
694 assert!(MaskingValidator::no_sensitive_substring(
695 "138****5678",
696 "1234"
697 ));
698 assert!(!MaskingValidator::no_sensitive_substring(
699 "13812345678",
700 "1234"
701 ));
702 }
703
704 #[test]
705 fn validator_no_sensitive_substring_short() {
706 assert!(MaskingValidator::no_sensitive_substring("ab", "ab"));
708 }
709
710 #[test]
711 fn validator_validate_map() {
712 let mut original = HashMap::new();
713 original.insert("phone".to_string(), "13812345678".to_string());
714 original.insert("name".to_string(), "Alice".to_string());
715 let mut masked = HashMap::new();
716 masked.insert("phone".to_string(), "138****5678".to_string());
717 masked.insert("name".to_string(), "Alice".to_string());
718 let unmasked = MaskingValidator::validate_map(&original, &masked);
719 assert!(unmasked.contains(&"name".to_string()));
721 assert!(!unmasked.contains(&"phone".to_string()));
722 }
723}