/*
 * handshake.c — BLE Challenge-Response Authentication implementation
 *
 * Key: SHA256("29C310F5")[0..15] XOR 0x5A
 * Challenge: R = 4-byte timestamp (big-endian) + 12 random bytes
 * Verify: base64(SHA256(R || Key)) == received_response
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <time.h>
#include <fcntl.h>
#include <unistd.h>

#include "sha256.h"
#include "base64.h"
#include "handshake.h"

/* Timestamp when challenge was sent (for timeout check) */
static time_t s_challenge_send_time = 0;

/* ── Key derivation ─────────────────────────────────────────────────────── */

void handshake_derive_key(uint8_t out_key[HANDSHAKE_KEY_LEN])
{
    uint8_t digest[32];
    sha256((const uint8_t *)HANDSHAKE_KN_STR, strlen(HANDSHAKE_KN_STR), digest);
    for (int i = 0; i < HANDSHAKE_KEY_LEN; i++)
        out_key[i] = digest[i] ^ 0x5A;
}

void handshake_print_key(const uint8_t key[HANDSHAKE_KEY_LEN])
{
    for (int i = 0; i < HANDSHAKE_KEY_LEN; i++)
        printf("%02x%s", key[i], (i < HANDSHAKE_KEY_LEN - 1) ? " " : "\n");
}

/* ── Init / Reset ────────────────────────────────────────────────────────── */

void handshake_init(handshake_ctx_t *ctx)
{
    memset(ctx, 0, sizeof(*ctx));
    ctx->state = HS_STATE_IDLE;
    handshake_derive_key(ctx->key);
}

void handshake_reset(handshake_ctx_t *ctx)
{
    ctx->state = HS_STATE_IDLE;
    memset(ctx->R, 0, sizeof(ctx->R));
    s_challenge_send_time = 0;
}

handshake_state_t handshake_get_state(const handshake_ctx_t *ctx)
{
    return ctx->state;
}

int handshake_check_timeout(handshake_ctx_t *ctx)
{
    if (ctx->state != HS_STATE_WAIT_RESPONSE)
        return 0;
    return (time(NULL) - s_challenge_send_time >= HANDSHAKE_TIMEOUT_SEC) ? 1 : 0;
}

/* ── Challenge generation ────────────────────────────────────────────────── */

int handshake_build_challenge(handshake_ctx_t *ctx,
                               char *out_json_buf, size_t buf_size)
{
#if 0  /* ── TEST VECTOR (Protocol 5.3): change to #if 0 to restore random R ── */
    /* R = 6a4636a0a1b2c3d4e5f60718293a4b5c
     * Expected challenge: akY2oKGyw9Tl9gcYKTpLXA==
     * Expected response:  NISSISwhM59eitBMS/stTnJMSLIU7YWZG7bd+NkQBAU= */
    static const uint8_t TEST_R[HANDSHAKE_R_LEN] = {
        0x6a, 0x46, 0x36, 0xa0, 0xa1, 0xb2, 0xc3, 0xd4,
        0xe5, 0xf6, 0x07, 0x18, 0x29, 0x3a, 0x4b, 0x5c
    };
    memcpy(ctx->R, TEST_R, HANDSHAKE_R_LEN);
    time_t now = time(NULL);   /* still needed for s_challenge_send_time */
#else  /* Normal operation: random R */
    /* R[0..3] = current time big-endian */
    time_t now = time(NULL);
    ctx->R[0] = (uint8_t)((now >> 24) & 0xFF);
    ctx->R[1] = (uint8_t)((now >> 16) & 0xFF);
    ctx->R[2] = (uint8_t)((now >>  8) & 0xFF);
    ctx->R[3] = (uint8_t)( now        & 0xFF);

    /* R[4..15] = 12 random bytes from /dev/urandom */
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd >= 0) {
        if (read(fd, ctx->R + 4, 12) != 12) {
            /* partial read fallback */
            for (int i = 4; i < 16; i++) ctx->R[i] ^= (uint8_t)(rand() ^ i);
        }
        close(fd);
    } else {
        /* /dev/urandom unavailable — simple fallback */
        srand((unsigned int)now);
        for (int i = 4; i < 16; i++)
            ctx->R[i] = (uint8_t)(rand() ^ i ^ (int)now);
    }
#endif

    /* Base64-encode R (16 bytes → 24 chars + NUL) */
    char r_b64[28];   /* ceil(16/3)*4+1 = 25 bytes, round up */
    base64_encode(ctx->R, HANDSHAKE_R_LEN, r_b64);

    /* Build JSON */
    int n = snprintf(out_json_buf, buf_size,
        "{\"response_type\":\"handshake_challenge\","
        "\"content\":{\"challenge\":\"%s\"}}",
        r_b64);
    if (n < 0 || (size_t)n >= buf_size)
        return -1;

    /* Record send time and advance state */
    s_challenge_send_time = now;
    ctx->state = HS_STATE_WAIT_RESPONSE;
    return 0;
}

/* Response validation*/

int handshake_process_response(handshake_ctx_t *ctx, const char *json_str)
{
    /* 1. Timeout check */
    if (handshake_check_timeout(ctx)) {
        ctx->state = HS_STATE_FAILED;
        return -2;
    }

    /* 2. Parse "response" field value from JSON
     *    Looks for: "response":"<base64>" */
    const char *p = strstr(json_str, "\"response\"");
    if (!p) return -1;
    p = strchr(p, ':');
    if (!p) return -1;
    /* skip whitespace and colon */
    while (*p == ':' || *p == ' ' || *p == '\t') p++;
    if (*p != '"') return -1;
    p++;  /* skip opening quote */
    const char *end = strchr(p, '"');
    if (!end) return -1;
    size_t val_len = (size_t)(end - p);
    if (val_len == 0 || val_len > 64) return -1;

    char received_b64[68];
    memcpy(received_b64, p, val_len);
    received_b64[val_len] = '\0';

    /* 3. Compute expected = base64(SHA256(R || Key)) */
    uint8_t input[HANDSHAKE_R_LEN + HANDSHAKE_KEY_LEN];  /* 32 bytes */
    memcpy(input,                   ctx->R,   HANDSHAKE_R_LEN);
    memcpy(input + HANDSHAKE_R_LEN, ctx->key, HANDSHAKE_KEY_LEN);

    uint8_t digest[32];
    sha256(input, HANDSHAKE_R_LEN + HANDSHAKE_KEY_LEN, digest);

    char expected_b64[48];  /* 32 bytes → 44 chars + NUL */
    base64_encode(digest, 32, expected_b64);

    printf("[HS DEBUG] R        : ");
    for (int i = 0; i < HANDSHAKE_R_LEN; i++) printf("%02x", ctx->R[i]);
    printf("\n");

    printf("[HS DEBUG] Key      : ");
    for (int i = 0; i < HANDSHAKE_KEY_LEN; i++) printf("%02x", ctx->key[i]);
    printf("\n");

    printf("[HS DEBUG] expected : %s\n", expected_b64);
    printf("[HS DEBUG] received : %s\n", received_b64);

    /* 4. Constant-time compare (simple for embedded) */
    if (strcmp(received_b64, expected_b64) == 0) {
        ctx->state = HS_STATE_SUCCESS;
        return 1;
    }

    ctx->state = HS_STATE_FAILED;
    return 0;
}

/* ── Confirm message ─────────────────────────────────────────────────────── */

int handshake_build_confirm(handshake_ctx_t *ctx, int success,
                             const char *reason_str,
                             char *out_json_buf, size_t buf_size)
{
    (void)ctx;
    if (success) {
        snprintf(out_json_buf, buf_size,
            "{\"response_type\":\"handshake_confirm\","
            "\"content\":{\"status\":\"success\"}}");
    } else {
        const char *reason = reason_str ? reason_str : "invalid_response";
        snprintf(out_json_buf, buf_size,
            "{\"response_type\":\"handshake_confirm\","
            "\"content\":{\"status\":\"failed\",\"reason\":\"%s\"}}",
            reason);
    }
    return 0;
}