//! Parser do handshake WPA/WPA2 (EAPOL-Key, IEEE 802.1X).
//!
//! Lê as quatro mensagens do 4-way handshake, identifica o número de cada
//! uma, extrai nonces, MIC, versão de chave e PMKID — matéria-prima para
//! montar um alvo de quebra por dicionário em [`crate::crack`].

pub const EAPOL_TYPE_KEY: u8 = 3;

const KEY_INFO_PAIRWISE: u16 = 0x0008;
const KEY_INFO_INSTALL: u16 = 0x0040;
const KEY_INFO_ACK: u16 = 0x0080;
const KEY_INFO_MIC: u16 = 0x0100;
const KEY_INFO_SECURE: u16 = 0x0200;

/// Deslocamento absoluto do campo MIC dentro do quadro EAPOL (já incluindo
/// o cabeçalho de 4 bytes).
pub const MIC_OFFSET: usize = 81;
pub const KEY_DATA_OFFSET: usize = 99;
pub const MIN_KEY_FRAME_LEN: usize = 99;

#[derive(Debug, Clone)]
pub struct EapolKey {
    pub version: u8,
    pub descriptor_type: u8,
    /// 1, 2, 3 ou 4; 0 quando não é uma das mensagens do 4-way handshake.
    pub message: u8,
    pub key_info: u16,
    /// Versão do descritor: 1 = TKIP, 2 = CCMP, 3 = AES-CMAC.
    pub key_version: u8,
    pub pairwise: bool,
    pub replay_counter: u64,
    pub nonce: [u8; 32],
    pub mic: [u8; 16],
    pub key_data: Vec<u8>,
    pub pmkid: Option<[u8; 16]>,
    /// Quadro EAPOL completo (cabeçalho + descritor + dados), com o MIC no
    /// lugar — usado para recalcular o MIC durante a quebra.
    pub frame: Vec<u8>,
}

pub fn parse(buf: &[u8]) -> Option<EapolKey> {
    if buf.len() < 4 || buf[1] != EAPOL_TYPE_KEY {
        return None;
    }
    let version = buf[0];
    let declared = u16::from_be_bytes([buf[2], buf[3]]) as usize;
    let total = 4 + declared;
    if total < MIN_KEY_FRAME_LEN || buf.len() < total {
        return None;
    }
    let frame = buf[..total].to_vec();

    let descriptor_type = frame[4];
    // 2 = RSN (WPA2/WPA3), 254 = WPA1. Outros descritores não interessam.
    if descriptor_type != 2 && descriptor_type != 254 {
        return None;
    }
    // Todo campo multi-byte do descritor EAPOL-Key (Key Information, Key
    // Length, Replay Counter, Key Data Length) é big-endian — é 802.1X,
    // segue a mesma ordem de rede do cabeçalho EAPOL em si. Ler como
    // little-endian (bug anterior) embaralhava os bits de ack/mic/secure e
    // fazia praticamente todo handshake real cair em message=0, descartado
    // silenciosamente — só não aparecia nos testes porque o gerador de
    // quadro sintético cometia o mesmo engano, simetricamente.
    let key_info = u16::from_be_bytes([frame[5], frame[6]]);
    let key_version = (key_info & 0x0007) as u8;
    let pairwise = key_info & KEY_INFO_PAIRWISE != 0;
    let ack = key_info & KEY_INFO_ACK != 0;
    let mic = key_info & KEY_INFO_MIC != 0;
    let install = key_info & KEY_INFO_INSTALL != 0;
    let secure = key_info & KEY_INFO_SECURE != 0;

    let message = if ack && !mic {
        1
    } else if !ack && mic && !secure {
        2
    } else if ack && mic && (install || !secure) {
        3
    } else if !ack && mic && secure {
        4
    } else {
        0
    };

    let replay_counter = u64::from_be_bytes(frame[9..17].try_into().ok()?);
    let mut nonce = [0u8; 32];
    nonce.copy_from_slice(&frame[17..49]);
    let mut mic_bytes = [0u8; 16];
    mic_bytes.copy_from_slice(&frame[MIC_OFFSET..MIC_OFFSET + 16]);
    let key_data_len = u16::from_be_bytes([frame[97], frame[98]]) as usize;
    let data_end = KEY_DATA_OFFSET.checked_add(key_data_len)?;
    if data_end > frame.len() {
        return None;
    }
    let key_data = frame[KEY_DATA_OFFSET..data_end].to_vec();
    let pmkid = parse_pmkid_kde(&key_data);

    Some(EapolKey {
        version,
        descriptor_type,
        message,
        key_info,
        key_version,
        pairwise,
        replay_counter,
        nonce,
        mic: mic_bytes,
        key_data,
        pmkid,
        frame,
    })
}

/// PMKID KDE: `dd 14 00 0f ac 04` + 16 bytes. Pode vir em qualquer posição
/// dos dados da chave (normalmente na mensagem 1).
pub fn parse_pmkid_kde(data: &[u8]) -> Option<[u8; 16]> {
    let mut i = 0usize;
    while i + 2 <= data.len() {
        let tag = data[i];
        let len = data[i + 1] as usize;
        let start = i + 2;
        if start + len > data.len() {
            break;
        }
        if tag == 0xdd && len >= 20 && data[start..start + 4] == [0x00, 0x0f, 0xac, 0x04] {
            let mut pmkid = [0u8; 16];
            pmkid.copy_from_slice(&data[start + 4..start + 20]);
            return Some(pmkid);
        }
        i = start + len;
    }
    None
}

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

    pub(crate) fn key_frame(message: u8, keyver: u8, nonce: [u8; 32], key_data: &[u8]) -> Vec<u8> {
        let mut key_info = keyver as u16 | KEY_INFO_PAIRWISE;
        match message {
            1 => key_info |= KEY_INFO_ACK,
            2 => key_info |= KEY_INFO_MIC,
            3 => key_info |= KEY_INFO_ACK | KEY_INFO_MIC | KEY_INFO_INSTALL,
            4 => key_info |= KEY_INFO_MIC | KEY_INFO_SECURE,
            _ => {}
        }
        let total = KEY_DATA_OFFSET + key_data.len();
        let mut f = vec![0u8; total];
        f[0] = 2; // EAPOL version
        f[1] = EAPOL_TYPE_KEY;
        f[2..4].copy_from_slice(&((total - 4) as u16).to_be_bytes());
        f[4] = 2; // RSN descriptor
        f[5..7].copy_from_slice(&key_info.to_be_bytes());
        f[7..9].copy_from_slice(&16u16.to_be_bytes());
        f[9..17].copy_from_slice(&1u64.to_be_bytes());
        f[17..49].copy_from_slice(&nonce);
        for (i, b) in f[MIC_OFFSET..MIC_OFFSET + 16].iter_mut().enumerate() {
            *b = i as u8;
        }
        f[97..99].copy_from_slice(&(key_data.len() as u16).to_be_bytes());
        f[KEY_DATA_OFFSET..].copy_from_slice(key_data);
        f
    }

    #[test]
    fn identifies_all_four_messages() {
        for msg in 1..=4u8 {
            let frame = key_frame(msg, 2, [msg; 32], &[]);
            let parsed = parse(&frame).expect("eapol");
            assert_eq!(parsed.message, msg);
            assert_eq!(parsed.key_version, 2);
            assert!(parsed.pairwise);
            assert_eq!(parsed.nonce, [msg; 32]);
            assert_eq!(parsed.mic, {
                let mut m = [0u8; 16];
                for (i, b) in m.iter_mut().enumerate() {
                    *b = i as u8;
                }
                m
            });
            assert_eq!(parsed.frame, frame);
        }
    }

    /// Byte a byte, independente do helper `key_frame` (que usa a mesma
    /// codificação do parser): confere que um quadro Mensagem 1 real — key
    /// info 0x008a (Ack=1, versão CCMP=2, pairwise), big-endian como é no
    /// fio de verdade — é lido como message=1/key_version=2, não message=0.
    /// É exatamente o valor citado pela documentação de ferramentas de
    /// quebra de WPA pra mensagem 1 com CCMP.
    #[test]
    fn parses_real_wire_message1_key_info_as_big_endian() {
        let mut frame = vec![0u8; MIN_KEY_FRAME_LEN];
        frame[0] = 2; // versão EAPOL
        frame[1] = EAPOL_TYPE_KEY;
        frame[2..4].copy_from_slice(&((MIN_KEY_FRAME_LEN - 4) as u16).to_be_bytes());
        frame[4] = 2; // descritor RSN
        frame[5..7].copy_from_slice(&0x008au16.to_be_bytes()); // Ack=1, versão=2(CCMP), pairwise=1
        frame[7..9].copy_from_slice(&16u16.to_be_bytes()); // key length
        frame[9..17].copy_from_slice(&1u64.to_be_bytes()); // replay counter
        frame[17..49].fill(0xaa); // nonce
        frame[97..99].copy_from_slice(&0u16.to_be_bytes()); // sem key data

        let parsed = parse(&frame).expect("mensagem 1 real deveria parsear");
        assert_eq!(parsed.message, 1);
        assert_eq!(parsed.key_version, 2);
        assert!(parsed.pairwise);
        assert_eq!(parsed.replay_counter, 1);
    }

    #[test]
    fn extracts_pmkid_from_key_data() {
        let mut data = vec![0xdd, 20, 0x00, 0x0f, 0xac, 0x04];
        data.extend_from_slice(&[0xab; 16]);
        let frame = key_frame(1, 2, [0x11; 32], &data);
        let parsed = parse(&frame).expect("eapol");
        assert_eq!(parsed.pmkid, Some([0xab; 16]));
    }

    #[test]
    fn rejects_non_key_and_truncated_frames() {
        assert!(parse(&[2, 0, 0, 0]).is_none());
        let frame = key_frame(2, 2, [0; 32], &[]);
        assert!(parse(&frame[..frame.len() - 1]).is_none());
    }
}