/* Simulated board (C, libsodium + 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 <sodium.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

/* How to read this file: main() at the bottom runs the ten steps in order. The functions above it are
 * small helpers: logging, hex conversion, a synchronous wrapper around libmosquitto, and the crypto
 * (all libsodium). Counters g_checks/g_failures drive the final PASS/FAIL and the process exit code. */
#define TAG "board-c"
#define INFO_LABEL "robot-market/mqtt-pw/v1|"

static int g_checks, g_failures, g_step_delay = 2;

static FILE *g_logfile; /* optional copy of every line for the web viewer (LOG_FILE) */

static void logf_(const char *fmt, ...)
{
    struct timespec ts;
    struct tm tm;
    char t[16], body[1024];
    clock_gettime(CLOCK_REALTIME, &ts);
    localtime_r(&ts.tv_sec, &tm);
    strftime(t, sizeof t, "%H:%M:%S", &tm);
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(body, sizeof body, fmt, ap);
    va_end(ap);
    printf("%s.%03ld [%s] %s\n", t, ts.tv_nsec / 1000000, TAG, body);
    fflush(stdout);
    if (g_logfile) {
        fprintf(g_logfile, "%s.%03ld [%s] %s\n", t, ts.tv_nsec / 1000000, TAG, body);
        fflush(g_logfile);
    }
}
#define LOG(...) logf_(__VA_ARGS__)

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

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

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

static void die(const char *msg)
{
    LOG("FATAL: %s", msg);
    exit(2);
}

static const char *env_str(const char *k, const char *def)
{
    const char *v = getenv(k);
    return (v && *v) ? v : def;
}

static int env_int(const char *k, int def)
{
    const char *v = getenv(k);
    return (v && *v) ? atoi(v) : def;
}

static size_t unhex(const char *h, uint8_t *out, size_t max)
{
    size_t n = strlen(h) / 2;
    if (n > max)
        return 0;
    for (size_t i = 0; i < n; i++) {
        unsigned v;
        if (sscanf(h + 2 * i, "%2x", &v) != 1)
            return 0;
        out[i] = (uint8_t)v;
    }
    return n;
}

static void tohex(const uint8_t *b, size_t n, char *out)
{
    for (size_t i = 0; i < n; i++)
        sprintf(out + 2 * i, "%02x", b[i]);
    out[2 * n] = 0;
}

/* Short SHA-256 of the password: lets the logs prove both sides derived the same value without
 * ever printing the password itself. */
static void fingerprint(const char *s, char out[17])
{
    uint8_t h[crypto_hash_sha256_BYTES];
    crypto_hash_sha256(h, (const unsigned char *)s, strlen(s));
    tohex(h, 8, out);
}

/* ---- MQTT plumbing: synchronous helpers on top of libmosquitto callbacks ---- */

/* One MQTT connection plus the flags its callbacks set. The helpers below wait on these flags, which
 * turns libmosquitto's callback style into simple "do it, then wait for the result" calls. */
typedef struct {
    struct mosquitto *m;
    volatile int connected, conn_rc, sub_done, sub_qos, got_msg;
    char msg_topic[256], msg[2048];
} ctx_t;

static void on_connect(struct mosquitto *m, void *o, int rc)
{
    (void)m;
    ctx_t *c = o;
    c->conn_rc = rc;
    c->connected = 1;
}

static void on_subscribe(struct mosquitto *m, void *o, int mid, int n, const int *granted)
{
    (void)m, (void)mid;
    ctx_t *c = o;
    c->sub_qos = n > 0 ? granted[0] : -1;
    c->sub_done = 1;
}

static void on_message(struct mosquitto *m, void *o, const struct mosquitto_message *msg)
{
    (void)m;
    ctx_t *c = o;
    snprintf(c->msg_topic, sizeof c->msg_topic, "%s", msg->topic);
    size_t n = (size_t)msg->payloadlen < sizeof c->msg - 1 ? (size_t)msg->payloadlen : sizeof c->msg - 1;
    memcpy(c->msg, msg->payload, n);
    c->msg[n] = 0;
    c->got_msg = 1;
}

/* WAIT runs the network loop until `cond` becomes true or `secs` pass. */
#define WAIT(c, cond, secs)                                          \
    do {                                                             \
        time_t t0_ = time(NULL);                                     \
        while (!(cond) && time(NULL) - t0_ < (secs)) {               \
            if (mosquitto_loop((c)->m, 100, 1) != MOSQ_ERR_SUCCESS)  \
                usleep(50000);                                       \
        }                                                            \
    } while (0)

/* Connects (retrying while the broker is not reachable) and waits for CONNACK. Returns the
 * CONNACK return code (0 = accepted), or -1 if no CONNACK arrived. */
/* mq_connect: keep trying until the broker is reachable, then wait for the CONNACK. Returns its
 * return code, so a refused login (wrong password, wrong client id) shows up as a non-zero value. */
static int mq_connect(ctx_t *c, const char *host, int port, const char *clientid, const char *user,
                      const char *pw)
{
    memset(c, 0, sizeof *c);
    c->m = mosquitto_new(clientid, true, c);
    if (!c->m)
        die("mosquitto_new");
    mosquitto_connect_callback_set(c->m, on_connect);
    mosquitto_subscribe_callback_set(c->m, on_subscribe);
    mosquitto_message_callback_set(c->m, on_message);
    if (user)
        mosquitto_username_pw_set(c->m, user, pw);
    for (int i = 0; i < 30; i++) {
        if (mosquitto_connect(c->m, host, port, 30) == MOSQ_ERR_SUCCESS)
            break;
        LOG("broker not reachable yet, retrying ...");
        sleep(2);
    }
    WAIT(c, c->connected, 10);
    return c->connected ? c->conn_rc : -1;
}

static void mq_close(ctx_t *c)
{
    mosquitto_disconnect(c->m);
    mosquitto_loop(c->m, 100, 1);
    mosquitto_destroy(c->m);
    c->m = NULL;
}

/* Returns the granted QoS (0..2), 128 when the broker denied the subscription, -1 on timeout. */
/* mq_subscribe: 0..2 = granted, 128 = the broker's ACL denied it. The tests rely on the 128 case. */
static int mq_subscribe(ctx_t *c, const char *topic)
{
    c->sub_done = 0;
    mosquitto_subscribe(c->m, NULL, topic, 1);
    WAIT(c, c->sub_done, 10);
    return c->sub_done ? c->sub_qos : -1;
}

static int mq_wait_message(ctx_t *c, const char *topic, int secs)
{
    c->got_msg = 0;
    c->msg[0] = 0;
    WAIT(c, c->got_msg && strcmp(c->msg_topic, topic) == 0, secs);
    return c->got_msg && strcmp(c->msg_topic, topic) == 0;
}

/* field: read the value of "key=value" from a multi-line message body. */
static int field(const char *payload, const char *key, char *out, size_t outsz)
{
    size_t kl = strlen(key);
    for (const char *p = payload; *p;) {
        const char *eol = strchr(p, '\n');
        size_t len = eol ? (size_t)(eol - p) : strlen(p);
        if (len > kl && strncmp(p, key, kl) == 0 && p[kl] == '=') {
            size_t vl = len - kl - 1;
            if (vl >= outsz)
                return 0;
            memcpy(out, p + kl + 1, vl);
            out[vl] = 0;
            return 1;
        }
        p += len + (eol ? 1 : 0);
    }
    return 0;
}

/* ---- crypto ---- */

/* derive_password: HKDF-SHA256(shared secret, salt = nonce, info = label + FN), hex-encoded. This must
 * match the backend byte for byte, otherwise the login in step 8 fails. */
static void derive_password(const uint8_t *shared, const uint8_t *nonce, const char *fn, char out_hex[65])
{
    char info[128];
    uint8_t prk[crypto_kdf_hkdf_sha256_KEYBYTES], pw[32];
    int il = snprintf(info, sizeof info, INFO_LABEL "%s", fn);
    if (crypto_kdf_hkdf_sha256_extract(prk, nonce, 16, shared, 32) != 0 ||
        crypto_kdf_hkdf_sha256_expand(pw, sizeof pw, info, (size_t)il, prk) != 0)
        die("hkdf");
    tohex(pw, sizeof pw, out_hex);
}

/* build_signed_message: the exact byte string the backend signed: S || B || n1 || clientid || FN. */
static size_t build_signed_message(uint8_t *out, const uint8_t *S, const uint8_t *B, const uint8_t *n1,
                                   const char *cid, const char *fn)
{
    size_t l = 0;
    memcpy(out + l, S, 32), l += 32;
    memcpy(out + l, B, 32), l += 32;
    memcpy(out + l, n1, 16), l += 16;
    memcpy(out + l, cid, strlen(cid)), l += strlen(cid);
    memcpy(out + l, fn, strlen(fn)), l += strlen(fn);
    return l;
}

int main(void)
{
    setvbuf(stdout, NULL, _IOLBF, 0);
    if (getenv("LOG_FILE") && *getenv("LOG_FILE"))
        g_logfile = fopen(getenv("LOG_FILE"), "w");
    g_step_delay = env_int("STEP_DELAY", 2);
    const char *host = env_str("BROKER_HOST", "broker");
    int port = env_int("BROKER_PORT", 1883);
    const char *mac = env_str("BOARD_MAC", "02:00:00:00:00:01");
    uint8_t pinned[32];
    if (unhex(env_str("BACKEND_PUBKEY_HEX", ""), pinned, 32) != 32)
        die("BACKEND_PUBKEY_HEX must be 32 bytes of hex");
    if (sodium_init() < 0)
        die("sodium_init");
    mosquitto_lib_init();

    int start = env_int("START_DELAY", 15);
    LOG("board starting; waiting %d s so the logs can be followed (libsodium %s)", start, sodium_version_string());
    for (int left = start; left > 0; left -= 5) {
        LOG("starting in %d s ...", left);
        sleep((unsigned)(left < 5 ? left : 5));
    }

    /* 1. identity material */
    step(1, "generate random Client ID, X25519 key pair and nonce (all from the CSPRNG)");
    uint8_t cid_raw[16], b[32], B[32], n1[16];
    char cid[33], Bhex[65], nhex[33];
    randombytes_buf(cid_raw, sizeof cid_raw);
    randombytes_buf(b, sizeof b);
    randombytes_buf(n1, sizeof n1);
    if (crypto_scalarmult_base(B, b) != 0)
        die("scalarmult_base");
    tohex(cid_raw, 16, cid);
    tohex(B, 32, Bhex);
    tohex(n1, 16, nhex);
    LOG("client id=%s", cid);
    LOG("board public key B=%s", Bhex);
    LOG("nonce n1=%s", nhex);
    LOG("board private key b stays in memory and is never logged or sent");
    pause_step();

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

    /* 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");
    char resp_topic[96], req_topic[96], other_topic[96];
    snprintf(resp_topic, sizeof resp_topic, "enroll/response/%s", cid);
    snprintf(req_topic, sizeof req_topic, "enroll/request/%s", cid);
    snprintf(other_topic, sizeof other_topic, "enroll/response/%s", "00000000000000000000000000000000");
    int q = -1;
    for (int tries = 0; tries < 10; tries++) { /* backend may still be provisioning */
        q = mq_subscribe(&c, resp_topic);
        if (q >= 0 && q != 128)
            break;
        LOG("subscribe to %s not granted yet (qos=%d), retrying ...", resp_topic, q);
        sleep(2);
    }
    check(q >= 0 && q != 128, "subscribe to own enroll/response/<clientid> granted");
    q = mq_subscribe(&c, other_topic);
    check(q == 128, "subscribe to another board's enroll/response/<other> is DENIED (SUBACK 0x80)");
    q = mq_subscribe(&c, "enroll/response/#");
    check(q == 128, "wildcard subscribe enroll/response/# is DENIED (SUBACK 0x80)");
    pause_step();

    /* 4. request */
    step(4, "publish the enrollment request (B, n1, mac)");
    char req[256];
    snprintf(req, sizeof req, "B=%s\nn1=%s\nmac=%s\n", Bhex, nhex, mac);
    mosquitto_publish(c.m, NULL, req_topic, (int)strlen(req), req, 1, false);
    LOG("published request on %s", req_topic);

    /* 5. wait for the response (re-send the identical request if it gets lost) */
    step(5, "wait for the backend's response");
    int got = 0;
    for (int attempt = 0; attempt < 6 && !got; attempt++) {
        got = mq_wait_message(&c, resp_topic, 15);
        if (!got) {
            LOG("no response yet, re-sending the same request (attempt %d)", attempt + 2);
            mosquitto_publish(c.m, NULL, req_topic, (int)strlen(req), req, 1, false);
        }
    }
    check(got, "response received on enroll/response/<clientid>");
    if (!got)
        finish();
    char fn[32], Shex[65], sighex[129];
    int parsed = field(c.msg, "fn", fn, sizeof fn) && field(c.msg, "S", Shex, sizeof Shex) &&
                 field(c.msg, "sig", sighex, sizeof sighex);
    check(parsed, "response contains fn, S and sig");
    if (!parsed)
        finish();
    LOG("response: fn=%s S=%s", fn, Shex);
    LOG("response: sig=%s", sighex);
    LOG("(the response carries no password and no secret of any kind)");
    pause_step();

    /* 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");
    uint8_t S[32], sig[64], msg[256];
    check(unhex(Shex, S, 32) == 32 && unhex(sighex, sig, 64) == 64, "S and signature decode to 32/64 bytes");
    size_t ml = build_signed_message(msg, S, B, n1, cid, fn);
    check(crypto_sign_verify_detached(sig, msg, ml, pinned) == 0, "signature over S||B||n1||clientid||FN is VALID");
    uint8_t bad[64];
    memcpy(bad, sig, 64);
    bad[10] ^= 0x01;
    check(crypto_sign_verify_detached(bad, msg, ml, pinned) != 0, "tampered signature (1 bit flipped) is REJECTED");
    uint8_t other_nonce[16];
    memcpy(other_nonce, n1, 16);
    other_nonce[0] ^= 0xff;
    ml = build_signed_message(msg, S, B, other_nonce, cid, fn);
    check(crypto_sign_verify_detached(sig, msg, ml, pinned) != 0, "reply bound to a different nonce is REJECTED (replay check)");
    ml = build_signed_message(msg, S, B, n1, cid, "FN-999999");
    check(crypto_sign_verify_detached(sig, msg, ml, pinned) != 0, "reply with a swapped FN is REJECTED");
    pause_step();

    /* 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)");
    uint8_t shared[32];
    check(crypto_scalarmult(shared, b, S) == 0, "X25519(b, S) succeeded (non all-zero)");
    char pw[65], fp[17];
    derive_password(shared, n1, fn, pw);
    fingerprint(pw, fp);
    LOG("derived password for %s locally; fingerprint=%s (compare with the backend log; password never logged)", fn, fp);
    pause_step();

    /* 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");
    mq_close(&c);
    rc = mq_connect(&c, host, port, fn, fn, "definitely-not-the-password");
    check(rc != 0, "login with a WRONG password is REFUSED");
    mq_close(&c);
    rc = mq_connect(&c, host, port, "some-other-client-id", fn, pw);
    check(rc != 0, "login with the right password but a different Client ID is REFUSED (Client ID binding)");
    mq_close(&c);
    rc = mq_connect(&c, host, port, fn, fn, pw);
    check(rc == 0, "login as FN with the derived password and bound Client ID is ACCEPTED");
    if (rc != 0)
        finish();
    pause_step();

    /* 9. FN-scoped access */
    step(9, "use the FN topics: subscribe board/<FN>/cmd, publish board/<FN>/hello, wait for ack");
    char cmd_topic[96], hello_topic[96], other_cmd[96];
    snprintf(cmd_topic, sizeof cmd_topic, "board/%s/cmd", fn);
    snprintf(hello_topic, sizeof hello_topic, "board/%s/hello", fn);
    snprintf(other_cmd, sizeof other_cmd, "board/%s/cmd", "FN-999999");
    q = mq_subscribe(&c, cmd_topic);
    check(q >= 0 && q != 128, "subscribe to own board/<FN>/cmd granted");
    q = mq_subscribe(&c, other_cmd);
    check(q == 128, "subscribe to another FN's board/<other>/cmd is DENIED");
    q = mq_subscribe(&c, "enroll/response/#");
    check(q == 128, "FN board can no longer use the enrollment topics (DENIED)");
    mosquitto_publish(c.m, NULL, hello_topic, (int)strlen("hello from board-c"), "hello from board-c", 1, false);
    LOG("published hello on %s", hello_topic);
    got = mq_wait_message(&c, cmd_topic, 15);
    check(got && strcmp(c.msg, "ack-hello") == 0, "backend's ack-hello received on board/<FN>/cmd");
    mq_close(&c);

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