#include "crypto.h"
#include "system.h"
#include "constants.h"
#include "storage.h"
#include "dev_name.h"

#include <string.h>
#include <mbedtls/build_info.h>
#include <mbedtls/platform.h>
#include <mbedtls/entropy.h>
#include <mbedtls/ctr_drbg.h>
#include <mbedtls/ecdsa.h>
#include <mbedtls/sha256.h>
#include <mbedtls/pk.h>
#include <mbedtls/x509_csr.h>
#include <mbedtls/base64.h>
#include "esp_log.h"

#define ECPARAMS MBEDTLS_ECP_DP_SECP256R1

static const char *TAG = "crypto";

static const char *pers = "ecdsa";

static void seed_ctr_drbg_context(mbedtls_entropy_context *entropy, mbedtls_ctr_drbg_context *ctr_drbg)
{
    int ret;

    mbedtls_entropy_init(entropy);
    mbedtls_ctr_drbg_init(ctr_drbg);

    ESP_LOGI(TAG, "Seed Mbedtls\n");
    if ((ret = mbedtls_ctr_drbg_seed(ctr_drbg, mbedtls_entropy_func, entropy,
                                     (const unsigned char *)pers,
                                     strlen(pers))) != 0)
    {
        ESP_LOGE(TAG, " failed\n  ! mbedtls_ctr_drbg_seed returned %d\n", ret);
        reboot();
    }
}

bool crypto_gen_priv_key()
{
    // Check if a private key has already been defined for this device
    if (storage_get_priv_key(NULL) > 0)
        return false;

    int ret = 1;

    mbedtls_pk_context key;
    mbedtls_pk_init(&key);

    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context ctr_drbg;
    seed_ctr_drbg_context(&entropy, &ctr_drbg);

    ESP_LOGI(TAG, "PK info from type\n");
    if ((ret = mbedtls_pk_setup(&key, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY))) != 0)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_pk_setup returned -0x%04x", (unsigned int)-ret);
        reboot();
    }

    // Generate private key
    ESP_LOGI(TAG, "Generate private key\n");
    ret = mbedtls_ecp_gen_key(ECPARAMS,
                              mbedtls_pk_ec(key),
                              mbedtls_ctr_drbg_random, &ctr_drbg);
    if (ret != 0)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_ecp_gen_key returned -0x%04x",
                 (unsigned int)-ret);
        reboot();
    }

    // Export private key
    ESP_LOGI(TAG, "Export private key\n");
    unsigned char *key_buff = malloc(PRV_KEY_DER_MAX_BYTES);
    if ((ret = mbedtls_pk_write_key_der(&key, key_buff, PRV_KEY_DER_MAX_BYTES)) < 1)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_pk_write_key_der returned -0x%04x",
                 (unsigned int)-ret);
        reboot();
    }

    storage_set_priv_key(key_buff + PRV_KEY_DER_MAX_BYTES - ret, ret);
    free(key_buff);

    mbedtls_pk_free(&key);
    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);

    return true;
}

void crypto_print_priv_key()
{
    int ret;

    unsigned char *key_buff = malloc(PRV_KEY_DER_MAX_BYTES);
    size_t key_len = storage_get_priv_key(key_buff);
    assert(key_len > 0);

    mbedtls_pk_context key;
    mbedtls_pk_init(&key);

    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context ctr_drbg;
    seed_ctr_drbg_context(&entropy, &ctr_drbg);

    ESP_LOGI(TAG, "Parse private key (len = %d)\n", key_len);
    if ((ret = mbedtls_pk_parse_key(&key, key_buff, key_len, NULL, 0, mbedtls_ctr_drbg_random, &ctr_drbg)) != 0)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_pk_parse_key returned -0x%04x",
                 (unsigned int)-ret);
        reboot();
    }
    free(key_buff);

    ESP_LOGI(TAG, "Show private key\n");
    unsigned char *out = malloc(16000);
    memset(out, 0, 16000);
    if ((ret = mbedtls_pk_write_key_pem(&key, out, 16000)) != 0)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_pk_write_key_pem returned -0x%04x",
                 (unsigned int)-ret);
        reboot();
    }

    ESP_LOGI(TAG, "%s", out);
    free(out);

    mbedtls_pk_free(&key);
    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);
}

char *crypto_get_csr()
{
    int ret;

    unsigned char *key_buff = malloc(PRV_KEY_DER_MAX_BYTES);
    size_t key_len = storage_get_priv_key(key_buff);
    assert(key_len > 0);

    mbedtls_pk_context key;
    mbedtls_pk_init(&key);

    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context ctr_drbg;
    seed_ctr_drbg_context(&entropy, &ctr_drbg);

    ESP_LOGI(TAG, "Parse private key (len = %d)\n", key_len);
    if ((ret = mbedtls_pk_parse_key(&key, key_buff, key_len, NULL, 0, mbedtls_ctr_drbg_random, &ctr_drbg)) != 0)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_pk_parse_key returned -0x%04x",
                 (unsigned int)-ret);
        reboot();
    }
    free(key_buff);

    // Create CSR
    mbedtls_x509write_csr req;
    mbedtls_x509write_csr_init(&req);
    mbedtls_x509write_csr_set_md_alg(&req, MBEDTLS_MD_SHA256);

    char subj[DEV_NAME_LEN + 4];
    char *n = dev_name();
    sprintf(subj, "CN=%s", n);
    free(n);
    if ((ret = mbedtls_x509write_csr_set_subject_name(&req, subj)) != 0)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_x509write_csr_set_subject_name returned %d", ret);
        reboot();
    }

    ESP_LOGI(TAG, "Sign CSR with private key\n");
    mbedtls_x509write_csr_set_key(&req, &key);

    char *csr = malloc(4096);
    if ((ret = mbedtls_x509write_csr_pem(&req, (u_char *)csr, 4096, mbedtls_ctr_drbg_random, &ctr_drbg)) < 0)
    {
        ESP_LOGE(TAG, " failed\n  !  mbedtls_x509write_csr_pem returned %d", ret);
        reboot();
    }

    mbedtls_x509write_csr_free(&req);
    mbedtls_pk_free(&key);
    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);

    return csr;
}

char *crypto_encode_base64_safe_url(const char *src, size_t srclen)
{
    size_t olen = 0;
    mbedtls_base64_encode(NULL, 0, &olen, (unsigned char *)src, srclen);

    if (olen < 1)
    {
        ESP_LOGE(TAG, "Failed to determine base64 buffer size!");
        return NULL;
    }

    char *out = calloc(1, olen);
    if (!out)
    {
        ESP_LOGE(TAG, "Failed to allocate memory for destination buffer!");
        return NULL;
    }

    if (mbedtls_base64_encode((unsigned char *)out, olen, &olen, (unsigned char *)src, srclen) != 0)
    {
        ESP_LOGE(TAG, "Failed to perfom base64 encoding!");
        free(out);
        return NULL;
    }

    if (out[olen - 1] == '=')
        out[olen - 1] = '\0';

    if (out[olen - 2] == '=')
        out[olen - 2] = '\0';

    return out;
}