/*
 * fb_stealth.c
 * OPSEC: MAC randomization, TX power control, passive scan.
 */
#include "sdkconfig.h"
#include "fb_stealth.h"

#ifdef CONFIG_MODULE_FALLBACK

#include <string.h>
#include "esp_log.h"
#include "esp_wifi.h"
#include "esp_random.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

static const char *TAG = "FB_STEALTH";

/* ============================================================
 * MAC randomization
 * ============================================================ */

static uint8_t s_orig_mac[6] = {0};
static bool s_mac_saved = false;

void fb_stealth_save_original_mac(void)
{
    if (esp_wifi_get_mac(WIFI_IF_STA, s_orig_mac) == ESP_OK) {
        s_mac_saved = true;
        ESP_LOGI(TAG, "Original MAC: %02X:%02X:%02X:%02X:%02X:%02X",
                 s_orig_mac[0], s_orig_mac[1], s_orig_mac[2],
                 s_orig_mac[3], s_orig_mac[4], s_orig_mac[5]);
    }
}

void fb_stealth_randomize_mac(void)
{
    uint8_t mac[6];
    esp_fill_random(mac, 6);
    mac[0] &= 0xFE;  /* unicast */
    mac[0] |= 0x02;  /* locally administered */

    esp_wifi_disconnect();
    vTaskDelay(pdMS_TO_TICKS(50));

    esp_err_t err = esp_wifi_set_mac(WIFI_IF_STA, mac);
    if (err == ESP_OK) {
        ESP_LOGI(TAG, "MAC randomized: %02X:%02X:%02X:%02X:%02X:%02X",
                 mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    } else {
        ESP_LOGW(TAG, "MAC set failed: %s", esp_err_to_name(err));
    }
}

void fb_stealth_restore_mac(void)
{
    if (s_mac_saved) {
        esp_wifi_disconnect();
        vTaskDelay(pdMS_TO_TICKS(50));
        esp_wifi_set_mac(WIFI_IF_STA, s_orig_mac);
        ESP_LOGI(TAG, "MAC restored: %02X:%02X:%02X:%02X:%02X:%02X",
                 s_orig_mac[0], s_orig_mac[1], s_orig_mac[2],
                 s_orig_mac[3], s_orig_mac[4], s_orig_mac[5]);
    }
}

void fb_stealth_get_current_mac(uint8_t mac[6])
{
    esp_wifi_get_mac(WIFI_IF_STA, mac);
}

/* ============================================================
 * TX power control
 * ============================================================ */

void fb_stealth_low_tx_power(void)
{
    esp_err_t err = esp_wifi_set_max_tx_power(32);  /* 8 dBm */
    if (err == ESP_OK) {
        ESP_LOGI(TAG, "TX power reduced to 8 dBm");
    } else {
        ESP_LOGW(TAG, "TX power set failed: %s", esp_err_to_name(err));
    }
}

void fb_stealth_restore_tx_power(void)
{
    esp_wifi_set_max_tx_power(80);  /* 20 dBm */
    ESP_LOGI(TAG, "TX power restored to 20 dBm");
}

/* ============================================================
 * Passive scan — promiscuous mode beacon capture
 * ============================================================ */

typedef struct {
    unsigned frame_ctrl:16;
    unsigned duration_id:16;
    uint8_t  addr1[6];
    uint8_t  addr2[6];
    uint8_t  addr3[6];
    unsigned seq_ctrl:16;
} __attribute__((packed)) wifi_mgmt_hdr_t;

#define BEACON_FIXED_LEN  12

static fb_scan_ap_t s_scan_results[FB_MAX_SCAN_APS];
static volatile int s_scan_count = 0;

static int find_bssid(const uint8_t bssid[6])
{
    for (int i = 0; i < s_scan_count; i++) {
        if (memcmp(s_scan_results[i].bssid, bssid, 6) == 0)
            return i;
    }
    return -1;
}

static void passive_scan_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;
    wifi_mgmt_hdr_t *hdr = (wifi_mgmt_hdr_t *)pkt->payload;

    uint16_t fc = hdr->frame_ctrl;
    uint8_t subtype = (fc >> 4) & 0x0F;
    if (subtype != 8 && subtype != 5) return;  /* beacon or probe_resp */

    const uint8_t *bssid = hdr->addr3;

    int idx = find_bssid(bssid);
    if (idx >= 0) {
        if (pkt->rx_ctrl.rssi > s_scan_results[idx].rssi) {
            s_scan_results[idx].rssi = pkt->rx_ctrl.rssi;
        }
        return;
    }

    if (s_scan_count >= FB_MAX_SCAN_APS) return;

    size_t hdr_len = sizeof(wifi_mgmt_hdr_t);
    size_t body_offset = hdr_len + BEACON_FIXED_LEN;

    if ((int)pkt->rx_ctrl.sig_len < (int)(body_offset + 2))
        return;

    const uint8_t *body = pkt->payload + body_offset;
    size_t body_len = pkt->rx_ctrl.sig_len - body_offset;
    if (body_len > 4) body_len -= 4;

    fb_scan_ap_t *ap = &s_scan_results[s_scan_count];
    memset(ap, 0, sizeof(*ap));
    memcpy(ap->bssid, bssid, 6);
    ap->rssi = pkt->rx_ctrl.rssi;
    ap->channel = pkt->rx_ctrl.channel;
    ap->auth_mode = 0;

    size_t pos = 0;
    while (pos + 2 <= body_len) {
        uint8_t tag_id = body[pos];
        uint8_t tag_len = body[pos + 1];

        if (pos + 2 + tag_len > body_len) break;

        if (tag_id == 0) {
            size_t ssid_len = tag_len;
            if (ssid_len > 32) ssid_len = 32;
            memcpy(ap->ssid, body + pos + 2, ssid_len);
            ap->ssid[ssid_len] = '\0';
        } else if (tag_id == 48) {
            ap->auth_mode = 3;
        } else if (tag_id == 221) {
            if (tag_len >= 4 &&
                body[pos + 2] == 0x00 && body[pos + 3] == 0x50 &&
                body[pos + 4] == 0xF2 && body[pos + 5] == 0x01) {
                if (ap->auth_mode == 0) ap->auth_mode = 2;
            }
        }

        pos += 2 + tag_len;
    }

    s_scan_count++;
}

int fb_stealth_passive_scan(int duration_ms)
{
    s_scan_count = 0;
    memset(s_scan_results, 0, sizeof(s_scan_results));

    esp_wifi_disconnect();
    vTaskDelay(pdMS_TO_TICKS(100));

    esp_err_t ret = esp_wifi_set_promiscuous_rx_cb(passive_scan_cb);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Promiscuous CB failed: %s", esp_err_to_name(ret));
        return 0;
    }

    wifi_promiscuous_filter_t filter = {
        .filter_mask = WIFI_PROMIS_FILTER_MASK_MGMT
    };
    esp_wifi_set_promiscuous_filter(&filter);

    ret = esp_wifi_set_promiscuous(true);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Promiscuous enable failed: %s", esp_err_to_name(ret));
        return 0;
    }

    ESP_LOGI(TAG, "Passive scan started (%d ms)", duration_ms);

    int channels = 13;
    int hop_ms = 200;
    int elapsed = 0;

    while (elapsed < duration_ms) {
        for (int ch = 1; ch <= channels && elapsed < duration_ms; ch++) {
            esp_wifi_set_channel(ch, WIFI_SECOND_CHAN_NONE);
            vTaskDelay(pdMS_TO_TICKS(hop_ms));
            elapsed += hop_ms;
        }
    }

    esp_wifi_set_promiscuous(false);

    ESP_LOGI(TAG, "Passive scan done: %d unique APs", s_scan_count);
    return s_scan_count;
}

int fb_stealth_get_scan_results(fb_scan_ap_t *out, int max_count)
{
    int count = s_scan_count;
    if (count > max_count) count = max_count;
    memcpy(out, s_scan_results, count * sizeof(fb_scan_ap_t));
    return count;
}

#else /* !CONFIG_MODULE_FALLBACK — empty stubs */

#include <string.h>

void fb_stealth_save_original_mac(void) {}
void fb_stealth_randomize_mac(void) {}
void fb_stealth_restore_mac(void) {}
void fb_stealth_get_current_mac(uint8_t mac[6]) { memset(mac, 0, 6); }
void fb_stealth_low_tx_power(void) {}
void fb_stealth_restore_tx_power(void) {}
int fb_stealth_passive_scan(int duration_ms) { (void)duration_ms; return 0; }
int fb_stealth_get_scan_results(fb_scan_ap_t *out, int max_count) { (void)out; (void)max_count; return 0; }

#endif /* CONFIG_MODULE_FALLBACK */