#include "RSA.h"
#include <assert.h>
#include "bearssl_x509.h"
#include "bearssl_pem.h"
#include "tlibc/base64.h"

// https://crypto.stackexchange.com/questions/3110/impacts-of-not-using-rsa-exponent-of-65537
#define DEFAULT_PUBLIC_EXPONENT 65537

Result(void) RSA_generateKeyPair(u32 key_size,
    br_rsa_private_key* sk, br_rsa_public_key* pk,
    const br_prng_class** rng_vtable_ptr)
{
    Deferral(4);
    bool success = false;
    
    void* sk_buf = malloc(BR_RSA_KBUF_PRIV_SIZE(key_size));
    void* pk_buf = malloc(BR_RSA_KBUF_PUB_SIZE(key_size));
    Defer(
        if(!success){
            free(sk_buf);
            free(pk_buf);
        }
    );
    
    success = br_rsa_i31_keygen(rng_vtable_ptr, sk, sk_buf, pk, pk_buf, key_size, DEFAULT_PUBLIC_EXPONENT);
    if(!success){
        Return RESULT_ERROR("br_rsa_i31_keygen() failed", false);
    }

    Return RESULT_VOID;
}

Result(void) RSA_generateKeyPairFromSystemRandom(u32 key_size, 
    br_rsa_private_key* sk, br_rsa_public_key* pk)
{
    Deferral(4);
    br_hmac_drbg_context time_based_rng = { .vtable = &br_hmac_drbg_vtable };
    rng_init_sha256_seedFromSystem(&time_based_rng.vtable);
    try_void(RSA_generateKeyPair(key_size, sk, pk, &time_based_rng.vtable));
    Return RESULT_VOID;
}

Result(void) RSA_generateKeyPairFromPassword(u32 key_size, 
    br_rsa_private_key* sk, br_rsa_public_key* pk, str password)
{
    Deferral(4);
    br_hmac_drbg_context password_based_rng = { .vtable = &br_hmac_drbg_vtable };
    br_hmac_drbg_init(&password_based_rng, &br_sha256_vtable, password.data, password.size);
    try_void(RSA_generateKeyPair(key_size, sk, pk, &password_based_rng.vtable));
    Return RESULT_VOID;
}

Result(void) RSA_computePublicKey(const br_rsa_private_key* sk, br_rsa_public_key* pk){
    Deferral(4);
    br_rsa_compute_modulus compute_modulus = br_rsa_i31_compute_modulus;
	br_rsa_compute_pubexp compute_pubexp = br_rsa_i31_compute_pubexp;
    
    size_t modulus_size = compute_modulus(NULL, sk);
	if (modulus_size == 0) {
		Return RESULT_ERROR("compute_modulus", false);
	}
    void* modulus = malloc(modulus_size);
    bool success = false;
    Defer(
        if(!success)
            free(modulus)
    );
	if (compute_modulus(modulus, sk) != modulus_size) {
		Return RESULT_ERROR("compute_modulus", false);
	}
    
	u32 pubexp_little_endian = compute_pubexp(sk);
	if (pubexp_little_endian == 0) {
		Return RESULT_ERROR("compute_pubexp", false);
	}
    u8 pubexp_big_endian[4];
    pubexp_big_endian[0] = pubexp_little_endian >> 24;
	pubexp_big_endian[1] = pubexp_little_endian >> 16;
	pubexp_big_endian[2] = pubexp_little_endian >> 8;
	pubexp_big_endian[3] = pubexp_little_endian;

	pk->n = modulus;
	pk->nlen = modulus_size;
	pk->e = pubexp_big_endian;
	pk->elen = sizeof pubexp_big_endian;
    success = true;
    Return RESULT_VOID;
}

str RSA_serializePrivateKey_base64(const br_rsa_private_key* sk){
    u32 key_buffer_size = BR_RSA_KBUF_PRIV_SIZE(sk->n_bitlen);
    u32 key_base64_size = base64_encodedSize(key_buffer_size);
    char* serialized_buf = malloc(32 + key_base64_size);
    sprintf(serialized_buf, "RSA-Private-%u:", sk->n_bitlen);
    u32 offset = strlen(serialized_buf);
    offset += base64_encode(sk->p, key_buffer_size, serialized_buf + offset);
    serialized_buf[offset] = '\0';
    return str_construct(serialized_buf, offset, true);
}

str RSA_serializePublicKey_base64(const br_rsa_public_key* pk){
    u32 n_bitlen = pk->nlen * 8;
    u32 key_buffer_size = BR_RSA_KBUF_PUB_SIZE(n_bitlen);
    u32 key_base64_size = base64_encodedSize(key_buffer_size);
    char* serialized_buf = malloc(32 + key_base64_size);
    sprintf(serialized_buf, "RSA-Public-%u:", n_bitlen);
    u32 offset = strlen(serialized_buf);
    offset += base64_encode(pk->n, key_buffer_size, serialized_buf + offset);
    serialized_buf[offset] = '\0';
    return str_construct(serialized_buf, offset, true);
}

Result(void) RSA_parsePublicKey_base64(cstr src, br_rsa_public_key* pk){
    Deferral(4);
    u32 n_bitlen = 0;
    if(sscanf(src, "RSA-Public-%u:", &n_bitlen) != 1){
        Return RESULT_ERROR("can't parse key size", false);
    }
    u32 key_buffer_size = BR_RSA_KBUF_PUB_SIZE(n_bitlen);
    pk->n = malloc(key_buffer_size);
    pk->elen = 4;
    pk->nlen = key_buffer_size - 4;
    pk->e = pk->n + pk->nlen;
    str src_str = str_from_cstr(src);
    u32 offset = str_seekChar(src_str, ':', 10) + 1;
    if(offset == 0){
        Return RESULT_ERROR("missing ':' before key data", false);
    }
    str key_base64_str = src_str;
    key_base64_str.data += offset;
    key_base64_str.size -= offset;
    u32 decoded_size = base64_decodedSize(key_base64_str.data, key_base64_str.size);
    if(decoded_size != key_buffer_size){
        Return RESULT_ERROR_FMT("decoded key size is %u, must be %u", decoded_size, key_buffer_size);
    }
    decoded_size = base64_decode(key_base64_str.data, key_base64_str.size, pk->n);
    if(decoded_size != key_buffer_size){
        Return RESULT_ERROR("key decoding failed", false);
    }
    Return RESULT_VOID;
}

Result(void) RSA_parsePrivateKey_base64(cstr src, br_rsa_private_key* sk){
    Deferral(4);
    u32 n_bitlen = 0;
    if(sscanf(src, "RSA-Private-%u:", &n_bitlen) != 1){
        Return RESULT_ERROR("can't parse key size", false);
    }
    sk->n_bitlen = n_bitlen;
    u32 key_buffer_size = BR_RSA_KBUF_PRIV_SIZE(n_bitlen);
    u32 field_len = key_buffer_size / 5;
    sk->plen = sk->qlen = sk->dplen = sk->dqlen = sk->iqlen = field_len;
    sk->p = malloc(key_buffer_size);
    sk->q = sk->p + field_len;
    sk->dp = sk->q + field_len;
    sk->dq = sk->dp + field_len;
    sk->iq = sk->dq + field_len;
    str src_str = str_from_cstr(src);
    u32 offset = str_seekChar(src_str, ':', 10) + 1;
    if(offset == 0){
        Return RESULT_ERROR("missing ':' before key data", false);
    }
    str key_base64_str = src_str;
    key_base64_str.data += offset;
    key_base64_str.size -= offset;
    u32 decoded_size = base64_decodedSize(key_base64_str.data, key_base64_str.size);
    if(decoded_size != key_buffer_size){
        Return RESULT_ERROR_FMT("decoded key size is %u, must be %u", decoded_size, key_buffer_size);
    }
    decoded_size = base64_decode(key_base64_str.data, key_base64_str.size, sk->p);
    if(decoded_size != key_buffer_size){
        Return RESULT_ERROR("key decoding failed", false);
    }
    Return RESULT_VOID;
}


//////////////////////////////////////////////////////////////////////////////
//                               RSAEncryptor                               //
//////////////////////////////////////////////////////////////////////////////

void RSAEncryptor_construct(RSAEncryptor* ptr, const br_rsa_public_key* pk){
    ptr->pk = pk;
    ptr->rng.vtable = &br_hmac_drbg_vtable;
    rng_init_sha256_seedFromSystem(&ptr->rng.vtable);
}

Result(u32) RSAEncryptor_encrypt(RSAEncryptor* ptr, Array(u8) src, Array(u8) dst){
    u32 key_size_bytes = ptr->pk->nlen;
    const u32 max_src_size = RSAEncryptor_calcMaxSrcSize(key_size_bytes * 8, 256);
    if(src.size > max_src_size){  
        return RESULT_ERROR_FMT("src.size (%u) must be <= RSAEncryptor_calcMaxSrcSize() (%u)", 
            src.size, max_src_size);
    }
    if(dst.size < key_size_bytes){
        return RESULT_ERROR_FMT("dst.size (%u) must be >= key length in bytes (%u)", 
            dst.size, key_size_bytes);
    }
    size_t sz = br_rsa_i31_oaep_encrypt(
        &ptr->rng.vtable, &br_sha256_vtable, 
        NULL, 0,
        ptr->pk,
        dst.data, dst.size,
        src.data, src.size);
    
    if(sz == 0){
        return RESULT_ERROR("RSA encryption failed", false);
    }
    return RESULT_VALUE(u, sz);
}


//////////////////////////////////////////////////////////////////////////////
//                               RSADecryptor                               //
//////////////////////////////////////////////////////////////////////////////

void RSADecryptor_construct(RSADecryptor* ptr, const br_rsa_private_key* sk){
    ptr->sk = sk;
}

Result(u32) RSADecryptor_decrypt(RSADecryptor* ptr, Array(u8) buffer){
    u32 key_size_bits = ptr->sk->n_bitlen;
    if(buffer.size != key_size_bits/8){
        return RESULT_ERROR_FMT("buffer.size (%u) must be == key length in bytes (%u)", 
            buffer.size, key_size_bits/8);
    }

    size_t sz = buffer.size;
    size_t r = br_rsa_i31_oaep_decrypt(
        &br_sha256_vtable, 
        NULL, 0,
        ptr->sk,
        buffer.data, &sz);
    
    if(r == 0){
        return RESULT_ERROR("RSA encryption failed", false);
    }
    return RESULT_VALUE(u, sz);
}