base_d/features/
detection.rs

1use crate::core::config::{DictionaryRegistry, EncodingMode};
2use crate::core::dictionary::Dictionary;
3use crate::decode;
4use std::collections::HashSet;
5
6/// A match result from dictionary detection.
7#[derive(Debug, Clone)]
8pub struct DictionaryMatch {
9    /// Name of the matched dictionary
10    pub name: String,
11    /// Confidence score (0.0 to 1.0)
12    pub confidence: f64,
13    /// The dictionary itself
14    pub dictionary: Dictionary,
15}
16
17/// Detector for automatically identifying which dictionary was used to encode data.
18pub struct DictionaryDetector {
19    dictionaries: Vec<(String, Dictionary)>,
20}
21
22impl DictionaryDetector {
23    /// Creates a new detector from a configuration.
24    pub fn new(config: &DictionaryRegistry) -> Result<Self, Box<dyn std::error::Error>> {
25        let mut dictionaries = Vec::new();
26
27        for (name, dict_config) in &config.dictionaries {
28            let dictionary = match dict_config.mode {
29                EncodingMode::ByteRange => {
30                    let start = dict_config
31                        .start_codepoint
32                        .ok_or("ByteRange mode requires start_codepoint")?;
33                    Dictionary::new_with_mode_and_range(
34                        Vec::new(),
35                        dict_config.mode.clone(),
36                        None,
37                        Some(start),
38                    )?
39                }
40                _ => {
41                    let chars: Vec<char> = dict_config.chars.chars().collect();
42                    let padding = dict_config.padding.as_ref().and_then(|s| s.chars().next());
43                    Dictionary::new_with_mode(chars, dict_config.mode.clone(), padding)?
44                }
45            };
46            dictionaries.push((name.clone(), dictionary));
47        }
48
49        Ok(DictionaryDetector { dictionaries })
50    }
51
52    /// Detect which dictionary was likely used to encode the input.
53    /// Returns matches sorted by confidence (highest first).
54    pub fn detect(&self, input: &str) -> Vec<DictionaryMatch> {
55        let input = input.trim();
56        if input.is_empty() {
57            return Vec::new();
58        }
59
60        let mut matches = Vec::new();
61
62        for (name, dict) in &self.dictionaries {
63            if let Some(confidence) = self.score_dictionary(input, dict) {
64                matches.push(DictionaryMatch {
65                    name: name.clone(),
66                    confidence,
67                    dictionary: dict.clone(),
68                });
69            }
70        }
71
72        // Sort by confidence descending
73        matches.sort_by(|a, b| b.confidence.partial_cmp(&a.confidence).unwrap());
74
75        matches
76    }
77
78    /// Score how likely a dictionary matches the input.
79    /// Returns Some(confidence) if it's a plausible match, None otherwise.
80    fn score_dictionary(&self, input: &str, dict: &Dictionary) -> Option<f64> {
81        let mut score = 0.0;
82        let mut weight_sum = 0.0;
83
84        // Weight for each scoring component
85        const CHARSET_WEIGHT: f64 = 0.25;
86        const SPECIFICITY_WEIGHT: f64 = 0.20; // Increased
87        const PADDING_WEIGHT: f64 = 0.30; // Increased (very important for RFC standards)
88        const LENGTH_WEIGHT: f64 = 0.15;
89        const DECODE_WEIGHT: f64 = 0.10;
90
91        // 1. Character set matching
92        let charset_score = self.score_charset(input, dict);
93        score += charset_score * CHARSET_WEIGHT;
94        weight_sum += CHARSET_WEIGHT;
95
96        // If character set score is too low, skip this dictionary
97        if charset_score < 0.5 {
98            return None;
99        }
100
101        // 1.5. Specificity - does this dictionary use a focused character set?
102        let specificity_score = self.score_specificity(input, dict);
103        score += specificity_score * SPECIFICITY_WEIGHT;
104        weight_sum += SPECIFICITY_WEIGHT;
105
106        // 2. Padding detection (for chunked modes)
107        if let Some(padding_score) = self.score_padding(input, dict) {
108            score += padding_score * PADDING_WEIGHT;
109            weight_sum += PADDING_WEIGHT;
110        }
111
112        // 3. Length validation
113        let length_score = self.score_length(input, dict);
114        score += length_score * LENGTH_WEIGHT;
115        weight_sum += LENGTH_WEIGHT;
116
117        // 4. Decode validation (try to actually decode)
118        if let Some(decode_score) = self.score_decode(input, dict) {
119            score += decode_score * DECODE_WEIGHT;
120            weight_sum += DECODE_WEIGHT;
121        }
122
123        // Normalize score
124        if weight_sum > 0.0 {
125            Some(score / weight_sum)
126        } else {
127            None
128        }
129    }
130
131    /// Score based on character set matching.
132    fn score_charset(&self, input: &str, dict: &Dictionary) -> f64 {
133        // Get all unique characters in input (excluding whitespace and padding)
134        let input_chars: HashSet<char> = input
135            .chars()
136            .filter(|c| !c.is_whitespace() && Some(*c) != dict.padding())
137            .collect();
138
139        if input_chars.is_empty() {
140            return 0.0;
141        }
142
143        // For ByteRange mode, check if characters are in the expected range
144        if let Some(start) = dict.start_codepoint() {
145            let in_range = input_chars
146                .iter()
147                .filter(|&&c| {
148                    let code = c as u32;
149                    code >= start && code < start + 256
150                })
151                .count();
152            return in_range as f64 / input_chars.len() as f64;
153        }
154
155        // Check if all input characters are in the dictionary
156        let mut valid_count = 0;
157        for c in &input_chars {
158            if dict.decode_char(*c).is_some() {
159                valid_count += 1;
160            }
161        }
162
163        if valid_count < input_chars.len() {
164            // Not all characters are valid - reject this dictionary
165            return 0.0;
166        }
167
168        // All characters are valid. Now check how well the dictionary size matches
169        let dict_size = dict.base();
170        let input_unique = input_chars.len();
171
172        // Calculate what percentage of the dictionary is actually used
173        let usage_ratio = input_unique as f64 / dict_size as f64;
174
175        // Prefer dictionaries where we use most of the character set
176        // This helps distinguish base64 (64 chars) from base85 (85 chars)
177        if usage_ratio > 0.7 {
178            // We're using >70% of dictionary - excellent match
179            1.0
180        } else if usage_ratio > 0.5 {
181            // We're using >50% of dictionary - good match
182            0.85
183        } else if usage_ratio > 0.3 {
184            // We're using >30% of dictionary - okay match
185            0.7
186        } else {
187            // We're using <30% of dictionary - probably wrong
188            // (e.g., using 20 chars of a 85-char dictionary)
189            0.5
190        }
191    }
192
193    /// Score based on how specific/focused the dictionary character set is.
194    /// Smaller, more focused dictionaries score higher.
195    fn score_specificity(&self, _input: &str, dict: &Dictionary) -> f64 {
196        let dict_size = dict.base();
197
198        // Prefer smaller, more common dictionaries
199        // This helps distinguish base64 (64) from base85 (85) when both match
200        match dict_size {
201            16 => 1.0,   // hex
202            32 => 0.95,  // base32
203            58 => 0.90,  // base58
204            62 => 0.88,  // base62
205            64 => 0.92,  // base64 (very common)
206            85 => 0.70,  // base85 (less common)
207            256 => 0.60, // base256
208            _ if dict_size < 64 => 0.85,
209            _ if dict_size < 128 => 0.75,
210            _ => 0.65,
211        }
212    }
213
214    /// Score based on padding character presence and position.
215    fn score_padding(&self, input: &str, dict: &Dictionary) -> Option<f64> {
216        let padding = dict.padding()?;
217
218        // Chunked modes should have padding at the end (or no padding)
219        if *dict.mode() == EncodingMode::Chunked {
220            let has_padding = input.ends_with(padding);
221            let padding_count = input.chars().filter(|c| *c == padding).count();
222
223            if has_padding {
224                // Padding should only be at the end
225                let trimmed = input.trim_end_matches(padding);
226                let internal_padding = trimmed.chars().any(|c| c == padding);
227
228                if internal_padding {
229                    Some(0.5) // Suspicious padding in middle
230                } else if padding_count <= 3 {
231                    Some(1.0) // Valid padding
232                } else {
233                    Some(0.3) // Too much padding
234                }
235            } else {
236                // No padding is also valid for chunked mode
237                Some(0.8)
238            }
239        } else {
240            None
241        }
242    }
243
244    /// Score based on input length validation for the encoding mode.
245    fn score_length(&self, input: &str, dict: &Dictionary) -> f64 {
246        let length = input.trim().len();
247
248        match dict.mode() {
249            EncodingMode::Chunked => {
250                // Chunked mode should have specific alignment
251                let base = dict.base();
252
253                // Remove padding to check alignment
254                let trimmed = if let Some(pad) = dict.padding() {
255                    input.trim_end_matches(pad)
256                } else {
257                    input
258                };
259
260                // For base64 (6 bits per char), output should be multiple of 4
261                // For base32 (5 bits per char), output should be multiple of 8
262                // For base16 (4 bits per char), output should be multiple of 2
263                let expected_multiple = match base {
264                    64 => 4,
265                    32 => 8,
266                    16 => 2,
267                    _ => return 0.5, // Unknown chunked base
268                };
269
270                if trimmed.len() % expected_multiple == 0 {
271                    1.0
272                } else {
273                    0.3
274                }
275            }
276            EncodingMode::ByteRange => {
277                // ByteRange is 1:1 mapping, any length is valid
278                1.0
279            }
280            EncodingMode::BaseConversion => {
281                // Mathematical conversion can produce any length
282                if length > 0 {
283                    1.0
284                } else {
285                    0.0
286                }
287            }
288        }
289    }
290
291    /// Score based on whether the input can be successfully decoded.
292    fn score_decode(&self, input: &str, dict: &Dictionary) -> Option<f64> {
293        match decode(input, dict) {
294            Ok(decoded) => {
295                if decoded.is_empty() {
296                    Some(0.5)
297                } else {
298                    // Successfully decoded!
299                    Some(1.0)
300                }
301            }
302            Err(_) => {
303                // Failed to decode
304                Some(0.0)
305            }
306        }
307    }
308}
309
310/// Convenience function to detect dictionary from input.
311pub fn detect_dictionary(input: &str) -> Result<Vec<DictionaryMatch>, Box<dyn std::error::Error>> {
312    let config = DictionaryRegistry::load_with_overrides()?;
313    let detector = DictionaryDetector::new(&config)?;
314    Ok(detector.detect(input))
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320    use crate::encode;
321
322    #[test]
323    fn test_detect_base64() {
324        let config = DictionaryRegistry::load_default().unwrap();
325        let detector = DictionaryDetector::new(&config).unwrap();
326
327        // Standard base64 with padding
328        let matches = detector.detect("SGVsbG8sIFdvcmxkIQ==");
329        assert!(!matches.is_empty());
330        // base64 and base64url are very similar, so either is acceptable
331        assert!(matches[0].name == "base64" || matches[0].name == "base64url");
332        assert!(matches[0].confidence > 0.7);
333    }
334
335    #[test]
336    fn test_detect_base32() {
337        let config = DictionaryRegistry::load_default().unwrap();
338        let detector = DictionaryDetector::new(&config).unwrap();
339
340        let matches = detector.detect("JBSWY3DPEBLW64TMMQ======");
341        assert!(!matches.is_empty());
342        // base32 should be in top 5 candidates
343        let base32_found = matches.iter().take(5).any(|m| m.name.starts_with("base32"));
344        assert!(base32_found, "base32 should be in top 5 candidates");
345    }
346
347    #[test]
348    fn test_detect_hex() {
349        let config = DictionaryRegistry::load_default().unwrap();
350        let detector = DictionaryDetector::new(&config).unwrap();
351
352        let matches = detector.detect("48656c6c6f");
353        assert!(!matches.is_empty());
354        // hex or hex_math are both correct
355        assert!(matches[0].name == "hex" || matches[0].name == "hex_math");
356        assert!(matches[0].confidence > 0.8);
357    }
358
359    #[test]
360    fn test_detect_from_encoded() {
361        let config = DictionaryRegistry::load_default().unwrap();
362
363        // Test with actual encoding
364        let dict_config = config.get_dictionary("base64").unwrap();
365        let chars: Vec<char> = dict_config.chars.chars().collect();
366        let padding = dict_config.padding.as_ref().and_then(|s| s.chars().next());
367        let dict = Dictionary::new_with_mode(chars, dict_config.mode.clone(), padding).unwrap();
368
369        let data = b"Hello, World!";
370        let encoded = encode(data, &dict);
371
372        let detector = DictionaryDetector::new(&config).unwrap();
373        let matches = detector.detect(&encoded);
374
375        assert!(!matches.is_empty());
376        // base64 and base64url only differ by 2 chars, so both are valid
377        assert!(matches[0].name == "base64" || matches[0].name == "base64url");
378    }
379
380    #[test]
381    fn test_detect_empty_input() {
382        let config = DictionaryRegistry::load_default().unwrap();
383        let detector = DictionaryDetector::new(&config).unwrap();
384
385        let matches = detector.detect("");
386        assert!(matches.is_empty());
387    }
388
389    #[test]
390    fn test_detect_invalid_input() {
391        let config = DictionaryRegistry::load_default().unwrap();
392        let detector = DictionaryDetector::new(&config).unwrap();
393
394        // Input with characters not in any dictionary
395        let matches = detector.detect("こんにちは世界");
396        // Should return few or no high-confidence matches
397        if !matches.is_empty() {
398            assert!(matches[0].confidence < 0.5);
399        }
400    }
401}
base_d/features/detection.rs

base_d/features/
detection.rs
