//! Helpers de teste: quadros 802.11/EAPOL sintéticos, radiotap e handshakes
//! válidos gerados a partir de uma senha conhecida. Só existe em `cfg(test)`.

use crate::crack::{compute_mic, compute_pmkid, compute_ptk, derive_pmk};
use crate::dot11;
use crate::eapol::{self, KEY_DATA_OFFSET, MIC_OFFSET};
use crate::handshake::Handshake;

const AP: [u8; 6] = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55];
const STA: [u8; 6] = [0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb];

pub fn key_frame(message: u8, keyver: u8, nonce: [u8; 32], key_data: &[u8]) -> Vec<u8> {
    use super::eapol::EAPOL_TYPE_KEY;
    let mut key_info = keyver as u16 | 0x0008; // pairwise
    match message {
        1 => key_info |= 0x0080,
        2 => key_info |= 0x0100,
        3 => key_info |= 0x0080 | 0x0100 | 0x0040,
        4 => key_info |= 0x0100 | 0x0200,
        _ => {}
    }
    let total = KEY_DATA_OFFSET + key_data.len();
    let mut f = vec![0u8; total];
    f[0] = 2;
    f[1] = EAPOL_TYPE_KEY;
    f[2..4].copy_from_slice(&((total - 4) as u16).to_be_bytes());
    f[4] = 2;
    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);
    f[97..99].copy_from_slice(&(key_data.len() as u16).to_be_bytes());
    f[KEY_DATA_OFFSET..].copy_from_slice(key_data);
    f
}

pub fn frame_with_eapol(
    from_ap: bool,
    ap: [u8; 6],
    sta: [u8; 6],
    da: [u8; 6],
    eapol: &[u8],
) -> Vec<u8> {
    let mut f = Vec::new();
    let fc: u16 = if from_ap { 0x0208 } else { 0x0108 };
    f.extend_from_slice(&fc.to_le_bytes());
    f.extend_from_slice(&0u16.to_le_bytes());
    if from_ap {
        f.extend_from_slice(&sta); // addr1 = destino
        f.extend_from_slice(&ap); // addr2 = BSSID
    } else {
        f.extend_from_slice(&ap); // addr1 = BSSID
        f.extend_from_slice(&sta); // addr2 = origem
    }
    f.extend_from_slice(&da);
    f.extend_from_slice(&[0u8; 2]);
    f.extend_from_slice(&[0xAA, 0xAA, 0x03, 0, 0, 0, 0x88, 0x8e]);
    f.extend_from_slice(eapol);
    f
}

pub fn radiotap_wrap(dot11: &[u8], freq: u16, signal: i8) -> Vec<u8> {
    // present: Channel | AntSignal
    let present: u32 = (1 << 3) | (1 << 5);
    let mut f = vec![0u8, 0, 13, 0];
    f.extend_from_slice(&present.to_le_bytes());
    f.extend_from_slice(&freq.to_le_bytes());
    f.extend_from_slice(&0xa0u16.to_le_bytes());
    f.push(signal as u8);
    f.extend_from_slice(dot11);
    f
}

fn ie(tag: u8, val: &[u8]) -> Vec<u8> {
    let mut v = vec![tag, val.len() as u8];
    v.extend_from_slice(val);
    v
}

fn rsn_ie(akms: &[u8]) -> Vec<u8> {
    let mut val = Vec::new();
    val.extend_from_slice(&1u16.to_le_bytes());
    val.extend_from_slice(&[0x00, 0x0f, 0xac, 4]);
    val.extend_from_slice(&1u16.to_le_bytes());
    val.extend_from_slice(&[0x00, 0x0f, 0xac, 4]);
    val.extend_from_slice(&(akms.len() as u16).to_le_bytes());
    for akm in akms {
        val.extend_from_slice(&[0x00, 0x0f, 0xac, *akm]);
    }
    val.extend_from_slice(&0x0040u16.to_le_bytes());
    ie(48, &val)
}

pub fn beacon_frame(bssid: [u8; 6], ssid: &str, channel: u8, akms: &[u8]) -> Vec<u8> {
    let mut f = Vec::new();
    f.extend_from_slice(&0x0080u16.to_le_bytes());
    f.extend_from_slice(&0u16.to_le_bytes());
    f.extend_from_slice(&[0xff; 6]);
    f.extend_from_slice(&bssid);
    f.extend_from_slice(&bssid);
    f.extend_from_slice(&[0u8; 2]);
    f.extend_from_slice(&[0u8; 8]);
    f.extend_from_slice(&100u16.to_le_bytes());
    f.extend_from_slice(&0x0031u16.to_le_bytes());
    f.extend_from_slice(&ie(0, ssid.as_bytes()));
    f.extend_from_slice(&ie(3, &[channel]));
    if !akms.is_empty() {
        f.extend_from_slice(&rsn_ie(akms));
    }
    f
}

pub fn probe_request(station: [u8; 6], ssid: &str) -> Vec<u8> {
    let mut f = Vec::new();
    f.extend_from_slice(&0x0040u16.to_le_bytes());
    f.extend_from_slice(&0u16.to_le_bytes());
    f.extend_from_slice(&[0xff; 6]);
    f.extend_from_slice(&station);
    f.extend_from_slice(&[0xff; 6]);
    f.extend_from_slice(&[0u8; 2]);
    f.extend_from_slice(&ie(0, ssid.as_bytes()));
    f
}

pub fn data_frame(
    to_ds: bool,
    from_ds: bool,
    addr1: [u8; 6],
    addr2: [u8; 6],
    addr3: [u8; 6],
    payload: &[u8],
) -> Vec<u8> {
    let mut fc: u16 = 0x0008;
    if to_ds {
        fc |= 0x0100;
    }
    if from_ds {
        fc |= 0x0200;
    }
    let mut f = Vec::new();
    f.extend_from_slice(&fc.to_le_bytes());
    f.extend_from_slice(&0u16.to_le_bytes());
    f.extend_from_slice(&addr1);
    f.extend_from_slice(&addr2);
    f.extend_from_slice(&addr3);
    f.extend_from_slice(&[0u8; 2]);
    f.extend_from_slice(payload);
    f
}

/// Gera (msg1, msg2) válidas para a senha informada, com MIC correto.
pub fn handshake_frames(
    ssid: &str,
    password: &str,
    keyver: u8,
    with_pmkid: bool,
) -> (Vec<u8>, Vec<u8>) {
    let anonce = [0xaa; 32];
    let snonce = [0xbb; 32];
    let pmk = derive_pmk(password, ssid);
    let mut m1_key_data = Vec::new();
    if with_pmkid {
        let pmkid = compute_pmkid(&pmk, &AP, &STA);
        m1_key_data.extend_from_slice(&[0xdd, 20, 0x00, 0x0f, 0xac, 0x04]);
        m1_key_data.extend_from_slice(&pmkid);
    }
    let m1 = key_frame(1, keyver, anonce, &m1_key_data);
    let mut m2 = key_frame(2, keyver, snonce, &[]);
    let ptk = compute_ptk(&pmk, &AP, &STA, &anonce, &snonce);
    let mic = compute_mic(keyver, &ptk[..16], &m2, MIC_OFFSET).expect("mic");
    m2[MIC_OFFSET..MIC_OFFSET + 16].copy_from_slice(&mic);
    (m1, m2)
}

/// Monta um `Handshake` diretamente dos quadros EAPOL, para testes de quebra
/// sem depender do pipeline de captura.
pub fn synthetic_handshake(
    ssid: &str,
    bssid: &str,
    station: &str,
    m1: &[u8],
    m2: &[u8],
) -> Handshake {
    let k1 = eapol::parse(m1).expect("msg1");
    let k2 = if m2.is_empty() {
        None
    } else {
        Some(eapol::parse(m2).expect("msg2"))
    };
    let mut messages = 1u8;
    if k2.is_some() {
        messages |= 2;
    }
    let hs = Handshake {
        bssid: bssid.to_string(),
        station: station.to_string(),
        ssid: Some(ssid.to_string()),
        anonce: Some(k1.nonce),
        snonce: k2.as_ref().map(|k| k.nonce),
        pmkid: k1.pmkid,
        key_version: k1.key_version,
        messages,
        eapol_msg2: k2.as_ref().map(|k| k.frame.clone()),
        raw: Vec::new(),
        first_seen_unix: 1,
        last_seen_unix: 2,
        password: None,
        tk: None,
    };
    // Exercita o parser de MAC do alvo.
    assert!(dot11::parse_mac(&hs.bssid).is_some());
    assert!(dot11::parse_mac(&hs.station).is_some());
    hs
}

pub fn ap() -> [u8; 6] {
    AP
}

pub fn sta() -> [u8; 6] {
    STA
}