// ________________________________________________
//
//  GenerateCert
//
//  PURPOSE: 
//  Generate certificate with public & private key.
//
//  PARAMETERS:
//  None
//
//  RETURN VALUE:
//  SSL_SUCCESS If successful the call will return.
//  BAD_MUTEX_E is an error that may be returned.
//  WC_INIT_E wolfCrypt initialization error returned.
// ________________________________________________
//
Export BOOL GenerateCert()
{
    /* Generate RSA Key */

    RsaKey    genKey;
    RNG        rng;
    int     ret;
    BOOL    bRet = FALSE;
    byte    derKey[4096];
    byte    pemKey[4096];
    int        derSize;
    int        pemSize;
    HANDLE hFile;
    DWORD    dwBytesWritten;

    wc_InitRng(&rng);
    wc_InitRsaKey(&genKey, NULL);
    ret = wc_MakeRsaKey(&genKey, 2048, 65537, &rng);

    if (ret != 0) {
        /* ret contains error */;
        return FALSE;
    }

    derSize = wc_RsaKeyToDer(&genKey, derKey, sizeof(derKey));
    if (derSize < 0) {
        /* derSz contains error */;
        MessageBoxA(0, "ERROR: failed to get DER RSA Key", "", MB_ICONERROR);
        return FALSE;
    }

    pemSize = wc_DerToPem(derKey, derSize, pemKey, sizeof(pemKey), PRIVATEKEY_TYPE);
    if (pemSize < 0) {
        /* pemCertSz contains error */;
        MessageBoxA(0, "ERROR: failed to get PEM RSA Key", "", MB_ICONERROR);
        return FALSE;
    }


    /* Write Key file  */

    hFile = CreateFile(L"key.key",                // name of the write
        GENERIC_WRITE,          // open for writing
        0,                      // do not share
        NULL,                   // default security
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,  // normal file
        NULL);

    bRet = WriteFile(
        hFile,           // open file handle
        pemKey,      // start of data to write
        pemSize,  // number of bytes to write
        &dwBytesWritten, // number of bytes that were written
        NULL);            // no overlapped structure

    CloseHandle(hFile);


    /* Generate Certificate */

    Cert myCert;
    byte derCert[4096];
    byte pemCert[4096];

    wc_InitCert(&myCert);

    derSize = wc_MakeSelfCert(&myCert, derCert, sizeof(derCert), &genKey, &rng);
    if (derSize < 0) {
        /* certSz contains the error */;
        return FALSE;
    }

    pemSize = wc_DerToPem(derCert, derSize, pemCert, sizeof(pemCert), CERT_TYPE);
    if (pemSize < 0) {
        /* pemCertSz contains error */;
        return FALSE;
    }

    /* Write Certificate file  */

    hFile = CreateFile(L"certificate.crt",                // name of the write
        GENERIC_WRITE,          // open for writing
        0,                      // do not share
        NULL,                   // default security
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,  // normal file
        NULL);

    bRet = WriteFile(
        hFile,           // open file handle
        pemCert,      // start of data to write
        pemSize,  // number of bytes to write
        &dwBytesWritten, // number of bytes that were written
        NULL);            // no overlapped structure

    CloseHandle(hFile);

    return bRet;
}

// ________________________________________________
//
//    InitTLS_Server
//
//  PURPOSE: 
//  Initialize TLS Server settings:
//  + Select Cipher Suite
//  + Enable Client Authentication
//    + Load RSA Key
//
//  PARAMETERS:
//  None
//
//  RETURN VALUE:
//  Pointer to WOLFSSL_CTX on success
//  NULL on failure
// ________________________________________________
//
Export WOLFSSL_CTX* InitTLS_Server()
{
    WOLFSSL_CTX*    ctx;

    /* Create and initialize WOLFSSL_CTX */
    if ((ctx = wolfSSL_CTX_new(wolfTLSv1_3_server_method())) == NULL) {
        MessageBoxA(0, "ERROR: failed to create WOLFSSL_CTX", "", MB_ICONERROR);
        //fprintf(stderr, "ERROR: failed to create WOLFSSL_CTX\n");
        return NULL;
    }

    /* Select cipher list to use */
    if (wolfSSL_CTX_set_cipher_list(ctx, "TLS_AES_128_GCM_SHA256") != SSL_SUCCESS)
    {
        MessageBoxA(0, "ERROR: failed to set cipher list.", "", MB_ICONERROR);
        wolfSSL_CTX_free(ctx);  /* Free the wolfSSL context object   */
        return NULL;
    }

    /* Load CA certificates into WOLFSSL_CTX */
    if (wolfSSL_CTX_load_verify_locations(ctx, CERT_FILE, NULL) != SSL_SUCCESS)
    {
        MessageBoxA(0, "ERROR: failed to load certificate.", "", MB_ICONERROR);
        wolfSSL_CTX_free(ctx);  // Free the wolfSSL context object
        return NULL;
    }

    /* Load server certificates into WOLFSSL_CTX */
    if (wolfSSL_CTX_use_certificate_file(ctx, CERT_FILE, SSL_FILETYPE_PEM) != SSL_SUCCESS)
    {
        MessageBoxA(0, "ERROR: failed to load server certificate.", "", MB_ICONERROR);
        wolfSSL_CTX_free(ctx);  /* Free the wolfSSL context object   */
        return NULL;
    }

    /* Load server key into WOLFSSL_CTX */
    if (wolfSSL_CTX_use_PrivateKey_file(ctx, KEY_FILE, SSL_FILETYPE_PEM)
        != SSL_SUCCESS) {
        //fprintf(stderr, "ERROR: failed to load %s, please check the file.\n", KEY_FILE);

        MessageBoxA(0, "ERROR: failed to load keyfile.", "", MB_ICONERROR);
        wolfSSL_CTX_free(ctx);
        return NULL;
    }

    /* Use trusted certificate to verify a peer when performing a TLS handshake (instead of CA certificate). */
    if (wolfSSL_CTX_trust_peer_cert(ctx, CERT_FILE, SSL_FILETYPE_PEM) != SSL_SUCCESS)
    {
        MessageBoxA(0, "ERROR: failed to load trusted peer certificate.", "", MB_ICONERROR);
        wolfSSL_CTX_free(ctx);  /* Free the wolfSSL context object   */
        return NULL;
    }

    /* Enable client authentication */
    wolfSSL_CTX_set_verify(ctx, SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, 0);

    return ctx;
}

// ________________________________________________
//
//  InitializeTLS (Client Side)
//
//  PURPOSE: 
//  One-time initialization of TLS.
//  Must be called just once at application start.
//
//  PARAMETERS:
//  Client Certificate buffer
//
//  RETURN VALUE:
//  TRUE on Success
//  FALSE on Failure
// ________________________________________________
//
bool CConnection::InitializeTLS(PBYTE pCert, int lenCert)
{
    int ret;

    /* Initialize wolfSSL */
    if ((ret = wolfSSL_Init()) != WOLFSSL_SUCCESS) {
        PrintEvent("[ERROR]", "Failed to initialize TLS");
        return false;
    }

    /* Create and initialize WOLFSSL_CTX */
    if ((ctxTLS = wolfSSL_CTX_new(wolfTLSv1_3_client_method())) == NULL) {
        PrintEvent("[ERROR]", "Failed to initialize TLS context");
        ret = -1;
        return false;
    }

    /* Load client certificates into WOLFSSL_CTX */
    if ((ret = wolfSSL_CTX_use_certificate_buffer(ctxTLS, pCert, lenCert, SSL_FILETYPE_ASN1))
        != SSL_SUCCESS) {
        PrintEvent("[ERROR]", "Failed to load TLS certificate");
        return false;
    }

    /* Select cipher list to use */
    if (wolfSSL_CTX_set_cipher_list(ctxTLS, "TLS_AES_128_GCM_SHA256") != SSL_SUCCESS)
    {
        PrintEvent("[ERROR]", "Failed to set cipher list");
        return false;
    }

    /* Use trusted certificate to verify a peer when performing a TLS handshake (instead of CA certificate). */
    if (wolfSSL_CTX_trust_peer_buffer(ctxTLS, pCert, lenCert, SSL_FILETYPE_ASN1) != SSL_SUCCESS)
    {
        PrintEvent("[ERROR]", "Failed to set peer certificate");
        return false;
    }

    bUseTLS = true;
    return true;
}