#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <stdarg.h>
#include "mbedtls/x509_crt.h"
#include "mbedtls/platform_time.h"
#include "mbedtls/gcm.h"
#include "mbedtls/ssl_ciphersuites.h"
#include "talkingq_wss.h"

void *mbedtls_calloc(size_t n, size_t size) { return calloc(n, size); }
void mbedtls_free(void *p) { free(p); }
static time_t test_time;
static time_t read_test_time(time_t *p) { if (p) *p = test_time; return test_time; }
mbedtls_time_t (*mbedtls_time)(mbedtls_time_t *) = read_test_time;
void CPUartLogPrintf(const char *format, ...) { (void)format; }
int mbedtls_snprintf(char *buffer, size_t size, const char *format, ...)
{
    int result;
    va_list args;
    va_start(args, format);
    result = vsnprintf(buffer, size, format, args);
    va_end(args);
    return result;
}

int main(int argc, char **argv)
{
    mbedtls_x509_crt root, chain;
    mbedtls_gcm_context gcm;
    uint32_t flags;
    int ret, i;
    char path[1024];
    char pem[16384];
    FILE *file;
    size_t length;
    clock_t verify_started;
    mbedtls_ecp_group group;
    unsigned char key[16] = {0}, iv[12] = {0}, plain[16] = {0}, cipher[16], tag[16];
    const unsigned char expected_cipher[16] = {0x03,0x88,0xda,0xce,0x60,0xb6,0xa3,0x92,0xf3,0x28,0xc2,0xb9,0x71,0xb2,0xfe,0x78};
    const unsigned char expected_tag[16] = {0xab,0x6e,0x47,0xd4,0x2c,0xec,0x13,0xbd,0xf5,0x3a,0x67,0xb2,0x12,0x57,0xbd,0xdf};
    assert(argc == 2);
    test_time = time(NULL);
    mbedtls_x509_crt_init(&root);
    mbedtls_x509_crt_init(&chain);
    ret = mbedtls_x509_crt_parse(&root, (const unsigned char *)talkingq_root_ca, talkingq_root_ca_size);
    printf("SDK root CA parse: %d (previous device failure: -9774)\n", ret);
    assert(ret == 0);
    for (i = 0; i < 3; ++i) {
        snprintf(path, sizeof(path), "%s/chain-%d.pem", argv[1], i);
        file = fopen(path, "rb");
        assert(file != NULL);
        length = fread(pem, 1, sizeof(pem)-1, file);
        assert(feof(file) && !ferror(file));
        fclose(file);
        pem[length] = 0;
        ret = mbedtls_x509_crt_parse(&chain, (const unsigned char *)pem, length + 1);
        assert(ret == 0);
    }
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    printf("SDK server chain + hostname verify: %d, flags=0x%08x\n", ret, (unsigned int)flags);
    assert(ret == 0 && flags == 0);
    verify_started = clock();
    for (i = 0; i < 20; ++i) {
        ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
        assert(ret == 0 && flags == 0);
    }
    printf("Host certificate chain verify average: %.2f ms (20 runs)\n",
           1000.0 * (clock() - verify_started) / CLOCKS_PER_SEC / 20);
    mbedtls_ecp_group_init(&group);
    assert(mbedtls_ecp_group_load(&group, MBEDTLS_ECP_DP_SECP384R1) == 0);
#ifdef MBEDTLS_ECP_NIST_OPTIM
    assert(group.modp != NULL);
#else
    assert(group.modp == NULL);
#endif
    mbedtls_ecp_group_free(&group);
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "wrong.invalid", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_CN_MISMATCH));
    ret = mbedtls_x509_crt_verify(&chain, NULL, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_NOT_TRUSTED));
    test_time = 1900000000; /* 2030: this server certificate is expired. */
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_EXPIRED));
    test_time = 1700000000; /* 2023: this server certificate is not yet valid. */
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_FUTURE));
    assert(mbedtls_ssl_ciphersuite_from_id(MBEDTLS_TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384) != NULL);
    mbedtls_gcm_init(&gcm);
    assert(mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, key, 128) == 0);
    assert(mbedtls_gcm_crypt_and_tag(&gcm, MBEDTLS_GCM_ENCRYPT, 16, iv, 12, NULL, 0,
                                   plain, cipher, 16, tag) == 0);
    assert(!memcmp(cipher, expected_cipher, 16) && !memcmp(tag, expected_tag, 16));
    mbedtls_gcm_free(&gcm);
    mbedtls_x509_crt_free(&chain);
    mbedtls_x509_crt_free(&root);
    puts("PASS: SDK ECDSA certificate chain, hostname, trust, dates, TLS suite and AES-GCM vector");
    return 0;
}