cleansh_core/engines/
entropy_engine.rs

1// cleansh-core/src/engines/entropy_engine.rs
2//! A `SanitizationEngine` implementation utilizing Shannon entropy.
3//! Features v0.2.0 Self-Healing Integration and Contiguous Match Merging.
4//!
5//! FIXED: Implemented 'Look-Ahead Stitcher' to prevent window fractures on long secrets.
6
7use std::collections::{HashMap, HashSet};
8use std::sync::Arc;
9use anyhow::Result;
10use strip_ansi_escapes::strip;
11use sha2::{Digest, Sha256};
12use hex;
13use chrono::Utc;
14use tokio::sync::mpsc;
15
16use crate::config::{RedactionConfig, RedactionSummaryItem, RedactionRule};
17use crate::redaction_match::{RedactionMatch, ensure_match_hashes};
18use crate::profiles::EngineOptions;
19use crate::engine::SanitizationEngine;
20use crate::sanitizers::compiler::{get_or_compile_rules, CompiledRules};
21use crate::remediation::fingerprint::SecretFingerprint;
22use cleansh_entropy::engine::EntropyEngine as LowLevelEntropyEngine;
23
24/// Improved Mapper: Handles ANSI escape offsets to prevent partial redaction.
25#[derive(Debug)]
26struct StrippedIndexMapper {
27    map: Vec<usize>,
28}
29
30impl StrippedIndexMapper {
31    fn new(original: &str) -> Self {
32        let stripped_bytes = strip(original.as_bytes());
33        let stripped_str = String::from_utf8_lossy(&stripped_bytes);
34        
35        let mut map: Vec<usize> = Vec::with_capacity(stripped_str.len() + 1);
36        let mut orig_char_indices = original.char_indices().peekable();
37        
38        for stripped_char in stripped_str.chars() {
39            while let Some(&(orig_index, orig_char)) = orig_char_indices.peek() {
40                if orig_char == stripped_char {
41                    map.push(orig_index);
42                    let _ = orig_char_indices.next();
43                    break;
44                } else {
45                    let _ = orig_char_indices.next();
46                }
47            }
48        }
49        map.push(original.len());
50        Self { map }
51    }
52
53    fn map_index(&self, stripped_index: usize) -> usize {
54        let idx = stripped_index.min(self.map.len().saturating_sub(1));
55        self.map[idx]
56    }
57}
58
59#[derive(Debug)]
60pub struct EntropyEngine {
61    config: RedactionConfig,
62    options: EngineOptions,
63    inner_engine: LowLevelEntropyEngine,
64    compiled_rules: Arc<CompiledRules>,
65    remediation_tx: Option<mpsc::Sender<RedactionMatch>>,
66    fingerprint_cache: HashSet<String>,
67}
68
69impl EntropyEngine {
70    pub fn new(config: RedactionConfig) -> Result<Self> {
71        Self::with_options(config, EngineOptions::default())
72    }
73
74    pub fn with_options(config: RedactionConfig, options: EngineOptions) -> Result<Self> {
75        let threshold = config.engines.entropy.threshold.unwrap_or(0.5);
76        let window_size = config.engines.entropy.window_size.unwrap_or(24);
77        let inner_engine = LowLevelEntropyEngine::new(threshold, window_size);
78        let compiled_rules = get_or_compile_rules(&config)?;
79        Ok(Self { 
80            config, 
81            options, 
82            inner_engine, 
83            compiled_rules, 
84            remediation_tx: None, 
85            fingerprint_cache: HashSet::new() 
86        })
87    }
88
89    pub fn update_fingerprints(&mut self, fingerprints: Vec<SecretFingerprint>) {
90        for fp in fingerprints { 
91            self.fingerprint_cache.insert(fp.hash); 
92        }
93    }
94
95    fn create_redaction_match(&self, original: &str, start: u64, end: u64, source_id: &str) -> RedactionMatch {
96        let mut sample_hash = None;
97        if self.options.post_processing.as_ref().map_or(false, |pp| pp.replace_with_token) {
98            let mut hasher = Sha256::new();
99            hasher.update(original.as_bytes());
100            sample_hash = Some(hex::encode(hasher.finalize()));
101        }
102        let rule = RedactionRule {
103            name: "high_entropy_secret".to_string(),
104            replace_with: "[ENTROPY_REDACTED]".to_string(),
105            pattern_type: "entropy".to_string(),
106            ..Default::default()
107        };
108        RedactionMatch {
109            rule_name: rule.name.clone(), 
110            original_string: original.to_string(),
111            sanitized_string: rule.replace_with.clone(), 
112            start, 
113            end, 
114            sample_hash,
115            match_context_hash: None, 
116            timestamp: Some(Utc::now().to_rfc3339()), 
117            rule,
118            source_id: source_id.to_string(), 
119            line_number: None,
120        }
121    }
122
123    fn find_matches_internal(&self, content: &str, source_id: &str) -> Vec<RedactionMatch> {
124        let stripped_bytes = strip(content.as_bytes());
125        let stripped_input = String::from_utf8_lossy(&stripped_bytes);
126        
127        let entropy_matches = self.inner_engine.scan(stripped_input.as_bytes());
128        if entropy_matches.is_empty() { return vec![]; }
129
130        // --- MERGE LOGIC START ---
131        let mut sorted_intervals = entropy_matches;
132        sorted_intervals.sort_by(|a, b| a.start.cmp(&b.start));
133
134        let mut merged_intervals = Vec::new();
135        let mut current_start = sorted_intervals[0].start;
136        let mut current_end = sorted_intervals[0].end;
137
138        for m in sorted_intervals.into_iter().skip(1) {
139            if m.start <= current_end {
140                current_end = std::cmp::max(current_end, m.end);
141            } else {
142                merged_intervals.push((current_start, current_end));
143                current_start = m.start;
144                current_end = m.end;
145            }
146        }
147        merged_intervals.push((current_start, current_end));
148        // --- MERGE LOGIC END ---
149
150        merged_intervals.into_iter().map(|(start, end)| {
151            // Apply refined surgical extraction AND Look-Ahead Stitcher
152            let (refined_start, refined_end) = self.extract_secret_core_indices(&stripped_input, start, end);
153            
154            let m = self.create_redaction_match(
155                &stripped_input[refined_start..refined_end], 
156                refined_start as u64, 
157                refined_end as u64, 
158                source_id
159            );
160            if let Some(tx) = &self.remediation_tx { 
161                let _ = tx.try_send(m.clone()); 
162            }
163            m
164        }).collect()
165    }
166
167    /// Heat-Seeker: Refines the match by anchoring to delimiters.
168    /// NEW: Look-Ahead Stitcher to extend redaction beyond the initial window if characters remain.
169    fn extract_secret_core_indices(&self, text: &str, raw_start: usize, raw_end: usize) -> (usize, usize) {
170        let bytes = text.as_bytes();
171        let len = bytes.len();
172        let mut start = raw_start;
173        let mut end = raw_end;
174
175        // 1. Semantic Anchor: Snap to label delimiters (e.g. key=)
176        let search_range = &bytes[start..end];
177        if let Some(pos) = search_range.iter().rposition(|&b| b == b':' || b == b'=') {
178            let potential_start = start + pos + 1;
179            if potential_start < end {
180                start = potential_start;
181            }
182        }
183
184        // 2. Character-Class Trimming (Leading)
185        while start < end && (
186            bytes[start].is_ascii_whitespace() || 
187            matches!(bytes[start], b'"' | b'\'' | b'[' | b'{' | b'<' | b'(' | b'-' | b'_')
188        ) {
189            start += 1;
190        }
191
192        // 3. LOOK-AHEAD STITCHER (The Fix for Partial Redaction)
193        // If we reached the end of the window, peek forward. If the next characters are
194        // alphanumeric/base64-safe, assume the window fractured a long secret and keep eating.
195        while end < len {
196            let b = bytes[end];
197            // Stop ONLY if we hit a hard delimiter or whitespace.
198            // If it's alphanumeric, underscore, or common secret chars (+, /, -), we extend.
199            if b.is_ascii_whitespace() || matches!(b, b'"' | b'\'' | b',' | b';' | b']' | b'}' | b')' | b'>') {
200                break;
201            }
202            end += 1;
203        }
204
205        // 4. Tail Trimming (Safety Check)
206        // Backtrack from the new end if we accidentally ate a trailing quote or punctuation.
207        while (end - start) > 2 {
208            let tail_byte = bytes[end - 1];
209            if tail_byte.is_ascii_whitespace() 
210                || matches!(tail_byte, b'.' | b',' | b'!' | b'?' | b']' | b'}' | b'>' | b')' | b'"' | b'\'' | b';') 
211            {
212                end -= 1;
213            } else {
214                break;
215            }
216        }
217
218        (start, end)
219    }
220}
221
222impl SanitizationEngine for EntropyEngine {
223    fn sanitize(
224        &self, 
225        content: &str, 
226        source_id: &str, 
227        _run_id: &str, 
228        _input_hash: &str, 
229        _user_id: &str, 
230        _reason: &str, 
231        _outcome: &str, 
232        _audit_log: Option<&mut crate::audit_log::AuditLog>
233    ) -> Result<(String, Vec<RedactionSummaryItem>)> {
234        let matches = self.find_matches_internal(content, source_id);
235        let mapper = StrippedIndexMapper::new(content);
236        let mut sanitized = String::with_capacity(content.len());
237        let mut last_end = 0usize;
238        let mut summary_map: HashMap<String, RedactionSummaryItem> = HashMap::new();
239        let mut sorted = matches;
240        
241        sorted.sort_by_key(|m| m.start);
242
243        for m in &sorted {
244            let original_start = mapper.map_index(m.start as usize);
245            let original_end = mapper.map_index(m.end as usize);
246            
247            if original_end <= last_end { continue; }
248            
249            sanitized.push_str(&content[last_end..original_start.max(last_end)]);
250            sanitized.push_str(&m.sanitized_string);
251            last_end = original_end;
252
253            let entry = summary_map.entry(m.rule_name.clone()).or_insert_with(|| RedactionSummaryItem {
254                rule_name: m.rule_name.clone(), 
255                occurrences: 0, 
256                original_texts: Vec::new(), 
257                sanitized_texts: Vec::new(),
258            });
259            entry.occurrences += 1;
260            entry.original_texts.push(m.original_string.clone());
261            entry.sanitized_texts.push(m.sanitized_string.clone());
262        }
263        
264        sanitized.push_str(&content[last_end..]);
265        Ok((sanitized, summary_map.into_values().collect()))
266    }
267
268    fn analyze_for_stats(&self, content: &str, source_id: &str) -> Result<Vec<RedactionSummaryItem>> {
269        let matches = self.find_matches_internal(content, source_id);
270        let mut summary_map: HashMap<String, RedactionSummaryItem> = HashMap::new();
271        for m in matches {
272            let entry = summary_map.entry(m.rule_name.clone()).or_insert_with(|| RedactionSummaryItem {
273                rule_name: m.rule_name.clone(), 
274                occurrences: 0, 
275                original_texts: Vec::new(), 
276                sanitized_texts: Vec::new(),
277            });
278            entry.occurrences += 1;
279        }
280        Ok(summary_map.into_values().collect())
281    }
282
283    fn find_matches_for_ui(&self, content: &str, source_id: &str) -> Result<Vec<RedactionMatch>> {
284        let mut matches = self.find_matches_internal(content, source_id);
285        ensure_match_hashes(&mut matches);
286        matches.sort_by_key(|m| m.start);
287        Ok(matches)
288    }
289
290    fn get_heat_scores(&self, content: &str) -> Vec<f64> {
291        let stripped_bytes = strip(content.as_bytes());
292        let mut scores = Vec::with_capacity(content.len());
293        
294        for i in 0..stripped_bytes.len() {
295            let start = i.saturating_sub(4);
296            let end = std::cmp::min(stripped_bytes.len(), i + 5);
297            scores.push(cleansh_entropy::entropy::calculate_shannon_entropy(&stripped_bytes[start..end]));
298        }
299        
300        while scores.len() < content.len() { scores.push(0.0); }
301        scores
302    }
303
304    fn compiled_rules(&self) -> &CompiledRules { &self.compiled_rules }
305    fn get_rules(&self) -> &RedactionConfig { &self.config }
306    fn get_options(&self) -> &EngineOptions { &self.options }
307    fn set_remediation_tx(&mut self, tx: mpsc::Sender<RedactionMatch>) { self.remediation_tx = Some(tx); }
308}
cleansh_core/engines/entropy_engine.rs

cleansh_core/engines/
entropy_engine.rs
