// Simulated board (C++17, OpenSSL 3 + libmosquitto): anonymous enrollment over MQTT, then login as FN.
// Every step is logged. The derived password is never printed, only a short SHA-256 fingerprint.
#include <mosquitto.h>
#include <openssl/core_names.h>
#include <openssl/evp.h>
#include <openssl/kdf.h>
#include <openssl/params.h>
#include <openssl/rand.h>

#include <chrono>
#include <cstdarg>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <map>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

// How to read this file: main() at the bottom runs the ten steps in order. Above it: logging and
// hex helpers, the crypto (OpenSSL EVP wrappers with automatic cleanup), and `Conn`, a synchronous
// wrapper around libmosquitto. g_checks/g_failures drive the final PASS/FAIL and the exit code.
using Bytes = std::vector<uint8_t>;

namespace {

constexpr const char *kTag = "board-cpp";
constexpr const char *kInfoLabel = "robot-market/mqtt-pw/v1|";

int g_checks = 0, g_failures = 0, g_step_delay = 2;
std::FILE *g_logfile = nullptr;

void sleepSec(int s) { std::this_thread::sleep_for(std::chrono::seconds(s)); }

void LOG(const char *fmt, ...) {
    auto now = std::chrono::system_clock::now();
    std::time_t t = std::chrono::system_clock::to_time_t(now);
    int ms = static_cast<int>(std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count() % 1000);
    std::tm tm{};
    localtime_r(&t, &tm);
    char ts[16];
    std::strftime(ts, sizeof ts, "%H:%M:%S", &tm);
    char body[1024];
    va_list ap;
    va_start(ap, fmt);
    std::vsnprintf(body, sizeof body, fmt, ap);
    va_end(ap);
    std::printf("%s.%03d [%s] %s\n", ts, ms, kTag, body);
    std::fflush(stdout);
    if (g_logfile) {  // optional copy of every line for the web viewer (LOG_FILE)
        std::fprintf(g_logfile, "%s.%03d [%s] %s\n", ts, ms, kTag, body);
        std::fflush(g_logfile);
    }
}

// step() prints a header line; the web page turns lines in this exact format into step bands.
void step(int n, const char *title) { LOG("---- step %d: %s ----", n, title); }
void pauseStep() { sleepSec(g_step_delay); }

// check() records one test result and prints PASS or FAIL for it.
void check(bool ok, const char *what) {
    g_checks++;
    if (!ok) g_failures++;
    LOG("%s  %s", ok ? "PASS" : "FAIL", what);
}

[[noreturn]] void finish() {
    LOG("RESULT: %s (%d checks, %d failed)", g_failures ? "FAIL" : "PASS", g_checks, g_failures);
    std::exit(g_failures ? 1 : 0);
}

[[noreturn]] void die(const std::string &m) {
    LOG("FATAL: %s", m.c_str());
    std::exit(2);
}

std::string envStr(const char *k, const char *def) {
    const char *v = std::getenv(k);
    return (v && *v) ? v : def;
}
int envInt(const char *k, int def) {
    const char *v = std::getenv(k);
    return (v && *v) ? std::atoi(v) : def;
}

Bytes unhex(const std::string &h) {
    Bytes out;
    for (size_t i = 0; i + 1 < h.size(); i += 2) out.push_back(static_cast<uint8_t>(std::stoi(h.substr(i, 2), nullptr, 16)));
    return out;
}

std::string tohex(const uint8_t *b, size_t n) {
    std::string s;
    char buf[3];
    for (size_t i = 0; i < n; i++) { std::snprintf(buf, sizeof buf, "%02x", b[i]); s += buf; }
    return s;
}
std::string tohex(const Bytes &b) { return tohex(b.data(), b.size()); }

Bytes randomBytes(size_t n) {
    Bytes b(n);
    if (RAND_bytes(b.data(), static_cast<int>(n)) != 1) die("RAND_bytes");
    return b;
}

// Short SHA-256 of the password: lets the logs prove both sides derived the same value without
// ever printing the password itself.
std::string fingerprint(const std::string &s) {
    uint8_t md[EVP_MAX_MD_SIZE];
    unsigned len = 0;
    if (EVP_Digest(s.data(), s.size(), md, &len, EVP_sha256(), nullptr) != 1) die("sha256");
    return tohex(md, 8);
}

// ---- crypto (OpenSSL EVP) ----
struct PkeyDel { void operator()(EVP_PKEY *p) const { EVP_PKEY_free(p); } };
struct PctxDel { void operator()(EVP_PKEY_CTX *p) const { EVP_PKEY_CTX_free(p); } };
struct MdDel { void operator()(EVP_MD_CTX *p) const { EVP_MD_CTX_free(p); } };
struct KdfDel { void operator()(EVP_KDF *p) const { EVP_KDF_free(p); } };
struct KdfCtxDel { void operator()(EVP_KDF_CTX *p) const { EVP_KDF_CTX_free(p); } };
using Pkey = std::unique_ptr<EVP_PKEY, PkeyDel>;

Bytes x25519Pub(const Bytes &priv) {
    Pkey k(EVP_PKEY_new_raw_private_key(EVP_PKEY_X25519, nullptr, priv.data(), priv.size()));
    Bytes pub(32);
    size_t l = pub.size();
    if (!k || EVP_PKEY_get_raw_public_key(k.get(), pub.data(), &l) != 1 || l != 32) die("x25519 pub");
    return pub;
}

// Returns an empty vector when the derivation fails (e.g. an all-zero shared secret).
// X25519 key agreement: my private key + their public key = the same secret on both sides.
Bytes x25519Shared(const Bytes &priv, const Bytes &peer) {
    Pkey k(EVP_PKEY_new_raw_private_key(EVP_PKEY_X25519, nullptr, priv.data(), priv.size()));
    Pkey p(EVP_PKEY_new_raw_public_key(EVP_PKEY_X25519, nullptr, peer.data(), peer.size()));
    if (!k || !p) return {};
    std::unique_ptr<EVP_PKEY_CTX, PctxDel> c(EVP_PKEY_CTX_new(k.get(), nullptr));
    Bytes out(32);
    size_t l = out.size();
    if (!c || EVP_PKEY_derive_init(c.get()) != 1 || EVP_PKEY_derive_set_peer(c.get(), p.get()) != 1 ||
        EVP_PKEY_derive(c.get(), out.data(), &l) != 1 || l != 32)
        return {};
    return out;
}

// HKDF-SHA256 (RFC 5869) through OpenSSL's KDF API; used to turn the shared secret into the password.
Bytes hkdf(const Bytes &ikm, const Bytes &salt, const Bytes &info, size_t n) {
    std::unique_ptr<EVP_KDF, KdfDel> kdf(EVP_KDF_fetch(nullptr, "HKDF", nullptr));
    if (!kdf) die("hkdf fetch");
    std::unique_ptr<EVP_KDF_CTX, KdfCtxDel> ctx(EVP_KDF_CTX_new(kdf.get()));
    std::string digest = "SHA256";
    OSSL_PARAM p[5];
    p[0] = OSSL_PARAM_construct_utf8_string(OSSL_KDF_PARAM_DIGEST, digest.data(), 0);
    p[1] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_KEY, const_cast<uint8_t *>(ikm.data()), ikm.size());
    p[2] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_SALT, const_cast<uint8_t *>(salt.data()), salt.size());
    p[3] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_INFO, const_cast<uint8_t *>(info.data()), info.size());
    p[4] = OSSL_PARAM_construct_end();
    Bytes out(n);
    if (!ctx || EVP_KDF_derive(ctx.get(), out.data(), out.size(), p) != 1) die("hkdf");
    return out;
}

// Ed25519 signature check against the backend's pinned public key.
bool edVerify(const Bytes &pub, const Bytes &msg, const Bytes &sig) {
    Pkey k(EVP_PKEY_new_raw_public_key(EVP_PKEY_ED25519, nullptr, pub.data(), pub.size()));
    std::unique_ptr<EVP_MD_CTX, MdDel> c(EVP_MD_CTX_new());
    return k && c && sig.size() == 64 && EVP_DigestVerifyInit(c.get(), nullptr, nullptr, nullptr, k.get()) == 1 &&
           EVP_DigestVerify(c.get(), sig.data(), sig.size(), msg.data(), msg.size()) == 1;
}

// The exact byte string the backend signed: S || B || n1 || clientid || FN.
Bytes signedMessage(const Bytes &S, const Bytes &B, const Bytes &n1, const std::string &cid, const std::string &fn) {
    Bytes m;
    for (const Bytes *p : {&S, &B, &n1}) m.insert(m.end(), p->begin(), p->end());
    m.insert(m.end(), cid.begin(), cid.end());
    m.insert(m.end(), fn.begin(), fn.end());
    return m;
}

// ---- MQTT plumbing ----
// Conn: one MQTT connection with blocking helpers. The callbacks only set flags/fields; connect(),
// subscribe() and waitMessage() run the network loop until the flag flips or a timeout passes.
// subscribe() returning 128 means the broker's ACL denied it, which several tests rely on.
class Conn {
public:
    Conn(const std::string &host, int port) : host_(host), port_(port) {}
    ~Conn() { close(); }

    // Returns the CONNACK return code (0 = accepted) or -1 if none arrived.
    int connect(const std::string &clientId, const char *user, const std::string &pw) {
        close();
        connected_ = false;
        rc_ = -1;
        m_ = mosquitto_new(clientId.c_str(), true, this);
        if (!m_) die("mosquitto_new");
        mosquitto_connect_callback_set(m_, [](mosquitto *, void *o, int rc) {
            auto *c = static_cast<Conn *>(o);
            c->rc_ = rc;
            c->connected_ = true;
        });
        mosquitto_subscribe_callback_set(m_, [](mosquitto *, void *o, int, int n, const int *granted) {
            auto *c = static_cast<Conn *>(o);
            c->subQos_ = n > 0 ? granted[0] : -1;
            c->subDone_ = true;
        });
        mosquitto_message_callback_set(m_, [](mosquitto *, void *o, const mosquitto_message *msg) {
            auto *c = static_cast<Conn *>(o);
            c->msgTopic_ = msg->topic;
            c->msg_.assign(static_cast<const char *>(msg->payload), static_cast<size_t>(msg->payloadlen));
            c->gotMsg_ = true;
        });
        if (user) mosquitto_username_pw_set(m_, user, pw.c_str());
        for (int i = 0; i < 30; i++) {
            if (mosquitto_connect(m_, host_.c_str(), port_, 30) == MOSQ_ERR_SUCCESS) break;
            LOG("broker not reachable yet, retrying ...");
            sleepSec(2);
        }
        waitFor([&] { return connected_; }, 10);
        return connected_ ? rc_ : -1;
    }

    void close() {
        if (!m_) return;
        mosquitto_disconnect(m_);
        mosquitto_loop(m_, 100, 1);
        mosquitto_destroy(m_);
        m_ = nullptr;
    }

    // Granted QoS (0..2), 128 when denied, -1 on timeout.
    int subscribe(const std::string &topic) {
        subDone_ = false;
        mosquitto_subscribe(m_, nullptr, topic.c_str(), 1);
        waitFor([&] { return subDone_; }, 10);
        return subDone_ ? subQos_ : -1;
    }

    void publish(const std::string &topic, const std::string &payload) {
        mosquitto_publish(m_, nullptr, topic.c_str(), static_cast<int>(payload.size()), payload.data(), 1, false);
    }

    bool waitMessage(const std::string &topic, int secs) {
        gotMsg_ = false;
        msg_.clear();
        waitFor([&] { return gotMsg_ && msgTopic_ == topic; }, secs);
        return gotMsg_ && msgTopic_ == topic;
    }

    const std::string &message() const { return msg_; }

private:
    template <class F>
    void waitFor(F cond, int secs) {
        auto end = std::chrono::steady_clock::now() + std::chrono::seconds(secs);
        while (!cond() && std::chrono::steady_clock::now() < end) {
            if (mosquitto_loop(m_, 100, 1) != MOSQ_ERR_SUCCESS) std::this_thread::sleep_for(std::chrono::milliseconds(50));
        }
    }

    std::string host_;
    int port_;
    mosquitto *m_ = nullptr;
    volatile bool connected_ = false, subDone_ = false, gotMsg_ = false;
    volatile int rc_ = -1, subQos_ = -1;
    std::string msgTopic_, msg_;
};

// parseFields: read the simple "key=value" lines used by the response message.
std::map<std::string, std::string> parseFields(const std::string &payload) {
    std::map<std::string, std::string> out;
    std::istringstream in(payload);
    for (std::string line; std::getline(in, line);) {
        auto eq = line.find('=');
        if (eq != std::string::npos) out[line.substr(0, eq)] = line.substr(eq + 1);
    }
    return out;
}

}  // namespace

int main() {
    if (const char *lf = std::getenv("LOG_FILE"); lf && *lf) g_logfile = std::fopen(lf, "w");
    g_step_delay = envInt("STEP_DELAY", 2);
    const std::string host = envStr("BROKER_HOST", "broker");
    const int port = envInt("BROKER_PORT", 1883);
    const std::string mac = envStr("BOARD_MAC", "02:00:00:00:00:02");
    const Bytes pinned = unhex(envStr("BACKEND_PUBKEY_HEX", ""));
    if (pinned.size() != 32) die("BACKEND_PUBKEY_HEX must be 32 bytes of hex");
    mosquitto_lib_init();

    const int start = envInt("START_DELAY", 15);
    LOG("board starting; waiting %d s so the logs can be followed (%s)", start, OpenSSL_version(OPENSSL_VERSION));
    for (int left = start; left > 0; left -= 5) {
        LOG("starting in %d s ...", left);
        sleepSec(std::min(5, left));
    }

    // 1. identity material
    step(1, "generate random Client ID, X25519 key pair and nonce (all from the CSPRNG)");
    const Bytes b = randomBytes(32), n1 = randomBytes(16);
    const std::string cid = tohex(randomBytes(16));
    const Bytes B = x25519Pub(b);
    LOG("client id=%s", cid.c_str());
    LOG("board public key B=%s", tohex(B).c_str());
    LOG("nonce n1=%s", tohex(n1).c_str());
    LOG("board private key b stays in memory and is never logged or sent");
    pauseStep();

    // 2. anonymous connect
    step(2, "connect anonymously to the broker using the random Client ID");
    Conn c(host, port);
    int rc = c.connect(cid, nullptr, "");
    check(rc == 0, "anonymous CONNECT accepted by the enrollment listener");
    if (rc != 0) finish();
    pauseStep();

    // The ACL only allows enroll/response/<my own client id>. Asking for someone else's topic or a
    // wildcard must be refused (SUBACK 128), otherwise another board could read my reply.
    // 3. subscribe to own response topic; try another board's
    step(3, "subscribe to own response topic; ACL must deny anyone else's");
    const std::string respTopic = "enroll/response/" + cid, reqTopic = "enroll/request/" + cid;
    int q = -1;
    for (int tries = 0; tries < 10; tries++) {  // backend may still be provisioning
        q = c.subscribe(respTopic);
        if (q >= 0 && q != 128) break;
        LOG("subscribe to %s not granted yet (qos=%d), retrying ...", respTopic.c_str(), q);
        sleepSec(2);
    }
    check(q >= 0 && q != 128, "subscribe to own enroll/response/<clientid> granted");
    q = c.subscribe("enroll/response/00000000000000000000000000000000");
    check(q == 128, "subscribe to another board's enroll/response/<other> is DENIED (SUBACK 0x80)");
    q = c.subscribe("enroll/response/#");
    check(q == 128, "wildcard subscribe enroll/response/# is DENIED (SUBACK 0x80)");
    pauseStep();

    // 4. request
    step(4, "publish the enrollment request (B, n1, mac)");
    const std::string req = "B=" + tohex(B) + "\nn1=" + tohex(n1) + "\nmac=" + mac + "\n";
    c.publish(reqTopic, req);
    LOG("published request on %s", reqTopic.c_str());

    // 5. wait for the response (re-send the identical request if it gets lost)
    step(5, "wait for the backend's response");
    bool got = false;
    for (int attempt = 0; attempt < 6 && !got; attempt++) {
        got = c.waitMessage(respTopic, 15);
        if (!got) {
            LOG("no response yet, re-sending the same request (attempt %d)", attempt + 2);
            c.publish(reqTopic, req);
        }
    }
    check(got, "response received on enroll/response/<clientid>");
    if (!got) finish();
    auto f = parseFields(c.message());
    const bool parsed = f.count("fn") && f.count("S") && f.count("sig");
    check(parsed, "response contains fn, S and sig");
    if (!parsed) finish();
    const std::string fn = f["fn"];
    LOG("response: fn=%s S=%s", fn.c_str(), f["S"].c_str());
    LOG("response: sig=%s", f["sig"].c_str());
    LOG("(the response carries no password and no secret of any kind)");
    pauseStep();

    // The backend's public key is built into the board (BACKEND_PUBKEY_HEX), like firmware would hold it.
    // The extra checks prove a forged, replayed or re-labelled reply would be rejected.
    // 6. verify signature, incl. tamper checks
    step(6, "verify the backend's Ed25519 signature with the pinned public key");
    const Bytes S = unhex(f["S"]);
    Bytes sig = unhex(f["sig"]);
    check(S.size() == 32 && sig.size() == 64, "S and signature decode to 32/64 bytes");
    Bytes msg = signedMessage(S, B, n1, cid, fn);
    check(edVerify(pinned, msg, sig), "signature over S||B||n1||clientid||FN is VALID");
    Bytes bad = sig;
    bad[10] ^= 0x01;
    check(!edVerify(pinned, msg, bad), "tampered signature (1 bit flipped) is REJECTED");
    Bytes otherNonce = n1;
    otherNonce[0] ^= 0xff;
    check(!edVerify(pinned, signedMessage(S, B, otherNonce, cid, fn), sig), "reply bound to a different nonce is REJECTED (replay check)");
    check(!edVerify(pinned, signedMessage(S, B, n1, cid, "FN-999999"), sig), "reply with a swapped FN is REJECTED");
    pauseStep();

    // Both sides now hold the other's public key: X25519(my private key, their public key) gives the
    // same secret on each side. Nothing secret was transmitted.
    // 7. key agreement
    step(7, "derive the shared secret (X25519) and the broker password (HKDF-SHA256)");
    const Bytes shared = x25519Shared(b, S);
    check(shared.size() == 32, "X25519(b, S) succeeded (non all-zero)");
    if (shared.size() != 32) finish();
    const std::string info = std::string(kInfoLabel) + fn;
    const std::string pw = tohex(hkdf(shared, n1, Bytes(info.begin(), info.end()), 32));
    LOG("derived password for %s locally; fingerprint=%s (compare with the backend log; password never logged)", fn.c_str(),
        fingerprint(pw).c_str());
    pauseStep();

    // The negative logins come first so a refused attempt cannot disturb the real session: the
    // broker checks credentials before it would take over a client id.
    // 8. leave enrollment, log in as FN
    step(8, "disconnect from enrollment and log in as the FN with the derived password");
    rc = c.connect(fn, fn.c_str(), "definitely-not-the-password");
    check(rc != 0, "login with a WRONG password is REFUSED");
    rc = c.connect("some-other-client-id", fn.c_str(), pw);
    check(rc != 0, "login with the right password but a different Client ID is REFUSED (Client ID binding)");
    rc = c.connect(fn, fn.c_str(), pw);
    check(rc == 0, "login as FN with the derived password and bound Client ID is ACCEPTED");
    if (rc != 0) finish();
    pauseStep();

    // 9. FN-scoped access
    step(9, "use the FN topics: subscribe board/<FN>/cmd, publish board/<FN>/hello, wait for ack");
    const std::string cmdTopic = "board/" + fn + "/cmd", helloTopic = "board/" + fn + "/hello";
    q = c.subscribe(cmdTopic);
    check(q >= 0 && q != 128, "subscribe to own board/<FN>/cmd granted");
    q = c.subscribe("board/FN-999999/cmd");
    check(q == 128, "subscribe to another FN's board/<other>/cmd is DENIED");
    q = c.subscribe("enroll/response/#");
    check(q == 128, "FN board can no longer use the enrollment topics (DENIED)");
    c.publish(helloTopic, "hello from board-cpp");
    LOG("published hello on %s", helloTopic.c_str());
    got = c.waitMessage(cmdTopic, 15);
    check(got && c.message() == "ack-hello", "backend's ack-hello received on board/<FN>/cmd");
    c.close();

    step(10, "done");
    mosquitto_lib_cleanup();
    finish();
}