/*
 * LITLS -- HKDF (RFC 5869) using HMAC-SHA256 and HMAC-SHA384
 */

#include "litls.h"
#include <string.h>

extern void litls_secure_zero(void *buf, size_t len);

void litls_hkdf_extract_sha256(const uint8_t *salt, size_t salt_len, const uint8_t *ikm, size_t ikm_len,
                               uint8_t prk[32]) {
    if (salt_len == 0) {
        uint8_t zero_salt[32];
        memset(zero_salt, 0, 32);
        litls_hmac_sha256(zero_salt, 32, ikm, ikm_len, prk);
    } else {
        litls_hmac_sha256(salt, salt_len, ikm, ikm_len, prk);
    }
}

int litls_hkdf_expand_sha256(const uint8_t prk[32], const uint8_t *info, size_t info_len, uint8_t *okm,
                             size_t okm_len) {
    if (okm_len > 255 * 32)
        return -1;

    uint8_t t[32];
    size_t t_len    = 0;
    uint8_t counter = 1;
    size_t offset   = 0;

    while (offset < okm_len) {
        /* T(n) = HMAC(PRK, T(n-1) || info || counter) */
        litls_sha256_ctx hctx;
        uint8_t k_pad[LITLS_SHA256_BLOCK_SIZE];
        uint8_t inner[32];

        /* Build HMAC input: T_prev || info || counter */
        /* We need to compute HMAC manually to stream the concatenation */
        uint8_t key_hash[32];
        const uint8_t *hmac_key = prk;
        size_t hmac_key_len     = 32;
        (void)key_hash; /* PRK is always 32 bytes, no hashing needed */

        /* ipad */
        memset(k_pad, 0x36, LITLS_SHA256_BLOCK_SIZE);
        for (size_t i = 0; i < hmac_key_len; i++)
            k_pad[i] ^= hmac_key[i];

        litls_sha256_init(&hctx);
        litls_sha256_update(&hctx, k_pad, LITLS_SHA256_BLOCK_SIZE);
        if (t_len > 0)
            litls_sha256_update(&hctx, t, t_len);
        litls_sha256_update(&hctx, info, info_len);
        litls_sha256_update(&hctx, &counter, 1);
        litls_sha256_final(&hctx, inner);

        /* opad */
        memset(k_pad, 0x5c, LITLS_SHA256_BLOCK_SIZE);
        for (size_t i = 0; i < hmac_key_len; i++)
            k_pad[i] ^= hmac_key[i];

        litls_sha256_init(&hctx);
        litls_sha256_update(&hctx, k_pad, LITLS_SHA256_BLOCK_SIZE);
        litls_sha256_update(&hctx, inner, 32);
        litls_sha256_final(&hctx, t);

        t_len          = 32;
        size_t to_copy = okm_len - offset;
        if (to_copy > 32)
            to_copy = 32;
        memcpy(okm + offset, t, to_copy);
        offset += to_copy;
        counter++;

        litls_secure_zero(k_pad, sizeof(k_pad));
        litls_secure_zero(inner, sizeof(inner));
    }

    litls_secure_zero(t, sizeof(t));
    return 0;
}

void litls_hkdf_extract_sha384(const uint8_t *salt, size_t salt_len, const uint8_t *ikm, size_t ikm_len,
                               uint8_t prk[48]) {
    if (salt_len == 0) {
        uint8_t zero_salt[48];
        memset(zero_salt, 0, 48);
        litls_hmac_sha384(zero_salt, 48, ikm, ikm_len, prk);
    } else {
        litls_hmac_sha384(salt, salt_len, ikm, ikm_len, prk);
    }
}

int litls_hkdf_expand_sha384(const uint8_t prk[48], const uint8_t *info, size_t info_len, uint8_t *okm,
                             size_t okm_len) {
    if (okm_len > 255 * 48)
        return -1;

    uint8_t t[48];
    size_t t_len    = 0;
    uint8_t counter = 1;
    size_t offset   = 0;

    while (offset < okm_len) {
        litls_sha512_ctx hctx; /* sha384 uses sha512 ctx */
        uint8_t k_pad[LITLS_SHA384_BLOCK_SIZE];
        uint8_t inner[48];

        const uint8_t *hmac_key = prk;
        size_t hmac_key_len     = 48;

        memset(k_pad, 0x36, LITLS_SHA384_BLOCK_SIZE);
        for (size_t i = 0; i < hmac_key_len; i++)
            k_pad[i] ^= hmac_key[i];

        litls_sha384_init(&hctx);
        litls_sha384_update(&hctx, k_pad, LITLS_SHA384_BLOCK_SIZE);
        if (t_len > 0)
            litls_sha384_update(&hctx, t, t_len);
        litls_sha384_update(&hctx, info, info_len);
        litls_sha384_update(&hctx, &counter, 1);
        litls_sha384_final(&hctx, inner);

        memset(k_pad, 0x5c, LITLS_SHA384_BLOCK_SIZE);
        for (size_t i = 0; i < hmac_key_len; i++)
            k_pad[i] ^= hmac_key[i];

        litls_sha384_init(&hctx);
        litls_sha384_update(&hctx, k_pad, LITLS_SHA384_BLOCK_SIZE);
        litls_sha384_update(&hctx, inner, 48);
        litls_sha384_final(&hctx, t);

        t_len          = 48;
        size_t to_copy = okm_len - offset;
        if (to_copy > 48)
            to_copy = 48;
        memcpy(okm + offset, t, to_copy);
        offset += to_copy;
        counter++;

        litls_secure_zero(k_pad, sizeof(k_pad));
        litls_secure_zero(inner, sizeof(inner));
    }

    litls_secure_zero(t, sizeof(t));
    return 0;
}