/*
 * Copyright (c) 2020-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * This software is available to you under a choice of one of two
 * licenses.  You may choose to be licensed under the terms of the GNU
 * General Public License (GPL) Version 2, available from the file
 * COPYING in the main directory of this source tree, or the
 * OpenIB.org BSD license below:
 *
 *     Redistribution and use in source and binary forms, with or
 *     without modification, are permitted provided that the following
 *     conditions are met:
 *
 *      - Redistributions of source code must retain the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer.
 *
 *      - Redistributions in binary form must reproduce the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer in the documentation and/or other materials
 *        provided with the distribution.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
 * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
 * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
#ifndef __WIN__
#include "hsmlunaclient.h"
#include <string>
using namespace std;
static unsigned char PublicExponent[] = {0x01, 0x00, 0x01};

CK_RV HSMLunaClient::CreatePublicKey(string publicKeyLabel, unsigned char* PublicModulus, unsigned int PublicModulusSize)
{
    CK_OBJECT_CLASS publicKey = CKO_PUBLIC_KEY;
    CK_KEY_TYPE rsaType = CKK_RSA;
    CK_BBOOL bTrue = 1;
    CK_BBOOL bSign = bTrue, bToken = bTrue, bPrivate = bTrue, bDecrypt = bTrue, bWrap = bTrue;
    // CK_BBOOL bToken = bTrue, bDecrypt = bTrue, bSign = bTrue, bWrap = bTrue;
    CK_OBJECT_HANDLE hPublicRSAKey = 0;
    CK_RV rv = CKR_TOKEN_NOT_PRESENT;
    unsigned int publicLabelSize = publicKeyLabel.size() + 1;
    CK_BYTE* pLabelPub = new CK_BYTE[publicLabelSize];
    memset(pLabelPub, 0, publicLabelSize);

    for (unsigned int i = 0; i < publicKeyLabel.size(); i++)
    {
        pLabelPub[i] = publicKeyLabel[i];
    }
    CK_ATTRIBUTE pPublicRSAKeyTemplate[] = {{CKA_CLASS, &publicKey, sizeof(publicKey)},
                                            {CKA_KEY_TYPE, &rsaType, sizeof(rsaType)},
                                            {CKA_TOKEN, &bToken, sizeof(bToken)},
                                            {CKA_PRIVATE, &bPrivate, sizeof(bPrivate)},
                                            {CKA_ENCRYPT, &bDecrypt, sizeof(bDecrypt)},
                                            {CKA_VERIFY, &bSign, sizeof(bSign)},
                                            {CKA_WRAP, &bWrap, sizeof(bWrap)},
                                            {CKA_MODULUS, PublicModulus, PublicModulusSize},
                                            {CKA_PUBLIC_EXPONENT, PublicExponent, sizeof(PublicExponent)},
                                            {CKA_LABEL, pLabelPub, publicLabelSize}};
    unsigned int usTemplateSize = DIM(pPublicRSAKeyTemplate);
    rv = P11Functions->C_CreateObject(m_hSession, pPublicRSAKeyTemplate, usTemplateSize, &hPublicRSAKey);
    if (rv != CKR_OK)
    {
        cout << "C_Create public Object FAILED...rv = " << hex << rv << endl;
        return -1;
    }
    cout << "Create public key success handle = " << hex << hPublicRSAKey << endl;
    delete[] pLabelPub;
    return CKR_OK;
}

CK_RV HSMLunaClient::CheckExistingLabel(string label, CK_ULONG& numOfLabels)
{
    CK_RV rv = CKR_TOKEN_NOT_PRESENT;
    unsigned int LabelSize = label.size() + 1;
    CK_BYTE* pLabel = new CK_BYTE[LabelSize];
    memset(pLabel, 0, LabelSize);
    for (unsigned int i = 0; i < label.size(); i++)
    {
        pLabel[i] = label[i];
    }
    CK_ATTRIBUTE Template[] = {{CKA_LABEL, pLabel, LabelSize}};

    CK_ULONG numHandles = 0;
    CK_OBJECT_HANDLE handles[MAX_NUM_OF_HANDLES] = {0};
    rv = P11Functions->C_FindObjectsInit(m_hSession, Template, 1);
    if (rv != CKR_OK)
    {
        delete[] pLabel;
        return rv;
    }
    rv = P11Functions->C_FindObjects(m_hSession, handles, MAX_NUM_OF_HANDLES, &numHandles);
    if (rv != CKR_OK)
    {
        delete[] pLabel;
        return rv;
    }
    numOfLabels = numHandles;
    if (numHandles != 0)
    {
        delete[] pLabel;
        return CKR_GENERAL_ERROR;
    }
    if (numHandles == 0)
    {
        Template[0].ulValueLen--;
        rv = P11Functions->C_FindObjectsInit(m_hSession, Template, 1);
        if (rv != CKR_OK)
        {
            delete[] pLabel;
            return rv;
        }
        rv = P11Functions->C_FindObjects(m_hSession, handles, MAX_NUM_OF_HANDLES, &numHandles);
        if (rv != CKR_OK)
        {
            delete[] pLabel;
            return rv;
        }
    }
    numOfLabels = numHandles;
    delete[] pLabel;
    if (numHandles != 0)
    {
        return CKR_GENERAL_ERROR;
    }
    return CKR_OK;
}

CK_RV HSMLunaClient::BurnPrivateKey(string privateKeyLabel, string publicKeyLabel, vector<unsigned char> PrivateKeyData)
{
    protectedInfo_T protectedInfo;
    CK_RV rv = CKR_TOKEN_NOT_PRESENT;
    CK_OBJECT_HANDLE hPrivateRSAKey = 0;
    char* pbWrappedKey;
    unsigned long ulWrappedKeySize;
    CK_OBJECT_CLASS privateKey = CKO_PRIVATE_KEY;

    CK_USHORT usOutDataLen = 0;
    char* pPlainData = 0;
    unsigned long ulPlainDataLength;
    char* pEncryptedData = 0;
    unsigned long ulEncryptedDataLength = 0;

    CK_MECHANISM mech;

    // CK_BYTE pLabel[] = "Super New RSA private Key222";
    CK_KEY_TYPE rsaType = CKK_RSA;
    CK_ULONG usTemplateSize;

    rv = P11Functions->C_GetInfo(&protectedInfo.info);
    if (rv != CKR_OK)
    {
        cout << "C_GetInfo FAILED..." << endl;
        return -1;
    }
    CK_BBOOL bFalse = 0;
    CK_BBOOL bTrue = 1;
    CK_BBOOL bToken = bTrue;
    CK_BBOOL bPrivate = bTrue;
    CK_BBOOL bEncrypt = bTrue;
    CK_BBOOL bSign = bTrue;
    CK_BBOOL bWrap = bTrue;

    CK_BBOOL bSensitive = bTrue;
    unsigned int privateLabelSize = privateKeyLabel.size() + 1;
    CK_BYTE* pLabel = new CK_BYTE[privateLabelSize];
    memset(pLabel, 0, privateLabelSize);

    unsigned int publicLabelSize = publicKeyLabel.size() + 1;
    CK_BYTE* pLabelPub = new CK_BYTE[publicLabelSize];
    memset(pLabelPub, 0, publicLabelSize);

    for (unsigned int i = 0; i < privateKeyLabel.size(); i++)
    {
        pLabel[i] = privateKeyLabel[i];
    }
    for (unsigned int i = 0; i < publicKeyLabel.size(); i++)
    {
        pLabelPub[i] = publicKeyLabel[i];
    }

    CK_ATTRIBUTE pPrivateKeyTemplate[] = {{CKA_CLASS, &privateKey, sizeof(privateKey)},
                                          {CKA_KEY_TYPE, &rsaType, sizeof(rsaType)},
                                          {CKA_TOKEN, &bToken, sizeof(bToken)},
                                          {CKA_SENSITIVE, &bSensitive, sizeof(bSensitive)},
                                          {CKA_PRIVATE, &bPrivate, sizeof(bPrivate)},
                                          {CKA_DECRYPT, &bEncrypt, sizeof(bEncrypt)},
                                          {CKA_SIGN, &bSign, sizeof(bSign)},
                                          {CKA_SIGN_RECOVER, &bTrue, sizeof(bTrue)}, // here may be remark!
                                          {CKA_UNWRAP, &bWrap, sizeof(bWrap)},
                                          {CKA_EXTRACTABLE, &bFalse, sizeof(bFalse)},
                                          {CKA_LABEL, pLabel, privateLabelSize}};

    rv = Generate3DESKey();
    if (rv != CKR_OK)
    {
        cout << "\n"
                "Error 0x"
             << hex << rv;
        cout << " generating the DES3 Key.\n";
        delete[] pLabelPub;
        return -11;
    }

    // Encrypt the RSA Key
    unsigned char iv[8] = {'1', '2', '3', '4', '5', '6', '7', '8'};
    mech.mechanism = CKM_DES3_CBC;
    mech.pParameter = iv;
    mech.usParameterLen = sizeof(iv);
    unsigned char* _PrivateKeyData = &PrivateKeyData[0];
    PrivateKeyDataSize = PrivateKeyData.size();
    pPlainData = (char*)(_PrivateKeyData);
    ulPlainDataLength = PrivateKeyDataSize;
    pEncryptedData = new char[ulPlainDataLength + 2048];
    rv = P11Functions->C_EncryptInit(m_hSession, &mech, hSymKey);
    CK_USHORT usInDataLen;
    usOutDataLen = (CK_USHORT)(ulPlainDataLength + 2048);
    CK_ULONG ulBytesRemaining = ulPlainDataLength;
    char* pPlainTextPointer = pPlainData;
    char* pEncryptedDataPointer = pEncryptedData;
    while (ulBytesRemaining > 0)
    {
        if (ulBytesRemaining > 0xfff0)
        {
            usInDataLen = 0xfff0;
            ulBytesRemaining -= usInDataLen;
        }
        else
        {
            usInDataLen = (CK_USHORT)ulBytesRemaining;
            ulBytesRemaining -= usInDataLen;
        }
        rv = P11Functions->C_EncryptUpdate(m_hSession, (CK_BYTE_PTR)pPlainTextPointer, usInDataLen,
                                           (CK_BYTE_PTR)pEncryptedDataPointer, &usOutDataLen);
        if (rv == CKR_OK)
        {
            cout << "Encrypting key..." << endl;
        }
        else
        {
            cout << "\n"
                    "Error 0x"
                 << hex << rv;
            cout << " while encrypting the RSA key.\n";
            return -1;
        }
        pPlainTextPointer += usInDataLen;
        pEncryptedDataPointer += usOutDataLen;
        ulEncryptedDataLength += usOutDataLen;
    }

    usOutDataLen = 0;
    CK_BYTE_PTR pOutData = (CK_BYTE_PTR)pEncryptedData;
    pOutData += ulEncryptedDataLength;
    rv = P11Functions->C_EncryptFinal(m_hSession, pOutData, &usOutDataLen);
    if (rv != CKR_OK)
    {
        cout << "C_EncryptFinal FAILED..." << endl;
        return -1;
    }
    ulEncryptedDataLength += usOutDataLen;

    mech.mechanism = CKM_DES3_CBC;
    mech.pParameter = (void*)"12345678"; // 8 byte IV
    mech.usParameterLen = 8;

    usTemplateSize = DIM(pPrivateKeyTemplate);
    pbWrappedKey = pEncryptedData;
    ulWrappedKeySize = ulEncryptedDataLength;
    rv = P11Functions->C_UnwrapKey(m_hSession, &mech, hSymKey, (CK_BYTE_PTR)pbWrappedKey, (CK_USHORT)ulWrappedKeySize,
                                   pPrivateKeyTemplate, usTemplateSize, &hPrivateRSAKey);
    if (rv != CKR_OK)
    {
        cout << "C_UnwrapKey FAILED...rv = " << hex << rv << endl;
        delete[] pEncryptedData;
        return -1;
    }
    else
    {
        cout << "Unwrapped Private key is " << hex << hPrivateRSAKey << endl;
    }
    if (pEncryptedData)
    {
        delete[] pEncryptedData;
    }
    /*rv = P11Functions->C_CreateObject(m_hSession, pPrivateKeyTemplate,
        usTemplateSize, &hPrivateRSAKey);
    if (rv != CKR_OK) {
        cout << "C_Create private Object FAILED...rv = " << hex << rv << endl;
        return -1;
    }*/
    CK_ATTRIBUTE RSAFindPriTemplate[] = {{CKA_LABEL, pLabel, privateLabelSize}};

    CK_ULONG numHandles = 0;
    CK_OBJECT_HANDLE handles[MAX_NUM_OF_HANDLES] = {0};
    CK_OBJECT_HANDLE handles1[MAX_NUM_OF_HANDLES] = {0};
    rv = P11Functions->C_FindObjectsInit(m_hSession, RSAFindPriTemplate, 1);
    rv = P11Functions->C_FindObjects(m_hSession, handles, MAX_NUM_OF_HANDLES, &numHandles);
    if (numHandles == 0)
    {
        cout << "Private key was not created " << endl;
        return CKR_CANCEL;
    }
    for (unsigned int i = 0; i < numHandles; i++)
    {
        cout << "Private key is " << handles[i] << endl;
    }
    numHandles = 0;
    CK_ATTRIBUTE RSAFindPubTemplate[] = {{CKA_LABEL, (CK_VOID_PTR)pLabelPub, publicLabelSize}};
    rv = P11Functions->C_FindObjectsInit(m_hSession, RSAFindPubTemplate, 1);
    rv = P11Functions->C_FindObjects(m_hSession, handles1, MAX_NUM_OF_HANDLES, &numHandles);
    for (unsigned int i = 0; i < numHandles; i++)
    {
        cout << "Public key is " << handles1[i] << endl;
    }
    return rv;
}
#endif
