// Host harness: drives the real kernel supplicant through a 4-way handshake.
// Frames come in / go out as hex on stdio so an independent Python AP can be
// the oracle.
#include <cstdio>
#include <cstring>
#include <cstdint>
#include <string>
#include <iostream>
#include "Drivers/Net/Wifi/Wpa.hpp"
#include "Drivers/Net/Wifi/Ieee80211.hpp"

namespace Timekeeping { static uint64_t g_ms = 0; uint64_t GetMilliseconds() { return g_ms; } }

using namespace Drivers::Net::Wifi;

static void puthex(const char* tag, const uint8_t* p, size_t n) {
    printf("%s ", tag);
    for (size_t i = 0; i < n; i++) printf("%02x", p[i]);
    printf("\n");
}

// --- hooks the supplicant calls back into ---
namespace Drivers::Net::Wifi {
    bool WpaTxEapol(const uint8_t* body, uint32_t len) {
        puthex("TX", body, len);
        return true;
    }
    bool WpaInstallPtk(const uint8_t* tk, uint32_t tkLen, uint8_t cipher) {
        printf("PTK-CIPHER %u\n", cipher);
        puthex("PTK", tk, tkLen);
        return true;
    }
    bool WpaInstallGtk(const uint8_t* gtk, uint32_t gtkLen, uint8_t keyIdx,
                       uint8_t cipher, const uint8_t* rsc) {
        printf("GTK-IDX %u\nGTK-CIPHER %u\n", keyIdx, cipher);
        puthex("GTK", gtk, gtkLen);
        puthex("GTK-RSC", rsc, 8);
        return true;
    }
}

static int unhex(const std::string& s, uint8_t* out) {
    int n = 0;
    for (size_t i = 0; i + 1 < s.size(); i += 2) {
        unsigned v; sscanf(s.c_str() + i, "%2x", &v); out[n++] = (uint8_t)v;
    }
    return n;
}

int main() {
    std::string line;
    uint8_t buf[2048];

    while (std::getline(std::cin, line)) {
        if (line.rfind("START ", 0) == 0) {
            // START <ownmac> <bssid> <ssid> <passphrase> <akm> <pcipher> <gcipher>
            char own[64], bss[64], ssid[64], pass[128];
            unsigned akm, pc, gc;
            sscanf(line.c_str(), "START %63s %63s %63s %127s %u %u %u",
                   own, bss, ssid, pass, &akm, &pc, &gc);
            WpaConfig cfg = {};
            unhex(own, cfg.OwnMac);
            unhex(bss, cfg.Bssid);
            cfg.SsidLen = (uint8_t)strlen(ssid);
            memcpy(cfg.Ssid, ssid, cfg.SsidLen);
            cfg.PassLen = (uint8_t)strlen(pass);
            memcpy(cfg.Passphrase, pass, cfg.PassLen);
            cfg.Akm = (uint8_t)akm;
            cfg.PairwiseCipher = (uint8_t)pc;
            cfg.GroupCipher = (uint8_t)gc;
            printf("START-OK %d\n", WpaStart(cfg) ? 1 : 0);
            uint8_t ie[64];
            uint32_t ieLen = WpaBuildRsnIe(ie, sizeof(ie));
            puthex("RSNIE", ie, ieLen);
        } else if (line.rfind("RX ", 0) == 0) {
            int n = unhex(line.substr(3), buf);
            bool consumed = WpaOnEapol(buf, (uint32_t)n);
            printf("RX-OK %d STATE %d\n", consumed ? 1 : 0, (int)WpaGetState());
        } else if (line.rfind("PARSE ", 0) == 0) {
            int n = unhex(line.substr(6), buf);
            WpaConfig cfg = {};
            bool ok = WpaParseApRsn(buf, (uint32_t)n, cfg);
            printf("PARSE-OK %d AKM %u PCIPHER %u GCIPHER %u\n", ok ? 1 : 0,
                   (unsigned)cfg.Akm, (unsigned)cfg.PairwiseCipher,
                   (unsigned)cfg.GroupCipher);
        } else if (line.rfind("STATE", 0) == 0) {
            printf("STATE %d\n", (int)WpaGetState());
        } else if (line.rfind("QUIT", 0) == 0) {
            break;
        }
        fflush(stdout);
    }
    return 0;
}