sabia/crates/sabia-core/src/rules.rs

//! Wordlist mutation rules (lightweight hashcat `-r` style) and mask brute-force
//! generator (hashcat `-a 3` style).
//!
//! Because each WPA-PSK candidate requires 4096 rounds of PBKDF2-HMAC-SHA1,
//! mutation rules focus on high-yield patterns (capitalization, common numeric
//! suffixes, single-character leetspeak) to keep the multiplier bounded.
//!
//!

/// Upper bound on variants produced by [`expand`] per base word.
///
pub const VARIANTS_PER_WORD: usize = 1 + 1 + 1 + SUFFIXES.len() + LEET_PAIRS.len();

const SUFFIXES: &[&str] = &[
    "1", "12", "123", "1234", "12345", "123456", "01", "007", "99", "2023", "2024", "2025", "!",
    "@", "#",
];

/// `(source_char, leet_replacement)` single-substitution pairs.
///
const LEET_PAIRS: &[(char, char)] = &[('a', '@'), ('e', '3'), ('i', '1'), ('o', '0'), ('s', '$')];

/// Generates deduplicated WPA-valid (8–63 char) mutations of `word` (including `word` itself).
///
pub fn expand(word: &str) -> Vec<String> {
    let mut out = Vec::with_capacity(VARIANTS_PER_WORD);
    let mut push = |s: String| {
        let len = s.chars().count();
        if (8..=63).contains(&len) && !out.contains(&s) {
            out.push(s);
        }
    };

    push(word.to_string());

    if let Some(first) = word.chars().next() {
        let capitalized: String = first.to_uppercase().chain(word.chars().skip(1)).collect();
        push(capitalized);
    }
    push(word.to_uppercase());

    for suffix in SUFFIXES {
        push(format!("{word}{suffix}"));
    }

    for (from, to) in LEET_PAIRS {
        if word.contains(*from) {
            push(word.replace(*from, &to.to_string()));
        }
    }

    out
}

/// Character class for a mask token (`?d` digits, `?l` lowercase, `?u` uppercase,
/// `?s` symbols, `?a` all printable ASCII, `??` literal `?`).
///
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CharClass {
    Digit,
    Lower,
    Upper,
    Symbol,
    All,
    Literal(char),
}

const DIGITS: &str = "0123456789";
const LOWER: &str = "abcdefghijklmnopqrstuvwxyz";
const UPPER: &str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
const SYMBOLS: &str = "!@#$%&*_-";

impl CharClass {
    fn alphabet(self) -> String {
        match self {
            CharClass::Digit => DIGITS.to_string(),
            CharClass::Lower => LOWER.to_string(),
            CharClass::Upper => UPPER.to_string(),
            CharClass::Symbol => SYMBOLS.to_string(),
            CharClass::All => format!("{DIGITS}{LOWER}{UPPER}{SYMBOLS}"),
            CharClass::Literal(c) => c.to_string(),
        }
    }
}

/// Compiled brute-force mask (one alphabet per character position), supporting
/// O(1) mixed-radix candidate generation by index.
///
#[derive(Debug, Clone)]
pub struct Mask {
    alphabets: Vec<Vec<char>>,
}

impl Mask {
    /// Compiles a hashcat-style mask pattern such as `"wifi?d?d?d?d"`.
    /// `??` escapes a literal `?`; unknown `?x` tokens return an error.
    pub fn parse(pattern: &str) -> Result<Self, String> {
        let mut alphabets = Vec::new();
        let mut chars = pattern.chars().peekable();
        while let Some(c) = chars.next() {
            let class = if c == '?' {
                match chars.next() {
                    Some('d') => CharClass::Digit,
                    Some('l') => CharClass::Lower,
                    Some('u') => CharClass::Upper,
                    Some('s') => CharClass::Symbol,
                    Some('a') => CharClass::All,
                    Some('?') => CharClass::Literal('?'),
                    Some(other) => return Err(format!("classe de máscara desconhecida: ?{other}")),
                    None => return Err("máscara termina com '?' sem classe".into()),
                }
            } else {
                CharClass::Literal(c)
            };
            alphabets.push(class.alphabet().chars().collect());
        }
        if alphabets.is_empty() {
            return Err("máscara vazia".into());
        }
        if alphabets.len() > 63 {
            return Err("máscara maior que o limite de uma senha WPA (63)".into());
        }
        Ok(Self { alphabets })
    }

    /// Total keyspace size, or `None` if it overflows `u64`.
    ///
    pub fn total(&self) -> Option<u64> {
        self.alphabets
            .iter()
            .try_fold(1u64, |acc, a| acc.checked_mul(a.len() as u64))
    }

    pub fn len(&self) -> usize {
        self.alphabets.len()
    }

    pub fn is_empty(&self) -> bool {
        self.alphabets.is_empty()
    }

    /// Materializes the candidate at `index` (mixed-radix odometer order), or
    /// `None` if `index >= total()`.
    ///
    pub fn candidate(&self, mut index: u64) -> Option<String> {
        let total = self.total()?;
        if index >= total {
            return None;
        }
        let mut chars = vec!['\0'; self.alphabets.len()];
        for i in (0..self.alphabets.len()).rev() {
            let base = self.alphabets[i].len() as u64;
            let digit = (index % base) as usize;
            chars[i] = self.alphabets[i][digit];
            index /= base;
        }
        Some(chars.into_iter().collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn expand_respects_min_length() {
        let variants = expand("senha");
        // `"senha"` is 5 chars (too short for WPA alone), but `"senha123"` (8 chars) is kept.
        //
        assert!(!variants.iter().any(|v| v == "senha"));
        assert!(variants.contains(&"senha123".to_string()));
    }

    #[test]
    fn expand_includes_original_when_long_enough() {
        let variants = expand("abacaxi1");
        assert!(variants.contains(&"abacaxi1".to_string()));
    }

    #[test]
    fn expand_applies_capitalization_suffix_and_leet() {
        let variants = expand("abacaxi1");
        assert!(variants.contains(&"Abacaxi1".to_string()));
        assert!(variants.contains(&"ABACAXI1".to_string()));
        assert!(variants.contains(&"abacaxi1123".to_string()));
        // Leetspeak substitution replaces all occurrences of `'a'` in `"abacaxi1"`.
        assert!(variants.contains(&"@b@c@xi1".to_string()));
    }

    #[test]
    fn expand_has_no_duplicates() {
        let variants = expand("12345678");
        let mut sorted = variants.clone();
        sorted.sort();
        sorted.dedup();
        assert_eq!(variants.len(), sorted.len());
    }

    #[test]
    fn mask_parses_literals_and_classes() {
        let mask = Mask::parse("wifi?d?d?d?d").expect("máscara válida");
        assert_eq!(mask.len(), 8);
        assert_eq!(mask.total(), Some(10_000));
        assert_eq!(mask.candidate(0).as_deref(), Some("wifi0000"));
        assert_eq!(mask.candidate(9_999).as_deref(), Some("wifi9999"));
        assert_eq!(mask.candidate(10_000), None);
    }

    #[test]
    fn mask_literal_question_mark_and_unknown_class() {
        // `"??"` escapes a literal `'?'`, so `"???d"` parses as literal `'?'` + `?d`.
        //
        assert_eq!(
            Mask::parse("???d").unwrap().candidate(0).as_deref(),
            Some("?0")
        );
        assert!(Mask::parse("?z").is_err());
        assert!(Mask::parse("").is_err());
    }

    #[test]
    fn mask_covers_every_combination_exactly_once() {
        let mask = Mask::parse("?d?l").expect("máscara válida");
        let total = mask.total().unwrap();
        assert_eq!(total, 10 * 26);
        let mut seen = std::collections::HashSet::new();
        for i in 0..total {
            seen.insert(mask.candidate(i).unwrap());
        }
        assert_eq!(seen.len() as u64, total);
    }
}