/**
 * @file frame_analyzer.c
 * @brief Frame analyzer implementation for EAPOL/Handshake/PMKID capture
 * 
 * Parses 802.11 data frames to extract WPA handshake components
 */
#include "frame_analyzer.h"

#include <string.h>
#include <stdlib.h>
#include <arpa/inet.h>
#include "esp_log.h"
#include "esp_err.h"
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"

static const char *TAG = "frame_analyzer";

// Global state
static uint8_t target_bssid[6] = {0};
static uint8_t target_ssid[33] = {0};
static uint8_t target_ssid_len = 0;
static search_type_t current_search_type = SEARCH_HANDSHAKE;
static handshake_data_t captured_handshake = {0};
static pmkid_item_t *captured_pmkids = NULL;
static SemaphoreHandle_t analyzer_mutex = NULL;
static bool capture_active = false;

// EAPOL/LLC constants
#define LLC_SNAP_HEADER_SIZE 8
#define EAPOL_TYPE_KEY 0x03

void frame_analyzer_init(void) {
    if (analyzer_mutex == NULL) {
        analyzer_mutex = xSemaphoreCreateMutex();
    }
    frame_analyzer_reset();
    ESP_LOGI(TAG, "Frame analyzer initialized");
}

// Internal reset function without mutex (called when mutex is already held)
static void frame_analyzer_reset_internal(void) {
    memset(&captured_handshake, 0, sizeof(captured_handshake));
    
    // Free PMKIDs
    while (captured_pmkids) {
        pmkid_item_t *next = captured_pmkids->next;
        free(captured_pmkids);
        captured_pmkids = next;
    }
    captured_pmkids = NULL;
}

void frame_analyzer_reset(void) {
    if (analyzer_mutex && xSemaphoreTake(analyzer_mutex, portMAX_DELAY)) {
        frame_analyzer_reset_internal();
        xSemaphoreGive(analyzer_mutex);
    }
}

void frame_analyzer_capture_start(search_type_t search_type, const uint8_t *bssid, const uint8_t *ssid, uint8_t ssid_len) {
    if (analyzer_mutex && xSemaphoreTake(analyzer_mutex, portMAX_DELAY)) {
        frame_analyzer_reset_internal();  // Use internal version (mutex already held)
        
        current_search_type = search_type;
        memcpy(target_bssid, bssid, 6);
        
        if (ssid && ssid_len > 0) {
            memcpy(target_ssid, ssid, ssid_len);
            target_ssid_len = ssid_len;
            memcpy(captured_handshake.ssid, ssid, ssid_len);
            captured_handshake.ssid_len = ssid_len;
        }
        
        memcpy(captured_handshake.ap_mac, bssid, 6);
        capture_active = true;
        
        ESP_LOGI(TAG, "Capture started for BSSID: %02x:%02x:%02x:%02x:%02x:%02x, SSID: %s",
                 bssid[0], bssid[1], bssid[2], bssid[3], bssid[4], bssid[5],
                 ssid_len > 0 ? (char*)target_ssid : "(unknown)");
        
        xSemaphoreGive(analyzer_mutex);
    }
}

void frame_analyzer_capture_stop(void) {
    if (analyzer_mutex && xSemaphoreTake(analyzer_mutex, portMAX_DELAY)) {
        capture_active = false;
        ESP_LOGI(TAG, "Capture stopped");
        xSemaphoreGive(analyzer_mutex);
    }
}

bool is_frame_bssid_matching(const wifi_promiscuous_pkt_t *frame, const uint8_t *target_bssid) {
    const data_frame_t *data_frame = (const data_frame_t *)frame->payload;
    
    // Check addr1, addr2, addr3 for BSSID match
    if (memcmp(data_frame->mac_header.addr1, target_bssid, 6) == 0 ||
        memcmp(data_frame->mac_header.addr2, target_bssid, 6) == 0 ||
        memcmp(data_frame->mac_header.addr3, target_bssid, 6) == 0) {
        return true;
    }
    return false;
}

eapol_packet_t* parse_eapol_packet(data_frame_t *frame) {
    // Skip MAC header to get LLC/SNAP header
    uint8_t *body = frame->body;
    
    // Check LLC/SNAP header
    llc_snap_header_t *llc = (llc_snap_header_t *)body;
    if (llc->snap_dsap != 0xAA || llc->snap_ssap != 0xAA || llc->control != 0x03) {
        return NULL;
    }
    
    // Check EtherType for EAPOL (0x888e)
    uint16_t *ether_type = (uint16_t *)(body + 6);
    if (ntohs(*ether_type) != ETHER_TYPE_EAPOL) {
        return NULL;
    }
    
    // Return pointer to EAPOL packet (after LLC/SNAP header)
    return (eapol_packet_t *)(body + LLC_SNAP_HEADER_SIZE);
}

eapol_key_packet_t* parse_eapol_key_packet(eapol_packet_t *eapol_packet) {
    if (eapol_packet->header.packet_type != EAPOL_KEY) {
        return NULL;
    }
    return (eapol_key_packet_t *)eapol_packet->packet_body;
}

pmkid_item_t* parse_pmkid(eapol_key_packet_t *eapol_key_packet) {
    uint16_t key_data_len = ntohs(eapol_key_packet->key_data_length);
    if (key_data_len == 0) {
        return NULL;
    }
    
    pmkid_item_t *head = NULL;
    pmkid_item_t *tail = NULL;
    
    uint8_t *key_data = eapol_key_packet->key_data;
    unsigned offset = 0;
    
    while (offset < key_data_len) {
        key_data_field_t *field = (key_data_field_t *)(key_data + offset);
        
        if (field->type == KEY_DATA_TYPE) {
            uint32_t oui = ntohl(field->oui << 8);
            if (oui == KEY_DATA_OUI_IEEE80211 && field->data_type == KEY_DATA_DATA_TYPE_PMKID_KDE) {
                // Found PMKID
                pmkid_item_t *item = malloc(sizeof(pmkid_item_t));
                if (item) {
                    memcpy(item->pmkid, field->data, 16);
                    item->next = NULL;
                    
                    if (!head) {
                        head = item;
                        tail = item;
                    } else {
                        tail->next = item;
                        tail = item;
                    }
                    ESP_LOGI(TAG, "PMKID captured!");
                }
            }
        }
        
        offset += 2 + field->length;  // type + length + data
        if (field->length == 0) break;  // Prevent infinite loop
    }
    
    return head;
}

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

// Process handshake message from AP (M1 or M3)
static void process_ap_message(data_frame_t *frame, eapol_packet_t *eapol, eapol_key_packet_t *eapol_key) {
    // Copy AP MAC if not set
    if (is_zero_array(captured_handshake.ap_mac, 6)) {
        memcpy(captured_handshake.ap_mac, frame->mac_header.addr2, 6);
    }
    
    // Determine M1 or M3 by checking Key MIC
    // M1: Key MIC is empty, M3: Key MIC is present
    if (is_zero_array(eapol_key->key_mic, 16)) {
        // This is M1 - contains ANonce
        ESP_LOGI(TAG, "Captured M1 (AP -> STA, ANonce)");
        memcpy(captured_handshake.anonce, eapol_key->key_nonce, 32);
        
        if (captured_handshake.state < HANDSHAKE_STATE_M1_CAPTURED) {
            captured_handshake.state = HANDSHAKE_STATE_M1_CAPTURED;
        }
    } else {
        // This is M3 - also contains ANonce
        ESP_LOGI(TAG, "Captured M3 (AP -> STA)");
        
        if (captured_handshake.state < HANDSHAKE_STATE_M3_CAPTURED) {
            if (captured_handshake.state < HANDSHAKE_STATE_M1_CAPTURED) {
                // Didn't see M1, copy ANonce from M3
                memcpy(captured_handshake.anonce, eapol_key->key_nonce, 32);
            }
            captured_handshake.state = HANDSHAKE_STATE_M3_CAPTURED;
        }
        
        // If we already have M2, we have a complete handshake
        if (captured_handshake.state >= HANDSHAKE_STATE_M2_CAPTURED) {
            captured_handshake.complete = true;
            captured_handshake.message_pair = 2;  // M1+M2 or M3+M2
            ESP_LOGI(TAG, "*** HANDSHAKE COMPLETE! ***");
        }
    }
}

// Process handshake message from STA (M2 or M4)
static void process_sta_message(data_frame_t *frame, eapol_packet_t *eapol, eapol_key_packet_t *eapol_key) {
    // Copy STA MAC if not set
    if (is_zero_array(captured_handshake.sta_mac, 6)) {
        memcpy(captured_handshake.sta_mac, frame->mac_header.addr2, 6);
    }
    
    // Determine M2 or M4 by checking SNonce
    // M2: SNonce is present, M4: SNonce is empty
    if (!is_zero_array(eapol_key->key_nonce, 32)) {
        // This is M2 - contains SNonce and MIC
        ESP_LOGI(TAG, "Captured M2 (STA -> AP, SNonce + MIC)");
        memcpy(captured_handshake.snonce, eapol_key->key_nonce, 32);
        memcpy(captured_handshake.mic, eapol_key->key_mic, 16);
        
        // Save EAPOL packet for cracking
        uint16_t eapol_len = sizeof(eapol_packet_header_t) + ntohs(eapol->header.packet_body_length);
        if (eapol_len <= sizeof(captured_handshake.eapol)) {
            memcpy(captured_handshake.eapol, eapol, eapol_len);
            captured_handshake.eapol_len = eapol_len;
            
            // Zero out MIC in saved EAPOL for verification during cracking
            memset(captured_handshake.eapol + 81, 0, 16);
        }
        
        if (captured_handshake.state < HANDSHAKE_STATE_M2_CAPTURED) {
            captured_handshake.state = HANDSHAKE_STATE_M2_CAPTURED;
        }
        
        // If we have M1, we have a complete handshake (M1+M2)
        if (captured_handshake.state >= HANDSHAKE_STATE_M1_CAPTURED) {
            captured_handshake.complete = true;
            captured_handshake.message_pair = 0;  // M1+M2
            ESP_LOGI(TAG, "*** HANDSHAKE COMPLETE! ***");
        }
    } else {
        // This is M4
        ESP_LOGI(TAG, "Captured M4 (STA -> AP)");
        
        if (captured_handshake.state < HANDSHAKE_STATE_M4_CAPTURED) {
            captured_handshake.state = HANDSHAKE_STATE_M4_CAPTURED;
        }
    }
}

void frame_analyzer_process_frame(wifi_promiscuous_pkt_t *pkt, wifi_promiscuous_pkt_type_t type) {
    if (!capture_active || type != WIFI_PKT_DATA) {
        return;
    }
    
    if (!is_frame_bssid_matching(pkt, target_bssid)) {
        return;
    }
    
    data_frame_t *frame = (data_frame_t *)pkt->payload;
    
    // Parse EAPOL packet
    eapol_packet_t *eapol = parse_eapol_packet(frame);
    if (!eapol) {
        return;
    }
    
    // Parse EAPOL-Key packet
    eapol_key_packet_t *eapol_key = parse_eapol_key_packet(eapol);
    if (!eapol_key) {
        return;
    }
    
    ESP_LOGD(TAG, "Got EAPOL-Key frame");
    
    if (analyzer_mutex && xSemaphoreTake(analyzer_mutex, pdMS_TO_TICKS(10))) {
        if (current_search_type == SEARCH_PMKID) {
            // Looking for PMKID
            pmkid_item_t *pmkids = parse_pmkid(eapol_key);
            if (pmkids) {
                // Add to list
                if (!captured_pmkids) {
                    captured_pmkids = pmkids;
                } else {
                    pmkid_item_t *tail = captured_pmkids;
                    while (tail->next) tail = tail->next;
                    tail->next = pmkids;
                }
            }
        } else if (current_search_type == SEARCH_HANDSHAKE) {
            // Looking for handshake
            // Determine frame direction by comparing addr2 (source) with addr3 (BSSID)
            if (memcmp(frame->mac_header.addr2, frame->mac_header.addr3, 6) == 0) {
                // Source == BSSID, this is from AP
                process_ap_message(frame, eapol, eapol_key);
            } else if (memcmp(frame->mac_header.addr1, frame->mac_header.addr3, 6) == 0) {
                // Dest == BSSID, this is from STA
                process_sta_message(frame, eapol, eapol_key);
            }
        }
        
        xSemaphoreGive(analyzer_mutex);
    }
}

handshake_state_t frame_analyzer_get_handshake_state(void) {
    return captured_handshake.state;
}

const handshake_data_t* frame_analyzer_get_handshake(void) {
    return &captured_handshake;
}

bool frame_analyzer_handshake_complete(void) {
    return captured_handshake.complete;
}

pmkid_item_t* frame_analyzer_get_pmkids(void) {
    return captured_pmkids;
}