//! Montagem dos handshakes WPA/WPA2 a partir dos quadros EAPOL vistos no ar.
//!
//! Um handshake é identificado pelo par (BSSID, estação). Guardamos as quatro
//! mensagens quando aparecem, o PMKID da mensagem 1 e os quadros crus
//! necessários para exportar um `.pcap` que o `aircrack-ng`/hashcat também
//! entendam.

use crate::dot11;
use crate::eapol::{self, EapolKey};
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub struct RawFrame {
    pub message: u8,
    pub raw: Vec<u8>,
}

#[derive(Debug, Clone, Default)]
pub struct Handshake {
    pub bssid: String,
    pub station: String,
    pub ssid: Option<String>,
    pub anonce: Option<[u8; 32]>,
    pub snonce: Option<[u8; 32]>,
    pub pmkid: Option<[u8; 16]>,
    pub key_version: u8,
    /// Bitmask das mensagens vistas (bit 0 = mensagem 1).
    pub messages: u8,
    /// Quadro EAPOL da mensagem 2, com o MIC — necessário para conferir a
    /// senha candidata.
    pub eapol_msg2: Option<Vec<u8>>,
    pub raw: Vec<RawFrame>,
    pub first_seen_unix: u64,
    pub last_seen_unix: u64,
    pub password: Option<String>,
    /// Temporal Key (CCMP/AES) derivada quando `password` é preenchido.
    /// Só pra TKIP (key_version 1) fica None: a decriptação aqui só cobre
    /// CCMP. Serve pra enxergar o tráfego em claro desse par (BSSID,
    /// estação) específico — não abre o tráfego de outros clientes, já que
    /// a PTK é efêmera por associação.
    pub tk: Option<[u8; 16]>,
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct HandshakeInfo {
    pub bssid: String,
    pub station: String,
    pub ssid: Option<String>,
    pub kind: String,
    pub messages: String,
    pub complete: bool,
    pub crackable: bool,
    pub pmkid: Option<String>,
    pub key_version: u8,
    pub password: Option<String>,
    pub raw_frames: usize,
    pub first_seen_unix: u64,
    pub last_seen_unix: u64,
}

impl Handshake {
    pub fn message_list(&self) -> String {
        let mut seen = Vec::new();
        for msg in 1..=4 {
            if self.messages & (1 << (msg - 1)) != 0 {
                seen.push(msg.to_string());
            }
        }
        seen.join(",")
    }

    pub fn kind(&self) -> &'static str {
        match (self.pmkid.is_some(), self.eapol_msg2.is_some()) {
            (true, true) => "EAPOL + PMKID",
            (true, false) => "PMKID",
            (false, _) => "EAPOL",
        }
    }

    pub fn is_complete(&self) -> bool {
        self.pmkid.is_some()
            || (self.anonce.is_some() && self.snonce.is_some() && self.eapol_msg2.is_some())
    }

    pub fn is_crackable(&self) -> bool {
        let ssid_ok = self.ssid.as_ref().is_some_and(|s| !s.is_empty());
        if !ssid_ok {
            return false;
        }
        if self.pmkid.is_some() {
            return true;
        }
        self.anonce.is_some() && self.eapol_msg2.is_some() && matches!(self.key_version, 1 | 2)
    }

    pub fn info(&self) -> HandshakeInfo {
        HandshakeInfo {
            bssid: self.bssid.clone(),
            station: self.station.clone(),
            ssid: self.ssid.clone(),
            kind: self.kind().to_string(),
            messages: self.message_list(),
            complete: self.is_complete(),
            crackable: self.is_crackable(),
            pmkid: self.pmkid.map(|p| hex(&p)),
            key_version: self.key_version,
            password: self.password.clone(),
            raw_frames: self.raw.len(),
            first_seen_unix: self.first_seen_unix,
            last_seen_unix: self.last_seen_unix,
        }
    }

    /// Quadros crus (radiotap + 802.11) das mensagens vistas, em ordem.
    pub fn raw_frames(&self) -> Vec<Vec<u8>> {
        let mut v = self.raw.clone();
        v.sort_by_key(|f| f.message);
        v.into_iter().map(|f| f.raw).collect()
    }

    fn store_raw(&mut self, message: u8, raw: &[u8]) {
        if let Some(existing) = self.raw.iter_mut().find(|f| f.message == message) {
            existing.raw = raw.to_vec();
        } else {
            self.raw.push(RawFrame {
                message,
                raw: raw.to_vec(),
            });
        }
    }
}

#[derive(Debug, Default)]
pub struct HandshakeTracker {
    map: HashMap<(String, String), Handshake>,
}

impl HandshakeTracker {
    pub fn new() -> Self {
        Self::default()
    }

    /// Processa um quadro 802.11. Retorna `true` quando o quadro era EAPOL e
    /// atualizou (ou criou) um handshake.
    pub fn ingest(
        &mut self,
        frame: &dot11::Frame<'_>,
        raw: &[u8],
        ssid_hint: Option<&str>,
        now: u64,
    ) -> bool {
        let Some(payload) = frame.eapol() else {
            return false;
        };
        let Some(key) = eapol::parse(payload) else {
            return false;
        };
        if key.message == 0 || !key.pairwise {
            return false;
        }

        // Em data frames com DS, o BSSID está em `frame.bssid`; a estação é
        // quem recebe (1/3, AP -> STA) ou quem transmite (2/4, STA -> AP).
        let bssid = frame.bssid_str();
        let station = match key.message {
            1 | 3 => frame.addr1.map(|m| dot11::mac_to_string(&m)),
            _ => frame.addr2.map(|m| dot11::mac_to_string(&m)),
        };
        let (Some(bssid), Some(station)) = (bssid, station) else {
            return false;
        };
        if bssid == station {
            return false;
        }

        let entry = self
            .map
            .entry((bssid.clone(), station.clone()))
            .or_insert_with(|| Handshake {
                bssid,
                station,
                first_seen_unix: now,
                ..Default::default()
            });
        if entry.ssid.is_none() {
            if let Some(ssid) = ssid_hint.filter(|s| !s.is_empty()) {
                entry.ssid = Some(ssid.to_string());
            }
        }
        if entry.first_seen_unix == 0 {
            entry.first_seen_unix = now;
        }
        entry.last_seen_unix = now;
        entry.messages |= 1 << (key.message - 1);
        entry.key_version = key.key_version;
        entry.store_raw(key.message, raw);
        apply_key(entry, &key);
        true
    }

    /// Atualiza o SSID de handshakes já vistos (quando o beacon chega depois
    /// do EAPOL, o que acontece em capturas começando no meio do handshake).
    pub fn update_ssid(&mut self, bssid: &str, ssid: &str) {
        if ssid.is_empty() {
            return;
        }
        for hs in self.map.values_mut() {
            if hs.bssid == bssid && hs.ssid.is_none() {
                hs.ssid = Some(ssid.to_string());
            }
        }
    }

    pub fn get(&self, bssid: &str, station: &str) -> Option<&Handshake> {
        self.map.get(&(bssid.to_string(), station.to_string()))
    }

    pub fn get_mut(&mut self, bssid: &str, station: &str) -> Option<&mut Handshake> {
        self.map.get_mut(&(bssid.to_string(), station.to_string()))
    }

    pub fn all(&self) -> Vec<&Handshake> {
        let mut v: Vec<_> = self.map.values().collect();
        v.sort_by(|a, b| {
            b.last_seen_unix
                .cmp(&a.last_seen_unix)
                .then_with(|| a.bssid.cmp(&b.bssid))
        });
        v
    }

    pub fn infos(&self) -> Vec<HandshakeInfo> {
        self.all().into_iter().map(Handshake::info).collect()
    }

    /// Primeiro handshake que pode ser quebrado; filtros opcionais por BSSID
    /// e estação.
    pub fn first_crackable(
        &self,
        bssid: Option<&str>,
        station: Option<&str>,
    ) -> Option<&Handshake> {
        self.all().into_iter().find(|h| {
            h.is_crackable()
                && bssid.map_or(true, |b| b.eq_ignore_ascii_case(&h.bssid))
                && station.map_or(true, |s| s.eq_ignore_ascii_case(&h.station))
        })
    }

    pub fn set_password(&mut self, bssid: &str, station: &str, password: &str) {
        if let Some(hs) = self.get_mut(bssid, station) {
            hs.password = Some(password.to_string());
            hs.tk = derive_tk(hs, password);
        }
    }

    /// TK já derivada pra esse par (BSSID, estação), se a senha da rede já
    /// foi quebrada e o handshake tinha os dois nonces (EAPOL completo — um
    /// handshake só de PMKID não dá pra derivar PTK/TK, só confirma a PSK).
    pub fn tk_for(&self, bssid: &str, station: &str) -> Option<[u8; 16]> {
        self.get(bssid, station).and_then(|hs| hs.tk)
    }

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

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

/// Deriva a TK (bytes 32..48 da PTK) a partir da senha quebrada, pra poder
/// decifrar CCMP do tráfego desse cliente. Só funciona com EAPOL/CCMP
/// completo (key_version 2); PMKID sozinho não tem os dois nonces pra
/// derivar PTK, e TKIP (key_version 1) usa RC4 + Michael, não implementado.
fn derive_tk(hs: &Handshake, password: &str) -> Option<[u8; 16]> {
    if hs.key_version != 2 {
        return None;
    }
    let ssid = hs.ssid.as_deref().filter(|s| !s.is_empty())?;
    let bssid = dot11::parse_mac(&hs.bssid)?;
    let station = dot11::parse_mac(&hs.station)?;
    let anonce = hs.anonce?;
    let snonce = hs.snonce?;
    let pmk = crate::crack::pmk_from_candidate(password, ssid.as_bytes())?;
    let ptk = crate::crack::compute_ptk(&pmk, &bssid, &station, &anonce, &snonce);
    let mut tk = [0u8; 16];
    tk.copy_from_slice(&ptk[32..48]);
    Some(tk)
}

fn apply_key(hs: &mut Handshake, key: &EapolKey) {
    match key.message {
        1 => {
            hs.anonce = Some(key.nonce);
            if key.pmkid.is_some() {
                hs.pmkid = key.pmkid;
            }
        }
        2 => {
            hs.snonce = Some(key.nonce);
            hs.eapol_msg2 = Some(key.frame.clone());
        }
        3 => hs.anonce = Some(key.nonce),
        4 => hs.snonce = Some(key.nonce),
        _ => {}
    }
}

fn hex(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        use std::fmt::Write;
        let _ = write!(s, "{b:02x}");
    }
    s
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eapol::KEY_DATA_OFFSET;
    use crate::testutil::{frame_with_eapol, key_frame};

    fn key_frame_for(message: u8, nonce: [u8; 32], key_data: &[u8]) -> Vec<u8> {
        key_frame(message, 2, nonce, key_data)
    }

    fn ap() -> [u8; 6] {
        [0x11; 6]
    }
    fn sta() -> [u8; 6] {
        [0x22; 6]
    }

    #[test]
    fn assembles_complete_handshake() {
        let mut tracker = HandshakeTracker::new();
        let m1 = key_frame_for(1, [0xaa; 32], &[]);
        let f1 = frame_with_eapol(true, ap(), sta(), [0x33; 6], &m1);
        let parsed1 = dot11::parse(&f1).expect("frame");
        assert!(tracker.ingest(&parsed1, &f1, Some("MinhaRede"), 10));

        let m2 = key_frame_for(2, [0xbb; 32], &[]);
        let f2 = frame_with_eapol(false, ap(), sta(), [0x33; 6], &m2);
        let parsed2 = dot11::parse(&f2).expect("frame");
        assert!(tracker.ingest(&parsed2, &f2, Some("MinhaRede"), 11));

        let hs = tracker
            .get("11:11:11:11:11:11", "22:22:22:22:22:22")
            .expect("hs");
        assert_eq!(hs.ssid.as_deref(), Some("MinhaRede"));
        assert_eq!(hs.messages, 0b11);
        assert!(hs.is_complete());
        assert!(hs.is_crackable());
        assert_eq!(hs.anonce, Some([0xaa; 32]));
        assert_eq!(hs.snonce, Some([0xbb; 32]));
        assert_eq!(hs.eapol_msg2.as_deref(), Some(&m2[..]));
        assert_eq!(hs.raw_frames().len(), 2);
    }

    #[test]
    fn captures_pmkid_from_message_one() {
        let mut data = vec![0xdd, 20, 0x00, 0x0f, 0xac, 0x04];
        data.extend_from_slice(&[0xcd; 16]);
        let m1 = key_frame_for(1, [0xaa; 32], &data);
        assert_eq!(m1.len(), KEY_DATA_OFFSET + data.len());
        let f1 = frame_with_eapol(true, ap(), sta(), [0x33; 6], &m1);
        let mut tracker = HandshakeTracker::new();
        let parsed = dot11::parse(&f1).expect("frame");
        assert!(tracker.ingest(&parsed, &f1, Some("Rede"), 1));
        let hs = tracker
            .get("11:11:11:11:11:11", "22:22:22:22:22:22")
            .expect("hs");
        assert_eq!(hs.pmkid, Some([0xcd; 16]));
        assert!(hs.is_complete());
        assert!(hs.is_crackable());
        assert_eq!(hs.kind(), "PMKID");
    }

    #[test]
    fn late_beacon_fills_ssid() {
        let m1 = key_frame_for(1, [1; 32], &[]);
        let f1 = frame_with_eapol(true, ap(), sta(), [0x33; 6], &m1);
        let mut tracker = HandshakeTracker::new();
        let parsed = dot11::parse(&f1).expect("frame");
        tracker.ingest(&parsed, &f1, None, 1);
        assert!(!tracker
            .get("11:11:11:11:11:11", "22:22:22:22:22:22")
            .unwrap()
            .is_crackable());
        tracker.update_ssid("11:11:11:11:11:11", "RedeTardia");
        assert!(
            tracker
                .get("11:11:11:11:11:11", "22:22:22:22:22:22")
                .unwrap()
                .ssid
                .as_deref()
                == Some("RedeTardia")
        );
    }
}