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

//! Regras de mutação sobre uma palavra do dicionário (estilo hashcat `-r`,
//! bem mais simples) e força bruta por máscara (estilo hashcat `-a 3`).
//!
//! Em quebra de hash normal, regras são "de graça" — um hash SHA256 custa
//! nanossegundos. Aqui não: cada candidato é um PBKDF2-HMAC-SHA1 de 4096
//! iterações, então cada variante multiplica o tempo de verdade. As regras
//! abaixo cobrem os padrões mais comuns (maiúscula inicial, sufixo numérico,
//! leet simples) sem explodir pra centenas de variantes por palavra.

/// Quantas variantes [`expand`] gera por palavra base (pra calcular o
/// multiplicador de tempo antes de rodar).
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", "!",
    "@", "#",
];

/// (char original, substituto) — leet simples, uma troca por vez, não
/// combinada, senão a combinatória explode rápido demais.
const LEET_PAIRS: &[(char, char)] = &[('a', '@'), ('e', '3'), ('i', '1'), ('o', '0'), ('s', '$')];

/// Gera as variantes de `word` (incluindo a palavra original) que respeitam
/// o comprimento válido de WPA (8–63), sem duplicatas.
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
}

/// Classe de caractere de um token `?x` de máscara (sintaxe hashcat: `?d`
/// dígito, `?l` minúscula, `?u` maiúscula, `?s` símbolo comum, `?a` os
/// quatro juntos). `??` escapa um `?` 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(),
        }
    }
}

/// Máscara compilada: uma lista de alfabetos, um por posição. Gera
/// candidatos por índice (contagem em base mista) sem materializar a lista
/// inteira — essencial, já que `?d?d?d?d?d?d?d?d` sozinho são 100 milhões.
#[derive(Debug, Clone)]
pub struct Mask {
    alphabets: Vec<Vec<char>>,
}

impl Mask {
    /// Compila uma máscara tipo `wifi?d?d?d?d` (literal "wifi" + 4 dígitos).
    /// `??` é um `?` literal; `?` seguido de letra desconhecida é erro.
    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 de candidatos. `None` se não couber num `u64` (máscara longa
    /// demais pra ser levada a sério de qualquer forma).
    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()
    }

    /// Candidato de índice `index` (contagem em base mista, como um
    /// odômetro: a última posição gira mais rápido). `None` se `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" tem 5 chars; sozinha não é candidata WPA válida (precisa
        // de 8+), mas com sufixo (ex.: "senha123", 8) sim.
        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()));
        // leet troca TODAS as ocorrências de 'a': as três em "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() {
        // "??" vira um '?' literal; os dois primeiros chars são consumidos
        // juntos, então "???d" é: '?' literal + "?d" (classe dígito).
        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);
    }
}