/* Motenpoche 
 *
 * (c) 2023 Daniele Lacamera <root@danielinux.net>
 *
 *
 * Motenpoche is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * Motenpoche is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
 *
 */
#include "user_settings.h"
#include <stdint.h>
#include "cryptoengine.h"
#include "fsm.h"
#include "flash.h"
#include "class/cdc/cdc.h"
#include "class/cdc/cdc_device.h"
#include "bsp/board.h"
#include "wolfssl/wolfcrypt/sha512.h"
#include "wolfssl/wolfcrypt/pwdbased.h"
#include "wolfssl/wolfcrypt/chacha.h"
#include "wolfssl/wolfcrypt/ecc.h"
#include "hardware/gpio.h"

static struct vault_header hdr_cache;

uint32_t sector_cache_id = 0xFFFFFFFF;
static ChaCha cha;
static ecc_key ecc;

static struct vault_service svc_cache;
static int svc_cache_id = -1;

int flash_decrypt_read_sector(uint32_t sector, uint8_t *buf)
{
    if (cryptoengine_check_vault() < 0)
        return -1;
    if (fsm_get() < VAULT_MAIN_MENU)
        return -1;
    if ((sector * SPI_FLASH_SECTOR_SIZE) >= DRIVE_FLASH_SIZE)
        return -1;
    flash_read(DRIVE_FLASH_OFFSET + sector * SPI_FLASH_SECTOR_SIZE, buf,
            SPI_FLASH_SECTOR_SIZE);
    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, sector);
    wc_Chacha_Process(&cha, buf, buf, SPI_FLASH_SECTOR_SIZE);
    return 0;
}

int flash_encrypt_write_sector(uint32_t sector, uint8_t *buf)
{
    if (cryptoengine_check_vault() < 0)
        return -1;
    if (fsm_get() < VAULT_MAIN_MENU)
        return -1;
    if ((sector * SPI_FLASH_SECTOR_SIZE) >= DRIVE_FLASH_SIZE)
        return -1;
    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, sector);
    wc_Chacha_Process(&cha, buf, buf, SPI_FLASH_SECTOR_SIZE);
    flash_write(DRIVE_FLASH_OFFSET + sector * SPI_FLASH_SECTOR_SIZE, buf,
            SPI_FLASH_SECTOR_SIZE);
    return 0;
}

int flash_decrypt_read_svc(struct vault_service *out, uint32_t addr)
{
    if (cryptoengine_check_vault() < 0)
        return -1;
    flash_read(addr, &svc_cache, SVC_SIZE);
    memcpy(out, (uint8_t *)(&svc_cache), SVC_ENC_OFF);
    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, addr);
    wc_Chacha_Process(&cha, (uint8_t *)out + SVC_ENC_OFF, ((uint8_t *)&svc_cache) + SVC_ENC_OFF, SVC_ENC_SIZE); 
    return 0;
}

static int flash_encrypt_write_svc(const struct vault_service *svc, uint32_t idx)
{
    uint32_t addr = idx * SVC_SIZE + SVC_FLASH_OFFSET;
    if (cryptoengine_check_vault() < 0)
        return -1;
    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, addr);
    wc_Chacha_Process(&cha, (uint8_t *)&svc_cache, (uint8_t *)svc, SVC_SIZE); 
    flash_write_svc(&svc_cache, idx);
}

int cryptoengine_verify_passphrase(const char *passphrase, int *res)
{
    int ret = -1;
    int i;
    uint8_t SecretKey[CRYPTO_KEY_SIZE];
    uint8_t DecryptedSignature[PK_SIGNATURE_SIZE];
    mp_int r, s;
    mp_init(&r);
    mp_init(&s);
    *res = 0;

    if (cryptoengine_check_vault() < 0) {
        return -1;
    }

    ret = wc_PBKDF2(SecretKey, passphrase, strlen(passphrase), hdr_cache.host_salt, SALT_LEN, 1000, 32, SHA512);
    if (ret < 0)
        return ret;
    wc_Chacha_SetKey(&cha, SecretKey, CRYPTO_KEY_SIZE);
    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, 0);
    wc_Chacha_Process(&cha, DecryptedSignature, hdr_cache.signature, PK_SIGNATURE_SIZE); 

    wc_ecc_init(&ecc);
    if (wc_ecc_import_unsigned(&ecc, hdr_cache.auth_pubkey,
            hdr_cache.auth_pubkey + CRYPTO_KEY_SIZE,
            NULL, ECC_SECP256R1) < 0) {
        mp_free(&r);
        mp_free(&s);
        wc_ecc_free(&ecc);
        return -1;
    }
    mp_read_unsigned_bin(&r, DecryptedSignature, CRYPTO_KEY_SIZE);
    mp_read_unsigned_bin(&s, DecryptedSignature + CRYPTO_KEY_SIZE,
            CRYPTO_KEY_SIZE);

    ret = wc_ecc_verify_hash_ex(&r, &s, hdr_cache.digest, VAULT_DIGEST_SIZE,
            res, &ecc);
    mp_free(&r);
    mp_free(&s);
    if ((ret < 0) || (*res != 1)) {
        if (*res == 0)
            gpio_put(16, 0);

        wc_ecc_free(&ecc);
    }
    return ret;
}


static void hdr_cache_load(void)
{
    flash_read(VAULT_FLASH_OFFSET, &hdr_cache, sizeof(struct vault_header));
}


uint8_t checked_digest[VAULT_DIGEST_SIZE];
int cryptoengine_hdr_sha_check(const struct vault_header *vh)
{
    Sha512 sha;
    int hash, hash_len;
    wc_InitSha512(&sha);
    for (hash = 0; hash < SHA_PAYLOAD_SIZE;) {
        hash_len = WC_SHA512_BLOCK_SIZE;
        if ((SHA_PAYLOAD_SIZE - hash) < hash_len)
            hash_len = SHA_PAYLOAD_SIZE - hash;
        wc_Sha512Update(&sha, ((uint8_t*)vh) + hash, hash_len);
        hash += hash_len;
    }
    wc_Sha512Final(&sha, checked_digest);
    if(memcmp(checked_digest, vh->digest, VAULT_DIGEST_SIZE) != 0) {
        return -1;
    }
    wc_Sha512Free(&sha);
    return 0;
}

int cryptoengine_svc_sha_check(const struct vault_service *vs)
{
    Sha512 sha;
    int hash, hash_len;
    const uint32_t svc_sha_payload_size = (sizeof(struct vault_service) - (VAULT_DIGEST_SIZE + PK_SIGNATURE_SIZE));
    wc_InitSha512(&sha);
    for (hash = 0; hash < svc_sha_payload_size;) {
        hash_len = WC_SHA512_BLOCK_SIZE;
        if ((svc_sha_payload_size - hash) < hash_len)
            hash_len = svc_sha_payload_size - hash;
        wc_Sha512Update(&sha, ((uint8_t*)vs) + hash, hash_len);
        hash += hash_len;
    }
    wc_Sha512Final(&sha, checked_digest);
    if(memcmp(checked_digest, vs->dig, VAULT_DIGEST_SIZE) != 0) {
        return -1;
    }
    wc_Sha512Free(&sha);
    return 0;
}

int cryptoengine_svc_sig_verify(const struct vault_service *vs, int *res)
{
    int ret;
    mp_int r, s;
    mp_init(&r);
    mp_init(&s);

    mp_read_unsigned_bin(&r, vs->sig, CRYPTO_KEY_SIZE);
    mp_read_unsigned_bin(&s, vs->sig + CRYPTO_KEY_SIZE, CRYPTO_KEY_SIZE);

    ret = wc_ecc_verify_hash_ex(&r, &s, vs->dig, VAULT_DIGEST_SIZE, res, &ecc);
    mp_free(&r);
    mp_free(&s);
    return ret;
}

int cryptoengine_check_vault(void)
{
    hdr_cache_load();
    if (hdr_cache.magic != VAULT_MAGIC) {
        return -1;
    }
    /* Checking hdr integrity */
    if (cryptoengine_hdr_sha_check(&hdr_cache) < 0)
        return -1;

    /* Vault seems valid, integrity check OK */
    return 0;
}

int cryptoengine_import_service(struct vault_service *vs, uint32_t idx)
{
    uint32_t address = SVC_FLASH_OFFSET + SVC_SIZE * idx;
    int res;
    if (cryptoengine_check_vault() < 0)
        return -1;
    memcpy(&svc_cache, ((uint8_t *)vs), SVC_ENC_OFF);
    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, address);
    wc_Chacha_Process(&cha, ((uint8_t *)&svc_cache) + SVC_ENC_OFF, 
            ((uint8_t *)vs) + SVC_ENC_OFF, SVC_ENC_SIZE);
    if (cryptoengine_svc_sha_check(&svc_cache) < 0) {
        return -1;
    }
    if ((cryptoengine_svc_sig_verify(&svc_cache, &res) < 0) || (res != 1))
    {
        return -2;
    }
    flash_write_svc(vs, idx);
    return 0;
}

int cryptoengine_service_count(uint16_t *n_srv, uint16_t *first_avail)
{
    int i = 0;
    *first_avail = (uint16_t)-1;
    *n_srv = 0;
    
    struct vault_service svc;
    if (cryptoengine_check_vault() < 0)
        return -1;

    while (1) {
        flash_read(SVC_FLASH_OFFSET + (i * SVC_SIZE), &svc, SVC_SIZE);
        if (svc.flags == SVC_FLAG_ACTIVE)
            (*n_srv)++;
        else if (svc.flags == SVC_FLAG_ERASED)
            *first_avail = i;
        else if (svc.flags == SVC_FLAG_UNUSED) {
            if (*first_avail == ((uint16_t)-1))
                *first_avail = i;
            break;
        }
        i++;
    }
    return *n_srv;
}
                
int cryptoengine_fill_vault_status(struct vault_status *vst)
{
    if (cryptoengine_check_vault() < 0)
        return -1;
    memcpy(vst->id, &hdr_cache.id, 2 * sizeof(uint32_t));
    memcpy(vst->salt, &hdr_cache.host_salt, SALT_LEN);
    memcpy(vst->seed, &hdr_cache.host_seed, SEED_LEN);
}