/*
 * hp_wifi_monitor.c
 * WiFi promiscuous-mode monitor: probe requests, deauth frames,
 * beacon flood, EAPOL capture detection.
 *
 * Sends EVT| events via event_send().
 * Conflict guard: refuses to start if the fakeAP sniffer is active.
 */
#include "sdkconfig.h"

#ifdef CONFIG_MODULE_HONEYPOT

#include <stdio.h>
#include <string.h>
#include <stdatomic.h>

#include "esp_wifi.h"
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#include "utils.h"
#include "event_format.h"
#include "hp_config.h"
#include "hp_wifi_monitor.h"

#define TAG "HP_WIFI"

#define WIFI_MON_STACK  4096
#define WIFI_MON_PRIO   4
#define WIFI_MON_CORE   1

/* Rate-limit counters (only report every N-th event) */
#define PROBE_RATE_LIMIT   10
#define DEAUTH_RATE_LIMIT   5
#define BEACON_RATE_LIMIT  20
#define EAPOL_RATE_LIMIT    3

/* Beacon flood detection: N beacons in BEACON_WINDOW_MS from same src */
#define BEACON_FLOOD_THRESHOLD 50
#define BEACON_WINDOW_MS       5000

/* ============================================================
 * State
 * ============================================================ */
static atomic_bool       mon_running = false;
static atomic_bool       mon_stop_req = false;
static TaskHandle_t      mon_task = NULL;

static uint32_t cnt_probe   = 0;
static uint32_t cnt_deauth  = 0;
static uint32_t cnt_beacon  = 0;
static uint32_t cnt_eapol   = 0;

/* Multi-source beacon flood tracker */
#define BEACON_TRACK_MAX 4

typedef struct {
    uint8_t  mac[6];
    uint32_t count;
    uint32_t start;
    bool     alerted;
} beacon_tracker_t;

static beacon_tracker_t  beacon_trackers[BEACON_TRACK_MAX];
static int               beacon_tracker_count = 0;

/* ============================================================
 * IEEE 802.11 helpers
 * ============================================================ */

/* Frame control subtypes */
#define WLAN_FC_TYPE_MGMT    0x00
#define WLAN_FC_STYPE_PROBE  0x40   /* Probe Request */
#define WLAN_FC_STYPE_BEACON 0x80   /* Beacon */
#define WLAN_FC_STYPE_DEAUTH 0xC0   /* Deauthentication */

/* EAPOL: data frame with ethertype 0x888E */
#define ETHERTYPE_EAPOL 0x888E

static void mac_to_str(const uint8_t *mac, char *buf, size_t len)
{
    snprintf(buf, len, "%02x:%02x:%02x:%02x:%02x:%02x",
        mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
}

/* ============================================================
 * Beacon flood helper — find or create tracker for MAC
 * ============================================================ */
static beacon_tracker_t *beacon_find_or_create(const uint8_t *mac, uint32_t now)
{
    /* Search existing */
    for (int i = 0; i < beacon_tracker_count; i++) {
        if (memcmp(beacon_trackers[i].mac, mac, 6) == 0)
            return &beacon_trackers[i];
    }

    /* Evict oldest if full */
    if (beacon_tracker_count >= BEACON_TRACK_MAX) {
        int oldest = 0;
        for (int i = 1; i < beacon_tracker_count; i++) {
            if (beacon_trackers[i].start < beacon_trackers[oldest].start)
                oldest = i;
        }
        if (oldest < beacon_tracker_count - 1)
            beacon_trackers[oldest] = beacon_trackers[beacon_tracker_count - 1];
        beacon_tracker_count--;
    }

    beacon_tracker_t *t = &beacon_trackers[beacon_tracker_count++];
    memcpy(t->mac, mac, 6);
    t->count = 0;
    t->start = now;
    t->alerted = false;
    return t;
}

/* ============================================================
 * Promiscuous RX callback
 * ============================================================ */
static void wifi_monitor_rx_cb(void *buf, wifi_promiscuous_pkt_type_t type)
{
    if (!mon_running)
        return;

    const wifi_promiscuous_pkt_t *pkt = (const wifi_promiscuous_pkt_t *)buf;
    const uint8_t *frame = pkt->payload;
    uint16_t frame_len = pkt->rx_ctrl.sig_len;

    if (frame_len < 24)
        return;

    uint8_t fc0 = frame[0];
    uint8_t fc_type = fc0 & 0x0C;       /* bits 2-3 */
    uint8_t fc_subtype = fc0 & 0xF0;    /* bits 4-7 */

    /* Source MAC (addr2 = transmitter) at offset 10 */
    const uint8_t *src_mac = &frame[10];
    char mac_str[18];

    if (type == WIFI_PKT_MGMT) {
        if (fc_type == WLAN_FC_TYPE_MGMT) {

            /* --- Probe Request --- */
            if (fc_subtype == WLAN_FC_STYPE_PROBE) {
                cnt_probe++;
                if ((cnt_probe % PROBE_RATE_LIMIT) == 1) {
                    mac_to_str(src_mac, mac_str, sizeof(mac_str));
                    char detail[64];
                    snprintf(detail, sizeof(detail), "count=%lu",
                        (unsigned long)cnt_probe);
                    event_send("WIFI_PROBE", "LOW",
                        mac_str, "0.0.0.0", 0, 0, detail, NULL);
                }
                return;
            }

            /* --- Deauthentication --- */
            if (fc_subtype == WLAN_FC_STYPE_DEAUTH) {
                cnt_deauth++;
                if ((cnt_deauth % DEAUTH_RATE_LIMIT) == 1) {
                    mac_to_str(src_mac, mac_str, sizeof(mac_str));
                    char detail[64];
                    snprintf(detail, sizeof(detail), "reason=%d count=%lu",
                        (frame_len >= 26) ? (frame[24] | (frame[25] << 8)) : 0,
                        (unsigned long)cnt_deauth);
                    event_send("WIFI_DEAUTH", "HIGH",
                        mac_str, "0.0.0.0", 0, 0, detail, NULL);
                }
                return;
            }

            /* --- Beacon flood detection (multi-source) --- */
            if (fc_subtype == WLAN_FC_STYPE_BEACON) {
                uint32_t now = (uint32_t)(xTaskGetTickCount() *
                    portTICK_PERIOD_MS);

                beacon_tracker_t *bt = beacon_find_or_create(src_mac, now);

                if ((now - bt->start) >= BEACON_WINDOW_MS) {
                    /* Window expired, reset */
                    bt->start = now;
                    bt->count = 1;
                    bt->alerted = false;
                } else {
                    bt->count++;
                    if (bt->count >= BEACON_FLOOD_THRESHOLD && !bt->alerted) {
                        bt->alerted = true;
                        cnt_beacon++;
                        mac_to_str(src_mac, mac_str, sizeof(mac_str));
                        char detail[64];
                        snprintf(detail, sizeof(detail),
                            "beacons=%lu window_ms=%d",
                            (unsigned long)bt->count,
                            BEACON_WINDOW_MS);
                        event_send("WIFI_BEACON_FLOOD", "HIGH",
                            mac_str, "0.0.0.0", 0, 0, detail, NULL);
                    }
                }
                return;
            }
        }
    }

    /* --- EAPOL detection (data frames with 802.1X ethertype) --- */
    if (type == WIFI_PKT_DATA && frame_len >= 36) {
        /* LLC/SNAP header starts at offset 24 for data frames:
         * 24: AA AA 03 00 00 00 [ethertype_hi] [ethertype_lo] */
        if (frame[24] == 0xAA && frame[25] == 0xAA && frame[26] == 0x03) {
            uint16_t ethertype = (frame[30] << 8) | frame[31];
            if (ethertype == ETHERTYPE_EAPOL) {
                cnt_eapol++;
                if ((cnt_eapol % EAPOL_RATE_LIMIT) == 1) {
                    mac_to_str(src_mac, mac_str, sizeof(mac_str));
                    char detail[64];
                    snprintf(detail, sizeof(detail), "count=%lu",
                        (unsigned long)cnt_eapol);
                    event_send("WIFI_EAPOL", "CRITICAL",
                        mac_str, "0.0.0.0", 0, 0, detail, NULL);
                }
            }
        }
    }
}

/* ============================================================
 * Monitor task (just keeps alive, callback does the work)
 * ============================================================ */
static void wifi_monitor_task(void *arg)
{
    (void)arg;

    esp_err_t err = esp_wifi_set_promiscuous_rx_cb(wifi_monitor_rx_cb);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "set_promiscuous_rx_cb failed: %s", esp_err_to_name(err));
        goto done;
    }

    err = esp_wifi_set_promiscuous(true);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "set_promiscuous(true) failed: %s", esp_err_to_name(err));
        goto done;
    }

    /* Filter: management + data frames only */
    wifi_promiscuous_filter_t filter = {
        .filter_mask = WIFI_PROMIS_FILTER_MASK_MGMT |
                       WIFI_PROMIS_FILTER_MASK_DATA
    };
    esp_wifi_set_promiscuous_filter(&filter);

    ESP_LOGI(TAG, "WiFi monitor started");
    mon_running = true;

    /* Idle loop, checking for stop request */
    while (!mon_stop_req) {
        vTaskDelay(pdMS_TO_TICKS(500));
    }

    esp_wifi_set_promiscuous(false);
    esp_wifi_set_promiscuous_rx_cb(NULL);

done:
    mon_running = false;
    mon_stop_req = false;
    ESP_LOGI(TAG, "WiFi monitor stopped");
    mon_task = NULL;
    vTaskDelete(NULL);
}

/* ============================================================
 * Public API
 * ============================================================ */
void hp_wifi_monitor_start(void)
{
    if (mon_running || mon_task) {
        ESP_LOGW(TAG, "WiFi monitor already running");
        return;
    }

    cnt_probe = cnt_deauth = cnt_beacon = cnt_eapol = 0;
    memset(beacon_trackers, 0, sizeof(beacon_trackers));
    beacon_tracker_count = 0;
    mon_stop_req = false;

    BaseType_t ret = xTaskCreatePinnedToCore(wifi_monitor_task, "hp_wifi",
        WIFI_MON_STACK, NULL, WIFI_MON_PRIO, &mon_task, WIFI_MON_CORE);
    if (ret != pdPASS) {
        ESP_LOGE(TAG, "Failed to create WiFi monitor task");
        mon_task = NULL;
    }
}

void hp_wifi_monitor_stop(void)
{
    if (!mon_running && !mon_task) {
        ESP_LOGW(TAG, "WiFi monitor not running");
        return;
    }
    mon_stop_req = true;
    ESP_LOGI(TAG, "WiFi monitor stop requested");
}

bool hp_wifi_monitor_running(void)
{
    return mon_running;
}

int hp_wifi_monitor_status(char *buf, size_t len)
{
    return snprintf(buf, len,
        "running=%s probes=%lu deauth=%lu beacon_flood=%lu eapol=%lu",
        mon_running ? "yes" : "no",
        (unsigned long)cnt_probe,
        (unsigned long)cnt_deauth,
        (unsigned long)cnt_beacon,
        (unsigned long)cnt_eapol);
}

#endif /* CONFIG_MODULE_HONEYPOT */