#include "cryptography.h"
#include <assert.h>
#include "bearssl_x509.h"
#include "bearssl_pem.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(16);
    bool success = false;
    rng_init_sha256_seedFromTime(rng_vtable_ptr);
    
    void* sk_buf = malloc(BR_RSA_KBUF_PRIV_SIZE(key_size));
    Defer(
        if(!success)
            free(sk_buf)
    );
    
    void* pk_buf = malloc(BR_RSA_KBUF_PUB_SIZE(key_size));
    Defer(
        if(!success)
            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_computePublicKey(const br_rsa_private_key* sk, br_rsa_public_key* pk){
    Deferral(16);
    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;
}

Result(void) RSA_serializePrivateKey_RawDER(
    const br_rsa_private_key* sk,
    NULLABLE(const br_rsa_public_key*) pk,
    Array(u8)* out_der)
{
    Deferral(32);
	br_rsa_compute_pubexp compute_pubexp = br_rsa_i31_compute_pubexp;
	br_rsa_compute_privexp compute_privexp = br_rsa_i31_compute_privexp;
	
    br_rsa_public_key pk_computed;
    if(pk == NULL){
        pk = &pk_computed;
        try_void(RSA_computePublicKey(sk, &pk_computed));
        Defer(free(pk->n));
    }
	
	u32 pubexp_little_endian = compute_pubexp(sk);
	if (pubexp_little_endian == 0) {
		Return RESULT_ERROR("compute_pubexp", false);
	}
	size_t privexp_size = compute_privexp(NULL, sk, pubexp_little_endian);
	if (privexp_size == 0) {
		Return RESULT_ERROR("compute_privexp", false);
	}
	void* privexp = malloc(privexp_size);
    Defer(free(privexp));
	if (compute_privexp(privexp, sk, pubexp_little_endian) != privexp_size) {
		Return RESULT_ERROR("compute_privexp", false);
	}
    

    size_t der_size = br_encode_rsa_raw_der(NULL, sk, pk, privexp, privexp_size);
    if (der_size == 0) {
		Return RESULT_ERROR("br_encode_rsa_raw_der", false);
    }
    void* der = malloc(der_size);
    bool success = false;
    Defer(
        if(!success)
            free(der)
    );
    if (br_encode_rsa_raw_der(der, sk, pk, privexp, privexp_size) != der_size) {
		Return RESULT_ERROR("br_encode_rsa_raw_der", false);
    }

    success = true;
    out_der->data = der;
    out_der->size = der_size;
    Return RESULT_VOID;
}

void PEM_encode(Array(u8) src, Array(u8)* dst, cstr label){
    u64 encoded_size = br_pem_encode(NULL, src.data, src.size, label, 0);
    // br_pem_encode doesn't count '\0' but writes it
    *dst = Array_alloc_size(encoded_size + 1);
    br_pem_encode(dst->data, src.data, src.size, label, 0);
}


Result(void) RSA_parsePrivateKey_DER(Array(u8) _src, br_rsa_private_key* sk){
    Deferral(16);
    // private key data will be written in this buffer on success
    Array(u8) buf = Array_copy(_src);
    bool success = false;
    Defer(free(buf.data));

    br_skey_decoder_context decoder;
    br_skey_decoder_init(&decoder);
    br_skey_decoder_push(&decoder, buf.data, buf.size);
    i32 errcode = br_skey_decoder_last_error(&decoder);
	if (errcode != 0) {
        Return RESULT_ERROR_FMT("br_skey_decoder error %i", errcode);
	}
    i32 parsed_type = br_skey_decoder_key_type(&decoder);
    if(parsed_type != BR_KEYTYPE_RSA){
        Return RESULT_ERROR_FMT("parsed key has unsupported type %i", parsed_type);
    }
    const br_rsa_private_key* decoded_key = br_skey_decoder_get_rsa(&decoder);
    if(decoded_key == NULL){
        Return RESULT_ERROR("decoder failed without errors", false);
    }

    success = true;
    memcpy(sk, decoded_key, sizeof(*decoded_key));
    
    // sk fields still point to stack array decoder.key_data
    // This code copies the data and adjusts sk fields to point to the copied chunk
    Array(u8) key_data_copy = Array_alloc_size(BR_RSA_KBUF_PRIV_SIZE(decoded_key->n_bitlen));
    Defer(
        if(!success)
            free(key_data_copy.data);
    );
    memcpy(key_data_copy.data, decoder.key_data, key_data_copy.size);
    u64 memory_distance =  (u64)key_data_copy.data - (u64)(void*)decoder.key_data;
    u8** sk_pointer_fields[] = {
        &sk->p, &sk->q, &sk->dp, &sk->dq, &sk->iq
    };
    for(u32 i = 0; i < ARRAY_LEN(sk_pointer_fields); i++){
        u8** field_place = sk_pointer_fields[i];
        u64 field_value = (u64)*field_place;
        field_value += memory_distance;
        *field_place = (void*)field_value;
    }

    Return RESULT_VOID;
}

void EncryptorRSA_construct(EncryptorRSA* ptr, const br_rsa_public_key* pk){
    ptr->pk = pk;
    ptr->rng.vtable = &br_hmac_drbg_vtable;
    rng_init_sha256_seedFromTime(&ptr->rng.vtable);
}

void EncryptorRSA_encrypt(EncryptorRSA* ptr, Array(u8) src, Array(u8) dst, u32* encrypted_size){
    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);
    *encrypted_size = sz;
}


void DecryptorRSA_construct(DecryptorRSA* ptr, const br_rsa_private_key* sk){
    ptr->sk = sk;
}

void DecryptorRSA_decrypt(DecryptorRSA* ptr, Array(u8) buf, u32* decrypted_size){
    size_t sz = buf.size;
    br_rsa_i31_oaep_decrypt(
        &br_sha256_vtable, 
        NULL, 0,
        ptr->sk,
        buf.data, &sz);
    *decrypted_size = sz;
}