1use std::collections::HashMap;
7use std::sync::Arc;
8use anyhow::{Result, Context, anyhow};
9use strip_ansi_escapes::strip;
10use sha2::{Digest, Sha256};
11use hex;
12use chrono::Utc;
13use tokio::sync::mpsc;
14
15use crate::config::{RedactionConfig, RedactionSummaryItem, RedactionRule};
16use crate::redaction_match::{RedactionMatch, RedactionLog, ensure_match_hashes};
17use crate::profiles::EngineOptions;
18use crate::engine::SanitizationEngine;
19use crate::sanitizers::compiler::{get_or_compile_rules, CompiledRules, CompiledRule};
20use crate::validators;
21
22#[derive(Debug)]
24struct StrippedIndexMapper {
25 map: Vec<usize>,
26}
27
28impl StrippedIndexMapper {
29 fn new(original: &str) -> Self {
30 let stripped_bytes = strip(original.as_bytes());
31 let stripped_str = String::from_utf8_lossy(&stripped_bytes);
32 let mut map: Vec<usize> = Vec::with_capacity(stripped_str.len() + 1);
33 let mut original_char_indices = original.char_indices();
34 let stripped_chars = stripped_str.chars();
35 let mut current_orig_char = original_char_indices.next();
36 for stripped_char in stripped_chars {
37 while let Some((orig_index, orig_char)) = current_orig_char {
38 if orig_char == stripped_char {
39 map.push(orig_index);
40 current_orig_char = original_char_indices.next();
41 break;
42 }
43 current_orig_char = original_char_indices.next();
44 }
45 }
46 map.push(original.len());
47 Self { map }
48 }
49
50 fn map_index(&self, stripped_index: usize) -> usize {
51 let idx = stripped_index.min(self.map.len().saturating_sub(1));
52 self.map[idx]
53 }
54}
55
56pub const BATCH_SIZE: usize = 4096;
57
58#[derive(Debug)]
59pub struct RegexEngine {
60 compiled_rules: Arc<CompiledRules>,
61 config: RedactionConfig,
62 options: EngineOptions,
63 remediation_tx: Option<mpsc::Sender<RedactionMatch>>,
64}
65
66impl RegexEngine {
67 pub fn new(config: RedactionConfig) -> Result<Self> {
68 Self::with_options(config, EngineOptions::default())
69 }
70
71 pub fn with_options(config: RedactionConfig, options: EngineOptions) -> Result<Self> {
72 let compiled_rules = get_or_compile_rules(&config)
73 .context("Failed to compile redaction rules for RegexEngine")?;
74
75 Ok(Self {
76 compiled_rules,
77 config,
78 options,
79 remediation_tx: None,
80 })
81 }
82
83 fn run_programmatic_validator(&self, compiled_rule: &CompiledRule, original_str: &str) -> bool {
84 if !compiled_rule.programmatic_validation {
85 return true;
86 }
87 match compiled_rule.name.as_str() {
88 "us_ssn" => validators::is_valid_ssn_programmatically(original_str),
89 "uk_nino" => validators::is_valid_uk_nino_programmatically(original_str),
90 "visa_card" | "mastercard_card" | "amex_card" | "discover_card" => {
91 validators::is_valid_credit_card_programmatically(original_str)
92 }
93 _ => true,
94 }
95 }
96
97 fn create_redaction_match(
98 &self,
99 rule_config: &RedactionRule,
100 original_match_str: &str,
101 start: u64,
102 end: u64,
103 replacement: String,
104 stripped_input: &str,
105 source_id: &str,
106 line_number: Option<u64>,
107 ) -> RedactionMatch {
108 let mut sample_hash = None;
109 let mut match_context_hash = None;
110 let needs_sample_hash = self.options.post_processing.as_ref().map_or(false, |pp| pp.replace_with_token) ||
111 self.options.samples_config.is_some();
112 let needs_context_hash = self.options.dedupe_config.as_ref().map_or(false, |dedupe| dedupe.use_hash);
113
114 if needs_sample_hash || needs_context_hash {
115 let mut hasher = Sha256::new();
116 if needs_sample_hash {
117 hasher.update(original_match_str.as_bytes());
118 sample_hash = Some(hex::encode(hasher.finalize_reset()));
119 }
120 if needs_context_hash {
121 let window = self.options.dedupe_config.as_ref().map(|d| d.window_bytes).unwrap_or(0);
122 let ctx_start = (start as usize).saturating_sub(window);
123 let ctx_end = std::cmp::min(stripped_input.len(), (end as usize).saturating_add(window));
124 let ctx = &stripped_input[ctx_start..ctx_end];
125 hasher.update(ctx.as_bytes());
126 match_context_hash = Some(hex::encode(hasher.finalize()));
127 }
128 }
129
130 RedactionMatch {
131 rule_name: rule_config.name.clone(),
132 original_string: original_match_str.to_string(),
133 sanitized_string: replacement,
134 start,
135 end,
136 sample_hash,
137 match_context_hash,
138 timestamp: Some(Utc::now().to_rfc3339()),
139 rule: rule_config.clone(),
140 source_id: source_id.to_string(),
141 line_number,
142 }
143 }
144
145 fn find_matches(&self, content: &str, source_id: &str) -> Result<HashMap<String, Vec<RedactionMatch>>> {
146 let stripped_bytes = strip(content.as_bytes());
147 let stripped_input = String::from_utf8_lossy(&stripped_bytes);
148 let original_rules_map: HashMap<&str, &RedactionRule> = self.config.rules.iter()
149 .map(|rule| (rule.name.as_str(), rule)).collect();
150 let mut all_matches: HashMap<String, Vec<RedactionMatch>> = HashMap::new();
151
152 for compiled_rule in &self.compiled_rules.rules {
153 if let Some(rule_config) = original_rules_map.get(compiled_rule.name.as_str()) {
154 if let Some(false) = rule_config.enabled { continue; }
155 for caps in compiled_rule.regex.captures_iter(&stripped_input) {
156 let original_match = caps.get(0).ok_or_else(|| anyhow!("Regex capture failed"))?;
157 if self.run_programmatic_validator(compiled_rule, original_match.as_str()) {
158 let mut replacement = compiled_rule.replace_with.clone();
159 for i in 1..caps.len() {
160 if let Some(group) = caps.get(i) {
161 replacement = replacement.replace(&format!("${}", i), group.as_str());
162 }
163 }
164 let m = self.create_redaction_match(
165 rule_config, original_match.as_str(), original_match.start() as u64,
166 original_match.end() as u64, replacement, &stripped_input, source_id, None,
167 );
168 if let Some(tx) = &self.remediation_tx { let _ = tx.try_send(m.clone()); }
169 all_matches.entry(compiled_rule.name.clone()).or_default().push(m);
170 }
171 }
172 }
173 }
174 Ok(all_matches)
175 }
176}
177
178impl SanitizationEngine for RegexEngine {
179 fn sanitize(
180 &self,
181 content: &str,
182 source_id: &str,
183 run_id: &str,
184 input_hash: &str,
185 user_id: &str,
186 reason: &str,
187 outcome: &str,
188 mut audit_log: Option<&mut crate::audit_log::AuditLog>,
189 ) -> Result<(String, Vec<RedactionSummaryItem>)> {
190 let all_matches = self.find_matches(content, source_id)?;
191 let mut sorted_matches: Vec<&RedactionMatch> = all_matches.values().flatten().collect();
192 sorted_matches.sort_by_key(|m| m.start);
193 let mapper = StrippedIndexMapper::new(content);
194 let mut sanitized_content = String::with_capacity(content.len());
195 let mut last_end = 0usize;
196
197 for m in sorted_matches.iter() {
198 let original_start_byte = mapper.map_index(m.start as usize);
199 let original_end_byte = mapper.map_index(m.end as usize);
200 if original_end_byte <= last_end { continue; }
201 let current_start = original_start_byte.max(last_end);
202 sanitized_content.push_str(&content[last_end..current_start]);
203 sanitized_content.push_str(&m.sanitized_string);
204 last_end = original_end_byte;
205
206 if let Some(log) = audit_log.as_mut() {
207 let _ = log.append(&RedactionLog {
208 timestamp: m.timestamp.clone().unwrap_or_default(),
209 run_id: run_id.to_string(), file_path: source_id.to_string(),
210 user_id: user_id.to_string(), reason_for_redaction: reason.to_string(),
211 redaction_outcome: outcome.to_string(), rule_name: m.rule_name.clone(),
212 input_hash: input_hash.to_string(), match_hash: m.sample_hash.clone().unwrap_or_default(),
213 start: m.start, end: m.end,
214 });
215 }
216 }
217 sanitized_content.push_str(&content[last_end..]);
218 let mut summary = Vec::new();
219 for (rule_name, matches) in all_matches.iter() {
220 summary.push(RedactionSummaryItem {
221 rule_name: rule_name.clone(), occurrences: matches.len(),
222 original_texts: matches.iter().map(|m| m.original_string.clone()).collect(),
223 sanitized_texts: matches.iter().map(|m| m.sanitized_string.clone()).collect(),
224 });
225 }
226 Ok((sanitized_content, summary))
227 }
228
229 fn analyze_for_stats(&self, content: &str, source_id: &str) -> Result<Vec<RedactionSummaryItem>> {
230 let all_matches = self.find_matches(content, source_id)?;
231 let mut summary = Vec::new();
232 for (rule_name, matches) in all_matches.iter() {
233 summary.push(RedactionSummaryItem {
234 rule_name: rule_name.clone(), occurrences: matches.len(),
235 original_texts: matches.iter().map(|m| m.original_string.clone()).collect(),
236 sanitized_texts: matches.iter().map(|m| m.sanitized_string.clone()).collect(),
237 });
238 }
239 Ok(summary)
240 }
241
242 fn find_matches_for_ui(&self, content: &str, source_id: &str) -> Result<Vec<RedactionMatch>> {
243 let all_map = self.find_matches(content, source_id)?;
244 let mut out: Vec<RedactionMatch> = all_map.into_values().flatten().collect();
245 ensure_match_hashes(&mut out);
246 out.sort_by_key(|m| m.start);
247 Ok(out)
248 }
249
250 fn get_heat_scores(&self, content: &str) -> Vec<f64> { vec![0.0; content.len()] }
251 fn compiled_rules(&self) -> &CompiledRules { &self.compiled_rules }
252 fn get_rules(&self) -> &RedactionConfig { &self.config }
253 fn get_options(&self) -> &EngineOptions { &self.options }
254 fn set_remediation_tx(&mut self, tx: mpsc::Sender<RedactionMatch>) { self.remediation_tx = Some(tx); }
255}