/**
 * @file    crypto.c
 * @brief   AES-128-CBC (HW CRYP) + HMAC-SHA256 implementation
 *
 * Hardware AES on STM32G474 uses CRYP peripheral.
 * HMAC-SHA256 uses a minimal standalone implementation (no mbedTLS dependency).
 */

#include "crypto.h"
#include "stm32g4xx.h"  /* CMSIS device header */

/* ==================== CRYP Peripheral Helpers ==================== */

/* CRYP configuration key size */
#define CRYP_KEYSIZE_128B     0x2

/* CRYP algorithm: AES-CBC */
#define CRYP_ALGO_AES         0x0
#define CRYP_ALGOMODE_CBC     0x2
#define CRYP_ALGODIR_DECRYPT  0x3
#define CRYP_ALGODIR_ENCRYPT  0x0

static uint8_t g_aes_key[16] __attribute__((aligned(4)));
static bool g_crypto_inited = false;

void CRYPTO_AES_Init(const uint8_t *key)
{
    /* Save key (aligned copy) */
    for (int i = 0; i < 16; i++) {
        g_aes_key[i] = key[i];
    }

    /* Enable CRYP clock */
    RCC->AHB2ENR |= RCC_AHB2ENR_CRYPEN;
    __DSB();

    /* Reset CRYP */
    CRYP->CR |= CRYP_CR_CRYPEN;
    while (CRYP->SR & CRYP_SR_BUSY);

    /* Configure for AES-128-CBC decryption */
    CRYP->CR = (CRYP_KEYSIZE_128B << CRYP_CR_KEYSIZE_Pos) |
               (CRYP_ALGO_AES     << CRYP_CR_ALGO_Pos)     |
               (CRYP_ALGOMODE_CBC << CRYP_CR_ALGOMODE_Pos) |
               (CRYP_ALGODIR_DECRYPT << CRYP_CR_ALGODIR_Pos);

    /* Wait for CRYP to be ready */
    while (CRYP->SR & CRYP_SR_BUSY);

    g_crypto_inited = true;
}

void CRYPTO_AES_Deinit(void)
{
    if (g_crypto_inited) {
        CRYP->CR = 0;
        RCC->AHB2ENR &= ~RCC_AHB2ENR_CRYPEN;
        g_crypto_inited = false;
    }
}

bool CRYPTO_AES_CBC_Decrypt(uint8_t *data, uint32_t len)
{
    uint32_t blocks;
    uint32_t i;
    uint32_t in_word;
    uint32_t *p32;

    if (!g_crypto_inited) return false;
    if (len < 32) return false;         /* Need at least IV + 1 block */
    if (len & 0x0F) return false;       /* Must be 16-byte aligned */

    blocks = len / 16;
    p32 = (uint32_t *)data;

    /* Wait for CRYP ready */
    while (CRYP->SR & CRYP_SR_BUSY);

    /* Write key (4 x 32-bit words) */
    CRYP->K[0] = __REV(*(uint32_t *)&g_aes_key[0]);
    CRYP->K[1] = __REV(*(uint32_t *)&g_aes_key[4]);
    CRYP->K[2] = __REV(*(uint32_t *)&g_aes_key[8]);
    CRYP->K[3] = __REV(*(uint32_t *)&g_aes_key[12]);

    /* Write IV (from first 16 bytes of data) */
    CRYP->IV[0] = __REV(p32[0]);
    CRYP->IV[1] = __REV(p32[1]);
    CRYP->IV[2] = __REV(p32[2]);
    CRYP->IV[3] = __REV(p32[3]);

    /* Process blocks */
    for (i = 0; i < (blocks - 1); i++) {
        /* Wait for output buffer empty */
        while (!(CRYP->SR & CRYP_SR_IFEM));
        /* Write input (p32 starts at IV, so first ciphertext is at offset 4) */
        CRYP->DIN = __REV(p32[(i * 4) + 4]);
        CRYP->DIN = __REV(p32[(i * 4) + 5]);
        CRYP->DIN = __REV(p32[(i * 4) + 6]);
        CRYP->DIN = __REV(p32[(i * 4) + 7]);

        /* Wait for output ready */
        while (!(CRYP->SR & CRYP_SR_OFNE));
        /* Read output back to same buffer (replacing ciphertext) */
        p32[(i * 4) + 0] = __REV(CRYP->DOUT);
        p32[(i * 4) + 1] = __REV(CRYP->DOUT);
        p32[(i * 4) + 2] = __REV(CRYP->DOUT);
        p32[(i * 4) + 3] = __REV(CRYP->DOUT);
    }

    /* Flush CRYP */
    while (CRYP->SR & CRYP_SR_BUSY);
    CRYP->CR &= ~CRYP_CR_CRYPEN;

    return true;
}


/* ==================== HMAC-SHA256 (Standalone) ==================== */

/* SHA-256 context */
typedef struct {
    uint32_t state[8];
    uint32_t count;
    uint8_t  buffer[64];
} sha256_ctx_t;

/* SHA-256 round constants */
static const uint32_t K256[64] = {
    0x428A2F98, 0x71374491, 0xB5C0FBCF, 0xE9B5DBA5,
    0x3956C25B, 0x59F111F1, 0x923F82A4, 0xAB1C5ED5,
    0xD807AA98, 0x12835B01, 0x243185BE, 0x550C7DC3,
    0x72BE5D74, 0x80DEB1FE, 0x9BDC06A7, 0xC19BF174,
    0xE49B69C1, 0xEFBE4786, 0x0FC19DC6, 0x240CA1CC,
    0x2DE92C6F, 0x4A7484AA, 0x5CB0A9DC, 0x76F988DA,
    0x983E5152, 0xA831C66D, 0xB00327C8, 0xBF597FC7,
    0xC6E00BF3, 0xD5A79147, 0x06CA6351, 0x14292967,
    0x27B70A85, 0x2E1B2138, 0x4D2C6DFC, 0x53380D13,
    0x650A7354, 0x766A0ABB, 0x81C2C92E, 0x92722C85,
    0xA2BFE8A1, 0xA81A664B, 0xC24B8B70, 0xC76C51A3,
    0xD192E819, 0xD6990624, 0xF40E3585, 0x106AA070,
    0x19A4C116, 0x1E376C08, 0x2748774C, 0x34B0BCB5,
    0x391C0CB3, 0x4ED8AA4A, 0x5B9CCA4F, 0x682E6FF3,
    0x748F82EE, 0x78A5636F, 0x84C87814, 0x8CC70208,
    0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2
};

#define ROTR(a, b) (((a) >> (b)) | ((a) << (32 - (b))))
#define CH(x, y, z) (((x) & (y)) ^ (~(x) & (z)))
#define MAJ(x, y, z) (((x) & (y)) ^ ((x) & (z)) ^ ((y) & (z)))
#define EP0(x) (ROTR(x, 2) ^ ROTR(x, 13) ^ ROTR(x, 22))
#define EP1(x) (ROTR(x, 6) ^ ROTR(x, 11) ^ ROTR(x, 25))
#define SIG0(x) (ROTR(x, 7) ^ ROTR(x, 18) ^ ((x) >> 3))
#define SIG1(x) (ROTR(x, 17) ^ ROTR(x, 19) ^ ((x) >> 10))

static void sha256_transform(sha256_ctx_t *ctx)
{
    uint32_t W[64], a, b, c, d, e, f, g, h, t1, t2;
    int i;

    for (i = 0; i < 16; i++) {
        W[i] = ((uint32_t)ctx->buffer[i * 4] << 24) |
               ((uint32_t)ctx->buffer[i * 4 + 1] << 16) |
               ((uint32_t)ctx->buffer[i * 4 + 2] << 8) |
               ((uint32_t)ctx->buffer[i * 4 + 3]);
    }
    for (i = 16; i < 64; i++) {
        W[i] = SIG1(W[i - 2]) + W[i - 7] + SIG0(W[i - 15]) + W[i - 16];
    }

    a = ctx->state[0]; b = ctx->state[1];
    c = ctx->state[2]; d = ctx->state[3];
    e = ctx->state[4]; f = ctx->state[5];
    g = ctx->state[6]; h = ctx->state[7];

    for (i = 0; i < 64; i++) {
        t1 = h + EP1(e) + CH(e, f, g) + K256[i] + W[i];
        t2 = EP0(a) + MAJ(a, b, c);
        h = g; g = f; f = e; e = d + t1;
        d = c; c = b; b = a; a = t1 + t2;
    }

    ctx->state[0] += a; ctx->state[1] += b;
    ctx->state[2] += c; ctx->state[3] += d;
    ctx->state[4] += e; ctx->state[5] += f;
    ctx->state[6] += g; ctx->state[7] += h;
}

static void sha256_init(sha256_ctx_t *ctx)
{
    ctx->state[0] = 0x6A09E667;
    ctx->state[1] = 0xBB67AE85;
    ctx->state[2] = 0x3C6EF372;
    ctx->state[3] = 0xA54FF53A;
    ctx->state[4] = 0x510E527F;
    ctx->state[5] = 0x9B05688C;
    ctx->state[6] = 0x1F83D9AB;
    ctx->state[7] = 0x5BE0CD19;
    ctx->count = 0;
}

static void sha256_update(sha256_ctx_t *ctx, const uint8_t *data, uint32_t len)
{
    uint32_t idx = ctx->count & 0x3F;
    ctx->count += len;

    while (len > 0) {
        uint32_t space = 64 - idx;
        uint32_t copy = (len < space) ? len : space;
        for (uint32_t i = 0; i < copy; i++) {
            ctx->buffer[idx + i] = data[i];
        }
        data += copy;
        len -= copy;
        idx += copy;
        if (idx == 64) {
            sha256_transform(ctx);
            idx = 0;
        }
    }
}

static void sha256_final(sha256_ctx_t *ctx, uint8_t *hash)
{
    uint32_t idx = ctx->count & 0x3F;
    uint32_t pad_len = (idx < 56) ? (56 - idx) : (120 - idx);
    uint8_t pad[64];
    uint32_t i;

    pad[0] = 0x80;
    for (i = 1; i < pad_len; i++) pad[i] = 0;
    sha256_update(ctx, pad, pad_len);

    /* Append bit count (big-endian) */
    uint64_t bit_count = ctx->count * 8;
    uint8_t count_buf[8];
    for (i = 0; i < 8; i++) {
        count_buf[i] = (uint8_t)(bit_count >> (56 - i * 8));
    }
    sha256_update(ctx, count_buf, 8);

    /* Output hash */
    for (i = 0; i < 8; i++) {
        hash[i * 4]     = (uint8_t)(ctx->state[i] >> 24);
        hash[i * 4 + 1] = (uint8_t)(ctx->state[i] >> 16);
        hash[i * 4 + 2] = (uint8_t)(ctx->state[i] >> 8);
        hash[i * 4 + 3] = (uint8_t)(ctx->state[i]);
    }
}

void CRYPTO_HMAC_SHA256(const uint8_t *key, uint32_t key_len,
                        const uint8_t *data, uint32_t data_len,
                        uint8_t *mac)
{
    sha256_ctx_t ctx;
    uint8_t k_ipad[64];
    uint8_t k_opad[64];
    uint8_t key_buf[64];
    uint32_t i;
    const uint8_t *k = key;

    /* If key is longer than block size, hash it first */
    if (key_len > 64) {
        sha256_init(&ctx);
        sha256_update(&ctx, key, key_len);
        sha256_final(&ctx, key_buf);
        k = key_buf;
        key_len = 32;
    }

    for (i = 0; i < key_len; i++) {
        k_ipad[i] = k[i] ^ 0x36;
        k_opad[i] = k[i] ^ 0x5C;
    }
    for (; i < 64; i++) {
        k_ipad[i] = 0x36;
        k_opad[i] = 0x5C;
    }

    /* Inner hash: H(K XOR ipad || data) */
    sha256_init(&ctx);
    sha256_update(&ctx, k_ipad, 64);
    sha256_update(&ctx, data, data_len);
    sha256_final(&ctx, mac);

    /* Outer hash: H(K XOR opad || H_inner) */
    sha256_init(&ctx);
    sha256_update(&ctx, k_opad, 64);
    sha256_update(&ctx, mac, 32);
    sha256_final(&ctx, mac);
}

/* Constant-time memory compare */
static bool memcmp_ct(const uint8_t *a, const uint8_t *b, uint32_t len)
{
    uint8_t diff = 0;
    for (uint32_t i = 0; i < len; i++) {
        diff |= (a[i] ^ b[i]);
    }
    return (diff == 0);
}

bool CRYPTO_HMAC_Verify(const uint8_t *key, uint32_t key_len,
                        const uint8_t *data, uint32_t data_len,
                        const uint8_t *tag)
{
    uint8_t computed[32];
    CRYPTO_HMAC_SHA256(key, key_len, data, data_len, computed);
    return memcmp_ct(computed, tag, 32);
}


/* ==================== High-level OTA Decryption ==================== */

bool CRYPTO_OTA_DecryptFirmware(const uint8_t *enc_buf, uint32_t enc_len,
                                const uint8_t *aes_key, const uint8_t *hmac_key,
                                uint8_t *out_buf, uint32_t *out_len)
{
    uint32_t payload_len;
    bool ret = false;

    if (enc_len <= OTA_ENC_OVERHEAD) return false;

    payload_len = enc_len - OTA_ENC_OVERHEAD;

    /* Step 1: Verify HMAC first (before decryption) */
    if (!CRYPTO_HMAC_Verify(hmac_key, 32, enc_buf, enc_len - OTA_ENC_HMAC_SIZE,
                            enc_buf + enc_len - OTA_ENC_HMAC_SIZE)) {
        return false;
    }

    /* Step 2: Copy IV + ciphertext to output buffer for in-place decrypt */
    for (uint32_t i = 0; i < payload_len + OTA_ENC_IV_SIZE; i++) {
        out_buf[i] = enc_buf[i];
    }

    /* Step 3: Initialize CRYP and decrypt */
    CRYPTO_AES_Init(aes_key);
    ret = CRYPTO_AES_CBC_Decrypt(out_buf, payload_len + OTA_ENC_IV_SIZE);
    CRYPTO_AES_Deinit();

    if (!ret) return false;

    /* Step 4: Remove PKCS7 padding */
    uint8_t pad_len = out_buf[payload_len - 1];
    if (pad_len == 0 || pad_len > 16) return false;
    /* Verify all padding bytes are correct */
    for (uint32_t i = payload_len - pad_len; i < payload_len; i++) {
        if (out_buf[i] != pad_len) return false;
    }

    *out_len = payload_len - pad_len;
    return true;
}