//! Decriptação CCMP (WPA2-AES) de quadros de dados, usando a TK já derivada
//! de um handshake quebrado.
//!
//! Isso não serve pra forjar nem reinjetar nada — só pra enxergar o conteúdo
//! em claro (DHCP, mDNS) de um cliente cuja senha da rede já foi recuperada
//! por dicionário, e assim extrair hostname/fabricante de verdade em vez de
//! só o MAC. A chave é par (BSSID, estação): só decifra tráfego do cliente
//! cujo próprio handshake foi capturado e quebrado, nunca de terceiros —
//! cada associação tem sua PTK efêmera, então a senha da rede sozinha não
//! abre o tráfego de quem a gente não viu entrar.
//!
//! Implementa só o modo CTR (a parte que decifra o conteúdo). O MIC
//! (CBC-MAC) não é conferido: a intenção aqui é mostrar o payload de um
//! tráfego que a gente já pode decifrar porque a senha é conhecida, não
//! provar integridade de um quadro adulterado — e implementar o AAD/CBC-MAC
//! certinho exigiria replicar o mascaramento do cabeçalho 802.11 inteiro
//! sem ganho prático pro que isso serve aqui.

use aes::cipher::{BlockEncrypt, KeyInit};
use aes::Aes128;

/// Decifra o corpo de um quadro de dados CCMP, a partir do cabeçalho CCMP
/// (8 bytes) em diante — ou seja, `body` é `frame.body[frame.header_len..]`.
/// `transmitter` é o endereço 2 (A2) do quadro — o rádio que o enviou, não o
/// BSSID/estação lógicos. `qos_tid` é o TID do QoS Control (0 se não-QoS).
///
/// Retorna o LLC/SNAP + payload em claro (o mesmo formato que um quadro não
/// protegido teria), sem os 8 bytes de MIC cifrado no final.
pub fn decrypt_data(
    tk: &[u8; 16],
    transmitter: &[u8; 6],
    qos_tid: u8,
    body: &[u8],
) -> Option<Vec<u8>> {
    const HEADER_LEN: usize = 8;
    const MIC_LEN: usize = 8;
    if body.len() < HEADER_LEN + MIC_LEN {
        return None;
    }
    let pn0 = body[0];
    let pn1 = body[1];
    let key_id_octet = body[3];
    // ExtIV (bit 5) precisa estar ligado — senão não é um cabeçalho CCMP de
    // verdade (podia ser WEP de 4 bytes por coincidência de layout).
    if key_id_octet & 0x20 == 0 {
        return None;
    }
    let pn2 = body[4];
    let pn3 = body[5];
    let pn4 = body[6];
    let pn5 = body[7];

    let ciphertext_and_mic = &body[HEADER_LEN..];
    let ciphertext = &ciphertext_and_mic[..ciphertext_and_mic.len() - MIC_LEN];

    // Nonce CCMP (13 bytes): flags/prioridade (1) || A2 (6) || PN, do byte
    // mais significativo (PN5) ao menos significativo (PN0).
    let mut nonce = [0u8; 13];
    nonce[0] = qos_tid & 0x0f;
    nonce[1..7].copy_from_slice(transmitter);
    nonce[7] = pn5;
    nonce[8] = pn4;
    nonce[9] = pn3;
    nonce[10] = pn2;
    nonce[11] = pn1;
    nonce[12] = pn0;

    let cipher = Aes128::new(tk.into());
    let mut out = Vec::with_capacity(ciphertext.len());
    let mut counter: u16 = 1; // o bloco de contador 0 cifra só o MIC, não dado.
    for chunk in ciphertext.chunks(16) {
        let mut block = [0u8; 16];
        block[0] = 0x01; // flag octet do CCMP counter mode (L=1, sem AAD aqui)
        block[1..14].copy_from_slice(&nonce);
        block[14..16].copy_from_slice(&counter.to_be_bytes());
        let mut keystream = block.into();
        cipher.encrypt_block(&mut keystream);
        for (c, k) in chunk.iter().zip(keystream.iter()) {
            out.push(c ^ k);
        }
        counter = counter.wrapping_add(1);
    }
    Some(out)
}

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

    /// AES-128 do FIPS-197 (vetor de referência oficial do NIST): confirma
    /// que a cifra de bloco usada embaixo está correta antes de confiar no
    /// contador CTR por cima dela.
    #[test]
    fn aes128_matches_fips197_vector() {
        let key: [u8; 16] = [
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
            0x0e, 0x0f,
        ];
        let plaintext: [u8; 16] = [
            0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
            0xee, 0xff,
        ];
        let expected: [u8; 16] = [
            0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30, 0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4,
            0xc5, 0x5a,
        ];
        let cipher = Aes128::new(&key.into());
        let mut block = plaintext.into();
        cipher.encrypt_block(&mut block);
        assert_eq!(block.as_slice(), &expected);
    }

    /// Round-trip de auto-consistência: decifrar um texto que a própria
    /// função "cifrou" (mesmo algoritmo, usado nos dois sentidos — CTR é
    /// simétrico) tem que devolver o original. Isso pega bug de indexação
    /// no nonce/contador mesmo sem um vetor de teste oficial do CCMP à mão.
    #[test]
    fn ctr_round_trip_recovers_plaintext() {
        let tk = [0x42u8; 16];
        let transmitter = [0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff];
        let plaintext = b"AA AA 03 00 00 00 08 00 payload de teste bem maior que um bloco de 16 bytes para cobrir mais de uma rodada do contador";

        // Monta um "quadro" CCMP cifrando manualmente com a mesma lógica de
        // keystream (ela é simétrica: XOR com o keystream cifra e decifra).
        let header = [0x11u8, 0x22, 0x00, 0x20, 0x33, 0x44, 0x55, 0x66]; // ExtIV ligado
        let mut nonce = [0u8; 13];
        nonce[0] = 3; // TID
        nonce[1..7].copy_from_slice(&transmitter);
        nonce[7] = header[7];
        nonce[8] = header[6];
        nonce[9] = header[5];
        nonce[10] = header[4];
        nonce[11] = header[1];
        nonce[12] = header[0];
        let cipher = Aes128::new(&tk.into());
        let mut ciphertext = Vec::with_capacity(plaintext.len());
        let mut counter: u16 = 1;
        for chunk in plaintext.chunks(16) {
            let mut block = [0u8; 16];
            block[0] = 0x01;
            block[1..14].copy_from_slice(&nonce);
            block[14..16].copy_from_slice(&counter.to_be_bytes());
            let mut keystream = block.into();
            cipher.encrypt_block(&mut keystream);
            for (p, k) in chunk.iter().zip(keystream.iter()) {
                ciphertext.push(p ^ k);
            }
            counter = counter.wrapping_add(1);
        }
        ciphertext.extend_from_slice(&[0u8; 8]); // MIC fictício, ignorado na decriptação

        let mut body = Vec::new();
        body.extend_from_slice(&header);
        body.extend_from_slice(&ciphertext);

        let out = decrypt_data(&tk, &transmitter, 3, &body).expect("decifra");
        assert_eq!(out, plaintext);
    }

    #[test]
    fn rejects_body_without_extiv_bit() {
        let tk = [0u8; 16];
        let body = [0u8; 20]; // key_id_octet (índice 3) sem o bit ExtIV
        assert!(decrypt_data(&tk, &[0; 6], 0, &body).is_none());
    }

    #[test]
    fn rejects_body_too_short() {
        let tk = [0u8; 16];
        let body = [0u8; 10];
        assert!(decrypt_data(&tk, &[0; 6], 0, &body).is_none());
    }
}