/*---------------------------------------------------------------------------*/
/* mbedtls_ecc.cpp                                                           */
/* Copyright (c) innovaphone 2025                                            */
/*                                                                           */
/* Elliptic curve cryptography using the mbedTLS library                     */
/*---------------------------------------------------------------------------*/

#include <platform.h>
#include "common/ilib/ecc.h"
#include "mbedtls_config.h"
#include "mbedtls/ecp.h"
#include "mbedtls/bignum.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/entropy.h"

USING_NAMESPACE

static bool innovaphone_mbedtls_initialized = false;
static mbedtls_ctr_drbg_context global_drbg_ctx;
static const unsigned char global_drbg_personalization_string[] = "innovaphone_mbedtls_personalization_string";

static void innovaphone_mbedtls_init()
{
    if (!innovaphone_mbedtls_initialized) {
        innovaphone_mbedtls_initialized = true;
        mbedtls_ctr_drbg_init(&global_drbg_ctx);
        mbedtls_ctr_drbg_seed(&global_drbg_ctx, innovaphone_mbedtls_get_random_bytes, NULL, global_drbg_personalization_string, sizeof(global_drbg_personalization_string) - 1);
    }
}

static mbedtls_ecp_group_id get_mbedtls_curve_id(elliptic_curve_t curve)
{
    switch (curve) {
    case EC_SECP256R1:
        return MBEDTLS_ECP_DP_SECP256R1;
    case EC_SECP384R1:
        return MBEDTLS_ECP_DP_SECP384R1;
    case EC_SECP521R1:
        return MBEDTLS_ECP_DP_SECP521R1;
    case EC_UNDEFINED:
    default:
        return MBEDTLS_ECP_DP_NONE;
    }
}

bool ecc::elliptic_curve_supported(elliptic_curve_t curve)
{
    switch (curve) {
    case EC_SECP256R1:
    case EC_SECP384R1:
    case EC_SECP521R1:
        return true;
    default:
        return false;
    }
    return false;
}

size_t ecc::elliptic_curve_size(elliptic_curve_t curve)
{
    switch (curve) {
    case EC_CURVE25519:
    case EC_SECP256R1:
        return 32;
    case EC_SECP384R1:
        return 48;
    case EC_SECP521R1:
        return 66;
    default:
        return 0;
    }
    return 0;
}

size_t ecc::signature_size(elliptic_curve_t curve)
{
    return elliptic_curve_size(curve) * 2;
}

size_t ecc::pkix_signature_size(elliptic_curve_t curve)
{
    return elliptic_curve_size(curve) * 2 + 10;
}

bool ecc::make_key(byte * public_key, byte * private_key, elliptic_curve_t curve)
{
    innovaphone_mbedtls_init();
    if (!ecc::elliptic_curve_supported(curve) || !public_key || !private_key) return false;
    
    size_t curve_size = ecc::elliptic_curve_size(curve);
    
    int ret;
    mbedtls_ecp_group grp;
    mbedtls_mpi d;
    mbedtls_ecp_point Q;
    mbedtls_ecp_group_init(&grp);
    mbedtls_mpi_init(&d);
    mbedtls_ecp_point_init(&Q);

    // load curve
    mbedtls_ecp_group_id curve_id = get_mbedtls_curve_id(curve);
    if (curve_id == MBEDTLS_ECP_DP_NONE) return false;
    ret = mbedtls_ecp_group_load(&grp, curve_id);
    if (ret != 0) goto error;

    // generate key pair
    ret = mbedtls_ecp_gen_keypair(&grp, &d, &Q, mbedtls_ctr_drbg_random, &global_drbg_ctx);
    if (ret != 0) goto error;

    // Export private key
    ret = mbedtls_mpi_write_binary(&d, private_key, curve_size);
    if (ret != 0) goto error;
    
    // Export public key
    byte tmp[EC_CURVE_SIZE_MAX * 2 + 1];
    size_t tmp_len;
    ret = mbedtls_ecp_point_write_binary(&grp, &Q, MBEDTLS_ECP_PF_UNCOMPRESSED, &tmp_len, tmp, sizeof(tmp));
    if (ret != 0) goto error;
    memcpy(public_key, tmp + 1, tmp_len - 1);

    mbedtls_ecp_group_free(&grp);
    mbedtls_mpi_free(&d);
    mbedtls_ecp_point_free(&Q);
    return true;

error:
    mbedtls_ecp_group_free(&grp);
    mbedtls_mpi_free(&d);
    mbedtls_ecp_point_free(&Q);
    return false;
}

bool ecc::shared_secret(byte * shared_secret, const byte * public_key, const byte * private_key, elliptic_curve_t curve)
{
    innovaphone_mbedtls_init();
    if (!shared_secret || !public_key || !private_key || !ecc::elliptic_curve_supported(curve)) return false;

    size_t curve_size = ecc::elliptic_curve_size(curve);
    if (curve_size == 0) return false;

    int ret;
    unsigned char tmp[1 + EC_CURVE_SIZE_MAX * 2];
    size_t tmp_len = 1 + 2 * curve_size;
    mbedtls_ecp_group grp;
    mbedtls_ecp_point Q, R;
    mbedtls_mpi d;

    mbedtls_ecp_group_init(&grp);
    mbedtls_ecp_point_init(&Q);
    mbedtls_ecp_point_init(&R);
    mbedtls_mpi_init(&d);

    mbedtls_ecp_group_id curve_id = get_mbedtls_curve_id(curve);
    if (curve_id == MBEDTLS_ECP_DP_NONE) goto error;

    ret = mbedtls_ecp_group_load(&grp, curve_id);
    if (ret != 0) goto error;

    // import private key
    ret = mbedtls_mpi_read_binary(&d, private_key, curve_size);
    if (ret != 0) goto error;

    // import public key
    tmp[0] = 0x04;
    memcpy(tmp + 1, public_key, 2 * curve_size);
    ret = mbedtls_ecp_point_read_binary(&grp, &Q, tmp, tmp_len);
    if (ret != 0) goto error;
    ret = mbedtls_ecp_check_pubkey(&grp, &Q);
    if (ret != 0) goto error;

    // caluclate R
    ret = mbedtls_ecp_mul(&grp, &R, &d, &Q, mbedtls_ctr_drbg_random, &global_drbg_ctx);
    if (ret != 0) goto error;

    // export shared secret (X coordinate of R)
    unsigned char outbuf[1 + EC_CURVE_SIZE_MAX * 2];
    size_t outlen;
    ret = mbedtls_ecp_point_write_binary(&grp, &R, MBEDTLS_ECP_PF_UNCOMPRESSED, &outlen, outbuf, sizeof(outbuf));
    if (ret != 0) goto error;
    if (outlen != 1 + 2 * curve_size) goto error;
    memcpy(shared_secret, outbuf + 1, curve_size);

    mbedtls_ecp_group_free(&grp);
    mbedtls_ecp_point_free(&Q);
    mbedtls_ecp_point_free(&R);
    mbedtls_mpi_free(&d);
    return true;

error:
    mbedtls_ecp_group_free(&grp);
    mbedtls_ecp_point_free(&Q);
    mbedtls_ecp_point_free(&R);
    mbedtls_mpi_free(&d);
    return false;
}

bool ecc::sign(byte * signature, const byte * hash, size_t hash_len, const byte * private_key, elliptic_curve_t curve)
{
    innovaphone_mbedtls_init();
    if (!signature || !hash || !private_key || !ecc::elliptic_curve_supported(curve)) return false;

    size_t curve_size = ecc::elliptic_curve_size(curve);
    if (curve_size == 0) return false;

    int ret;
    mbedtls_ecp_group grp;
    mbedtls_mpi d, k, r, s, h_mpi, k_inv, tmp1;
    mbedtls_ecp_point R;
    byte tmp_point_buffer[1 + EC_CURVE_SIZE_MAX * 2];
    size_t tmp_len;

    mbedtls_ecp_group_init(&grp);
    mbedtls_mpi_init(&d); mbedtls_mpi_init(&k); mbedtls_mpi_init(&r);
    mbedtls_mpi_init(&s); mbedtls_mpi_init(&h_mpi); mbedtls_mpi_init(&k_inv);
    mbedtls_mpi_init(&tmp1);
    mbedtls_ecp_point_init(&R);

    // load group
    mbedtls_ecp_group_id curve_id = get_mbedtls_curve_id(curve);
    if (curve_id == MBEDTLS_ECP_DP_NONE) { ret = -1; goto error; }
    ret = mbedtls_ecp_group_load(&grp, curve_id);
    if (ret != 0) goto error;

    // import private key
    ret = mbedtls_mpi_read_binary(&d, private_key, curve_size);
    if (ret != 0) goto error;

    // import hash
    if (hash_len > mbedtls_mpi_size(&grp.N)) {
        hash_len = mbedtls_mpi_size(&grp.N);
    }
    ret = mbedtls_mpi_read_binary(&h_mpi, hash, hash_len);
    if (ret != 0) goto error;

    // sign
    mbedtls_mpi_lset(&r, 0);
    mbedtls_mpi_lset(&s, 0);

    do {
        ret = mbedtls_ecp_gen_privkey(&grp, &k, mbedtls_ctr_drbg_random, &global_drbg_ctx);
        if (ret != 0) goto error;
        ret = mbedtls_ecp_mul(&grp, &R, &k, &grp.G, mbedtls_ctr_drbg_random, &global_drbg_ctx);
        if (ret != 0) goto error;
        ret = mbedtls_ecp_point_write_binary(
            &grp, &R, MBEDTLS_ECP_PF_UNCOMPRESSED,
            &tmp_len, tmp_point_buffer, sizeof(tmp_point_buffer)
        );
        if (ret != 0) goto error;
        ret = mbedtls_mpi_read_binary(&r, tmp_point_buffer + 1, curve_size);
        if (ret != 0) goto error;
        ret = mbedtls_mpi_mod_mpi(&r, &r, &grp.N);
        if (ret != 0) goto error;
        if (mbedtls_mpi_cmp_int(&r, 0) == 0) {
            mbedtls_mpi_lset(&s, 0);
            continue;
        }
        ret = mbedtls_mpi_inv_mod(&k_inv, &k, &grp.N);
        if (ret != 0) goto error;
        ret = mbedtls_mpi_mul_mpi(&tmp1, &r, &d);
        if (ret != 0) goto error;
        ret = mbedtls_mpi_mod_mpi(&tmp1, &tmp1, &grp.N);
        if (ret != 0) goto error;
        ret = mbedtls_mpi_add_mpi(&tmp1, &h_mpi, &tmp1);
        if (ret != 0) goto error;
        ret = mbedtls_mpi_mod_mpi(&tmp1, &tmp1, &grp.N);
        if (ret != 0) goto error;
        ret = mbedtls_mpi_mul_mpi(&s, &k_inv, &tmp1);
        if (ret != 0) goto error;
        ret = mbedtls_mpi_mod_mpi(&s, &s, &grp.N);
        if (ret != 0) goto error;

    } while (mbedtls_mpi_cmp_int(&r, 0) == 0 || mbedtls_mpi_cmp_int(&s, 0) == 0);

    ret = mbedtls_mpi_write_binary(&r, signature, curve_size);
    if (ret != 0) goto error;
    ret = mbedtls_mpi_write_binary(&s, signature + curve_size, curve_size);
    if (ret != 0) goto error;

    mbedtls_ecp_point_free(&R);
    mbedtls_mpi_free(&tmp1); mbedtls_mpi_free(&k_inv); mbedtls_mpi_free(&h_mpi);
    mbedtls_mpi_free(&s); mbedtls_mpi_free(&r); mbedtls_mpi_free(&k);
    mbedtls_mpi_free(&d);
    mbedtls_ecp_group_free(&grp);
    return true;

error:
    mbedtls_ecp_point_free(&R);
    mbedtls_mpi_free(&tmp1); mbedtls_mpi_free(&k_inv); mbedtls_mpi_free(&h_mpi);
    mbedtls_mpi_free(&s); mbedtls_mpi_free(&r); mbedtls_mpi_free(&k);
    mbedtls_mpi_free(&d);
    mbedtls_ecp_group_free(&grp);
    return false;
}

bool ecc::verify(const byte * signature, const byte * hash, size_t hash_len, const byte * public_key, elliptic_curve_t curve)
{
    innovaphone_mbedtls_init();
    if (!signature || !hash || !public_key || !ecc::elliptic_curve_supported(curve)) return false;

    size_t curve_size = ecc::elliptic_curve_size(curve);
    if (curve_size == 0) return false;

    bool success = false;
    int ret;
    mbedtls_ecp_group grp;
    mbedtls_mpi r, s, h_mpi, w, u1, u2, v, tmp1, tmp2;
    mbedtls_ecp_point Q, V;
    byte tmp_point_buffer[1 + EC_CURVE_SIZE_MAX * 2];
    size_t tmp_len;

    mbedtls_ecp_group_init(&grp);
    mbedtls_mpi_init(&r); mbedtls_mpi_init(&s); mbedtls_mpi_init(&h_mpi);
    mbedtls_mpi_init(&w); mbedtls_mpi_init(&u1); mbedtls_mpi_init(&u2);
    mbedtls_mpi_init(&v); mbedtls_mpi_init(&tmp1); mbedtls_mpi_init(&tmp2);
    mbedtls_ecp_point_init(&Q); mbedtls_ecp_point_init(&V);

    // load group
    mbedtls_ecp_group_id curve_id = get_mbedtls_curve_id(curve);
    if (curve_id == MBEDTLS_ECP_DP_NONE) { ret = -1; goto error; }
    ret = mbedtls_ecp_group_load(&grp, curve_id);
    if (ret != 0) goto error;

    // import signature
    ret = mbedtls_mpi_read_binary(&r, signature, curve_size);
    if (ret != 0) goto error;
    ret = mbedtls_mpi_read_binary(&s, signature + curve_size, curve_size);
    if (ret != 0) goto error;

    if (hash_len > mbedtls_mpi_size(&grp.N)) {
        hash_len = mbedtls_mpi_size(&grp.N);
    }
    ret = mbedtls_mpi_read_binary(&h_mpi, hash, hash_len);
    if (ret != 0) goto error;

    // check validity
    if (mbedtls_mpi_cmp_int(&r, 1) < 0 || mbedtls_mpi_cmp_mpi(&r, &grp.N) >= 0 ||
        mbedtls_mpi_cmp_int(&s, 1) < 0 || mbedtls_mpi_cmp_mpi(&s, &grp.N) >= 0) {
        ret = MBEDTLS_ERR_ECP_INVALID_KEY;
        goto error;
    }

    // import public key
    unsigned char tmp_pub[1 + EC_CURVE_SIZE_MAX * 2];
    tmp_pub[0] = 0x04;
    memcpy(tmp_pub + 1, public_key, 2 * curve_size);
    ret = mbedtls_ecp_point_read_binary(&grp, &Q, tmp_pub, 1 + 2 * curve_size);
    if (ret != 0) goto error;
    ret = mbedtls_ecp_check_pubkey(&grp, &Q);
    if (ret != 0) goto error;

    // verify
    ret = mbedtls_mpi_inv_mod(&w, &s, &grp.N);
    if (ret != 0) goto error;
    ret = mbedtls_mpi_mul_mpi(&u1, &h_mpi, &w);
    if (ret != 0) goto error;
    ret = mbedtls_mpi_mod_mpi(&u1, &u1, &grp.N);
    if (ret != 0) goto error;
    ret = mbedtls_mpi_mul_mpi(&u2, &r, &w);
    if (ret != 0) goto error;
    ret = mbedtls_mpi_mod_mpi(&u2, &u2, &grp.N);
    if (ret != 0) goto error;
    ret = mbedtls_ecp_muladd(&grp, &V, &u1, &grp.G, &u2, &Q);
    if (ret != 0) goto error;
    ret = mbedtls_ecp_point_write_binary(&grp, &V, MBEDTLS_ECP_PF_UNCOMPRESSED, &tmp_len, tmp_point_buffer, sizeof(tmp_point_buffer));
    if (ret != 0) goto error;
    ret = mbedtls_mpi_read_binary(&v, tmp_point_buffer + 1, curve_size);
    if (ret != 0) goto error;
    ret = mbedtls_mpi_mod_mpi(&v, &v, &grp.N);
    if (ret != 0) goto error;
    if (mbedtls_ecp_is_zero(&V)) {
        ret = MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
        goto error;
    }
    success = (mbedtls_mpi_cmp_mpi(&v, &r) == 0);

    mbedtls_ecp_point_free(&V); mbedtls_ecp_point_free(&Q);
    mbedtls_mpi_free(&tmp2); mbedtls_mpi_free(&tmp1); mbedtls_mpi_free(&v);
    mbedtls_mpi_free(&u2); mbedtls_mpi_free(&u1); mbedtls_mpi_free(&w);
    mbedtls_mpi_free(&h_mpi); mbedtls_mpi_free(&s); mbedtls_mpi_free(&r);
    mbedtls_ecp_group_free(&grp);
    return success;

error:
    mbedtls_ecp_point_free(&V); mbedtls_ecp_point_free(&Q);
    mbedtls_mpi_free(&tmp2); mbedtls_mpi_free(&tmp1); mbedtls_mpi_free(&v);
    mbedtls_mpi_free(&u2); mbedtls_mpi_free(&u1); mbedtls_mpi_free(&w);
    mbedtls_mpi_free(&h_mpi); mbedtls_mpi_free(&s); mbedtls_mpi_free(&r);
    mbedtls_ecp_group_free(&grp);
    return false;
}

static size_t pkix_sign_write_point(byte * out, const byte * data, size_t size)
{
    size_t offset = 0;
    while (offset < size - 1 && data[offset] == 0x00) offset++;
    size_t len = size - offset;
    size_t padding_len = (data[offset] & 0x80) != 0 ? 1 : 0;

    out[0] = 0x02;
    out[1] = (byte)(len + padding_len);
    if (padding_len) out[2] = 0;
    memcpy(out + 2 + padding_len, data + offset, len);
    return 2 + padding_len + len;
}

static size_t pkix_verify_read_point(const byte * signature, size_t size, byte * out)
{
    if (*signature != 0x02) return 0;
    size_t len_data = signature[1];
    size_t pos = 2;
    size_t len_copy = len_data;
    if (len_data > size && signature[pos] == 0x00) {
        pos++;
        len_copy--;
    }
    if (len_copy == 0 || len_copy > size) return 0;
    memcpy(out + size - len_copy, signature + pos, len_copy);
    return 2 + len_data;
}

/*
    https://datatracker.ietf.org/doc/html/rfc3279#section-2.2.3
          Ecdsa-Sig-Value  ::=  SEQUENCE  {
           r     INTEGER,
           s     INTEGER  }
*/

size_t ecc::pkix_sign(byte * signature, const byte * hash, size_t hash_len, const byte * private_key, elliptic_curve_t curve)
{
    size_t size = elliptic_curve_size(curve);
    size_t header_len = 0;
    byte * sig = (byte *)alloca(size * 2);
    memset(sig, 0, size * 2);
    if (!sign(sig, hash, hash_len, private_key, curve)) return 0;

    byte tmp[512];
    size_t len = 0;
    len += pkix_sign_write_point(tmp + len, sig, size);
    len += pkix_sign_write_point(tmp + len, sig + size, size);
    signature[0] = 0x30;
    if (len <= 127) {
        header_len = 2;
        signature[1] = (byte)len;
    }
    else {
        header_len = 3;
        signature[1] = 0x81;
        signature[2] = (byte)len;
    }
    memcpy(signature + header_len, tmp, len);
    return len + header_len;
}

bool ecc::pkix_verify(const byte * signature, size_t signature_len, const byte * hash, size_t hash_len, const byte * public_key, elliptic_curve_t curve)
{
    size_t curve_size = elliptic_curve_size(curve);
    if (signature_len < 10 || signature[0] != 0x30) return false;

    size_t pos;
    size_t len;
    size_t point_len;

    if ((signature[1] & 0x80) == 0) {
        pos = 2;
        len = signature[1];
    }
    else {
        if (signature[1] != 0x81) return false;
        if (signature_len < 3) return false;
        pos = 3;
        len = signature[2];
    }
    if (len + pos != signature_len) return false;

    byte * sig = (byte *)alloca(curve_size * 2);
    memset(sig, 0, curve_size * 2);

    point_len = pkix_verify_read_point(signature + pos, curve_size, sig);
    if (point_len == 0) return false;
    pos += point_len;

    point_len = pkix_verify_read_point(signature + pos, curve_size, sig + curve_size);
    if (point_len == 0) return false;
    pos += point_len;

    if (pos != signature_len) return false;
    return verify(sig, hash, hash_len, public_key, curve);
}
