/**
 * @file hccapx_serializer.c
 * @brief HCCAPX serializer for Hashcat-compatible output
 * 
 * Generates HCCAPX format files for password cracking with Hashcat
 */
#include "hccapx_serializer.h"

#include <string.h>
#include <stdlib.h>
#include <arpa/inet.h>
#include "esp_log.h"
#include "frame_analyzer.h"

static const char *TAG = "hccapx_serializer";

// HCCAPX constants
#define HCCAPX_SIGNATURE 0x58504348  // "HCPX"
#define HCCAPX_VERSION 4
#define HCCAPX_MAX_EAPOL_SIZE 256

// Static HCCAPX buffer
static hccapx_t hccapx = {
    .signature = HCCAPX_SIGNATURE,
    .version = HCCAPX_VERSION,
    .message_pair = 255,  // Invalid until we have a complete handshake
    .keyver = HCCAPX_KEYVER_WPA2
};

// State tracking
static unsigned message_ap = 0;
static unsigned message_sta = 0;
static unsigned eapol_source = 0;

// Helper: Check if array is all zeros
static bool is_array_zero(const uint8_t *array, unsigned size) {
    for (unsigned i = 0; i < size; i++) {
        if (array[i] != 0) return false;
    }
    return true;
}

void hccapx_serializer_init(const uint8_t *ssid, unsigned ssid_len) {
    hccapx_serializer_reset();
    
    if (ssid && ssid_len > 0 && ssid_len <= 32) {
        hccapx.essid_len = ssid_len;
        memcpy(hccapx.essid, ssid, ssid_len);
    }
    
    ESP_LOGI(TAG, "HCCAPX serializer initialized for SSID: %.*s", ssid_len, ssid);
}

void hccapx_serializer_reset(void) {
    memset(&hccapx, 0, sizeof(hccapx_t));
    hccapx.signature = HCCAPX_SIGNATURE;
    hccapx.version = HCCAPX_VERSION;
    hccapx.message_pair = 255;
    hccapx.keyver = HCCAPX_KEYVER_WPA2;
    
    message_ap = 0;
    message_sta = 0;
    eapol_source = 0;
    
    ESP_LOGI(TAG, "HCCAPX serializer reset");
}

// Save EAPOL packet to HCCAPX
static unsigned save_eapol(eapol_packet_t *eapol_packet, eapol_key_packet_t *eapol_key) {
    unsigned eapol_len = sizeof(eapol_packet_header_t) + ntohs(eapol_packet->header.packet_body_length);
    
    if (eapol_len > HCCAPX_MAX_EAPOL_SIZE) {
        ESP_LOGW(TAG, "EAPOL too long (%u > %u)", eapol_len, HCCAPX_MAX_EAPOL_SIZE);
        return 1;
    }
    
    hccapx.eapol_len = eapol_len;
    memcpy(hccapx.eapol, eapol_packet, eapol_len);
    memcpy(hccapx.keymic, eapol_key->key_mic, 16);
    
    // Clear MIC in saved EAPOL so Hashcat can calculate it
    // MIC is at offset 81 in EAPOL-Key packet (after 4-byte EAPOL header)
    memset(&hccapx.eapol[81], 0, 16);
    
    return 0;
}

// Handle M1 from AP
static void ap_message_m1(eapol_key_packet_t *eapol_key) {
    ESP_LOGI(TAG, "Processing M1 (AP)");
    message_ap = 1;
    memcpy(hccapx.nonce_ap, eapol_key->key_nonce, 32);
}

// Handle M3 from AP
static void ap_message_m3(eapol_packet_t *eapol, eapol_key_packet_t *eapol_key) {
    ESP_LOGI(TAG, "Processing M3 (AP)");
    message_ap = 3;
    
    if (message_ap == 0) {
        // No M1 seen, copy ANonce from M3
        memcpy(hccapx.nonce_ap, eapol_key->key_nonce, 32);
    }
    
    if (eapol_source == 2) {
        // Already have EAPOL from M2
        hccapx.message_pair = 2;
        return;
    }
    
    if (save_eapol(eapol, eapol_key) != 0) return;
    
    eapol_source = 3;
    if (message_sta == 2) {
        hccapx.message_pair = 3;
    }
}

// Handle AP messages (M1 or M3)
static void process_ap_message(data_frame_t *frame, eapol_packet_t *eapol, eapol_key_packet_t *eapol_key) {
    // Verify STA MAC consistency
    if (!is_array_zero(hccapx.mac_sta, 6) && 
        memcmp(frame->mac_header.addr1, hccapx.mac_sta, 6) != 0) {
        ESP_LOGW(TAG, "Different STA, ignoring");
        return;
    }
    
    if (message_ap == 0) {
        memcpy(hccapx.mac_ap, frame->mac_header.addr2, 6);
    }
    
    // M1 has empty Key MIC, M3 has filled Key MIC
    if (is_array_zero(eapol_key->key_mic, 16)) {
        ap_message_m1(eapol_key);
    } else {
        ap_message_m3(eapol, eapol_key);
    }
}

// Handle M2 from STA
static void sta_message_m2(eapol_packet_t *eapol, eapol_key_packet_t *eapol_key) {
    ESP_LOGI(TAG, "Processing M2 (STA)");
    message_sta = 2;
    memcpy(hccapx.nonce_sta, eapol_key->key_nonce, 32);
    
    if (save_eapol(eapol, eapol_key) != 0) return;
    
    eapol_source = 2;
    if (message_ap == 1) {
        hccapx.message_pair = 0;  // M1+M2
    }
}

// Handle M4 from STA
static void sta_message_m4(eapol_packet_t *eapol, eapol_key_packet_t *eapol_key) {
    ESP_LOGI(TAG, "Processing M4 (STA)");
    
    if (message_sta == 2 && eapol_source != 0) {
        ESP_LOGD(TAG, "Already have M2, M4 not needed");
        return;
    }
    
    if (message_ap == 0) {
        ESP_LOGW(TAG, "No AP message received yet");
        return;
    }
    
    if (eapol_source == 3) {
        hccapx.message_pair = 4;
        return;
    }
    
    if (save_eapol(eapol, eapol_key) != 0) return;
    
    eapol_source = 4;
    if (message_ap == 1) {
        hccapx.message_pair = 1;  // M1+M4
    }
    if (message_ap == 3) {
        hccapx.message_pair = 5;  // M3+M4
    }
}

// Handle STA messages (M2 or M4)
static void process_sta_message(data_frame_t *frame, eapol_packet_t *eapol, eapol_key_packet_t *eapol_key) {
    if (is_array_zero(hccapx.mac_sta, 6)) {
        memcpy(hccapx.mac_sta, frame->mac_header.addr2, 6);
    } else if (memcmp(frame->mac_header.addr2, hccapx.mac_sta, 6) != 0) {
        ESP_LOGW(TAG, "Different STA, ignoring");
        return;
    }
    
    // M2 has SNonce, M4 has empty SNonce
    if (!is_array_zero(eapol_key->key_nonce, 32)) {
        sta_message_m2(eapol, eapol_key);
    } else {
        sta_message_m4(eapol, eapol_key);
    }
}

void hccapx_serializer_add_frame(data_frame_t *frame) {
    // Parse EAPOL
    eapol_packet_t *eapol = parse_eapol_packet(frame);
    if (!eapol) return;
    
    eapol_key_packet_t *eapol_key = parse_eapol_key_packet(eapol);
    if (!eapol_key) return;
    
    // Determine direction: compare addr2 (source) with addr3 (BSSID)
    if (memcmp(frame->mac_header.addr2, frame->mac_header.addr3, 6) == 0) {
        // Source == BSSID => From AP
        process_ap_message(frame, eapol, eapol_key);
    } else if (memcmp(frame->mac_header.addr1, frame->mac_header.addr3, 6) == 0) {
        // Dest == BSSID => From STA
        process_sta_message(frame, eapol, eapol_key);
    } else {
        ESP_LOGW(TAG, "Unknown frame format");
    }
}

void hccapx_serializer_build(const handshake_data_t *handshake) {
    if (!handshake || !handshake->complete) {
        ESP_LOGW(TAG, "Incomplete handshake, cannot build HCCAPX");
        return;
    }
    
    hccapx_serializer_reset();
    
    // Copy data from handshake
    hccapx.essid_len = handshake->ssid_len;
    memcpy(hccapx.essid, handshake->ssid, handshake->ssid_len);
    
    memcpy(hccapx.mac_ap, handshake->ap_mac, 6);
    memcpy(hccapx.mac_sta, handshake->sta_mac, 6);
    
    memcpy(hccapx.nonce_ap, handshake->anonce, 32);
    memcpy(hccapx.nonce_sta, handshake->snonce, 32);
    
    memcpy(hccapx.keymic, handshake->mic, 16);
    
    if (handshake->eapol_len > 0 && handshake->eapol_len <= HCCAPX_MAX_EAPOL_SIZE) {
        hccapx.eapol_len = handshake->eapol_len;
        memcpy(hccapx.eapol, handshake->eapol, handshake->eapol_len);
    }
    
    hccapx.message_pair = handshake->message_pair;
    hccapx.keyver = HCCAPX_KEYVER_WPA2;
    
    ESP_LOGI(TAG, "HCCAPX built from handshake data");
}

hccapx_t* hccapx_serializer_get(void) {
    if (hccapx.message_pair == 255) {
        return NULL;  // No valid handshake
    }
    return &hccapx;
}

unsigned hccapx_serializer_get_size(void) {
    return sizeof(hccapx_t);
}

bool hccapx_serializer_is_valid(void) {
    return hccapx.message_pair != 255;
}