/*
 * 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 <iostream>
#include <vector>
using namespace std;

CK_RV HSMLunaClient::RSA_CreateSignature(vector<CK_BYTE> data,
                                         string privateKeyLabel,
                                         string publicKeyLabel,
                                         vector<u_int8_t>& result_signature)
{
    CK_ULONG dataLength = data.size();
    unsigned int labelSize = privateKeyLabel.size() + 1;
    CK_BYTE* pLabel = new CK_BYTE[labelSize];
    memset(pLabel, 0, labelSize);

    for (unsigned int i = 0; i < privateKeyLabel.size(); i++)
    {
        pLabel[i] = privateKeyLabel[i];
    }
    CK_ATTRIBUTE RSAFindPriTemplate = {CKA_LABEL, pLabel, labelSize};
    CK_OBJECT_HANDLE priv_key;
    CK_ULONG numHandles = 0;
    CK_OBJECT_HANDLE handles[MAX_NUM_OF_HANDLES] = {0};
    CK_RV rv = P11Functions->C_FindObjectsInit(m_hSession, &RSAFindPriTemplate, 1);
    rv = P11Functions->C_FindObjects(m_hSession, handles, MAX_NUM_OF_HANDLES, &numHandles);
    if (rv != CKR_OK)
    {
        delete[] pLabel;
        return rv;
    }
    if (numHandles == 0)
    {
        RSAFindPriTemplate.ulValueLen--;
        rv = P11Functions->C_FindObjectsInit(m_hSession, &RSAFindPriTemplate, 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;
        }
    }
    if (numHandles == 0)
    {
        cout << "Private key not found!" << endl;
        delete[] pLabel;
        return CKR_FUNCTION_FAILED;
    }
    priv_key = handles[0];
    cout << " The private key is " << hex << priv_key << endl;

    CK_MECHANISM mechanism = {CKM_SHA512_RSA_PKCS, NULL_PTR, 0};
    CK_BYTE* signature = NULL;

    CK_ULONG signatureLength = 0;
    if (rv != CKR_OK)
    {
        delete[] pLabel;
        return rv;
    }

    rv = P11Functions->C_SignInit(m_hSession, &mechanism, priv_key);
    if (rv != CKR_OK)
    {
        cout << " C_SignInit has failed rv = " << hex << rv << endl;
        delete[] pLabel;
        return rv;
    }
    rv = P11Functions->C_Sign(m_hSession, &data[0], dataLength, NULL, &signatureLength);
    if (rv != CKR_OK)
    {
        cout << " C_Sign1 has failed rv = " << hex << rv << endl;
        delete[] pLabel;
        return rv;
    }
    signature = new CK_BYTE[signatureLength];
    memset(signature, 0, signatureLength * sizeof(CK_BYTE));
    rv = P11Functions->C_Sign(m_hSession, &data[0], dataLength, signature, &signatureLength);
    if (rv != CKR_OK)
    {
        cout << " C_Sign2 has failed rv = " << hex << rv << endl;
        delete[] signature;
        delete[] pLabel;
        return rv;
    }
    printf("Signature length is %lu bytes \n", (unsigned long)signatureLength);
    for (unsigned int i = 0; i < signatureLength; i++)
    {
        printf("0x%02x ", signature[i]);
    }
    cout << endl;

    if (publicKeyLabel.empty() == false)
    {
        unsigned int publicLabelSize = publicKeyLabel.size() + 1;
        CK_BYTE* pPublicLabel = new CK_BYTE[publicLabelSize];
        memset(pPublicLabel, 0, publicLabelSize);
        for (unsigned int i = 0; i < publicKeyLabel.size(); i++)
        {
            pPublicLabel[i] = publicKeyLabel[i];
        }
        CK_ATTRIBUTE RSAFindPubTemplate = {CKA_LABEL, pPublicLabel, publicLabelSize};
        CK_OBJECT_HANDLE pub_key;
        numHandles = 0;
        rv = P11Functions->C_FindObjectsInit(m_hSession, &RSAFindPubTemplate, 1);
        rv = P11Functions->C_FindObjects(m_hSession, handles, MAX_NUM_OF_HANDLES, &numHandles);
        if (numHandles == 0)
        {
            RSAFindPubTemplate.ulValueLen--;
            rv = P11Functions->C_FindObjectsInit(m_hSession, &RSAFindPubTemplate, 1);
            if (rv != CKR_OK)
            {
                delete[] pPublicLabel;
                delete[] pLabel;
                return rv;
            }
            rv = P11Functions->C_FindObjects(m_hSession, handles, MAX_NUM_OF_HANDLES, &numHandles);
            if (rv != CKR_OK)
            {
                delete[] pPublicLabel;
                return rv;
            }
        }
        if (numHandles == 0)
        {
            cout << "Public key not found!" << endl;
            delete[] pPublicLabel;
            return CKR_FUNCTION_FAILED;
        }
        pub_key = handles[0];
        cout << " The public key is " << hex << pub_key << endl;
        /* Verify the known hash using the public key. */
        rv = P11Functions->C_VerifyInit(m_hSession, &mechanism, pub_key);
        if (rv != CKR_OK)
        {
            cout << "Signature C_VerifyInit failed. Error is " << hex << rv << endl;
            delete[] signature;
            return rv;
        }
        rv = P11Functions->C_Verify(m_hSession, &data[0], dataLength, signature, signatureLength);

        if (rv != CKR_OK)
        {
            cout << "Signature verifying failed. Error is " << hex << rv << endl;
            return CKR_GENERAL_ERROR;
        }
        cout << " RSA_CreateSignature is verified!" << endl;
    }
    result_signature.resize(signatureLength);
    for (unsigned int i = 0; i < signatureLength; i++)
    {
        result_signature[i] = signature[i];
    }
    delete[] signature;
    delete[] pLabel;
    return rv;
}

CK_RV HSMLunaClient::RSA_EncryptDataTest(string privateKeyLabel, string publicKeyLebel)
{
    CK_MECHANISM mech2;
    CK_OBJECT_HANDLE publ_key, priv_key;
    /** generate data **/
    CK_BYTE cipher[1024 * 4] = {0};
    unsigned int orig_len = 256;
    CK_BYTE original[256] = {0};
    CK_ULONG cipher_len, decipher_len;
    for (unsigned int i = 0; i < orig_len; i++)
        original[i] = i % 255;

    /** set mech2 for RSA crypto **/
    mech2.mechanism = CKM_RSA_PKCS;
    mech2.ulParameterLen = 0;
    mech2.pParameter = NULL;

    CK_ATTRIBUTE RSAFindPriTemplate[] = {{CKA_LABEL, (void*)privateKeyLabel.c_str(), privateKeyLabel.size()}};
    CK_ULONG numHandles = 0;
    CK_OBJECT_HANDLE handles[MAX_NUM_OF_HANDLES] = {0};
    CK_OBJECT_HANDLE handles1[MAX_NUM_OF_HANDLES] = {0};

    CK_RV rv = P11Functions->C_FindObjectsInit(m_hSession, RSAFindPriTemplate, 1);
    rv = P11Functions->C_FindObjects(m_hSession, handles, MAX_NUM_OF_HANDLES, &numHandles);

    priv_key = handles[0];
    numHandles = 0;
    CK_ATTRIBUTE RSAFindPubTemplate[] = {{CKA_LABEL, (void*)publicKeyLebel.c_str(), publicKeyLebel.size()}};
    rv = P11Functions->C_FindObjectsInit(m_hSession, RSAFindPubTemplate, 1);
    rv = P11Functions->C_FindObjects(m_hSession, handles1, MAX_NUM_OF_HANDLES, &numHandles);
    publ_key = handles1[0];

    if (rv == CKR_OK)
    {
        /** initialize RSA encryption (with private key) **/
        rv = P11Functions->C_EncryptInit(m_hSession, &mech2, publ_key);
        if (rv != CKR_OK)
        {
            cout << "C_EncryptInit " << hex << rv << endl;
            // testcase_error("C_EncryptInit rc=%s", p11_get_ckr(rc));
            return rv;
        }
        cipher_len = sizeof(cipher); // set cipher buffer size
        /** do RSA decryption (with public key) **/
        rv = P11Functions->C_Encrypt(m_hSession, original, orig_len, cipher, (CK_ULONG_PTR)&cipher_len);
        if (rv != CKR_OK)
        {
            cout << "C_Encrypt " << hex << rv << endl;
            return rv;
        }
        rv = P11Functions->C_DecryptInit(m_hSession, &mech2, priv_key);
        if (rv != CKR_OK)
        {
            cout << "C_DecryptInit " << hex << rv << endl;
            return rv;
        }
        // do (private key) decryption
        decipher_len = 0;
        rv = P11Functions->C_Decrypt(m_hSession, cipher, cipher_len, NULL, &decipher_len);
        if (rv != CKR_OK)
        {
            cout << "C_Decrypt " << hex << rv << endl;
            cout << "Error in decryption process 1 !" << endl;
            return rv;
        }
        CK_BYTE* decipher = new CK_BYTE[decipher_len];
        memset(decipher, 0, decipher_len * sizeof(CK_BYTE));
        rv = P11Functions->C_Decrypt(m_hSession, cipher, cipher_len, decipher, &decipher_len);
        if (rv != CKR_OK)
        {
            cout << "Error in decryption process 2 !" << hex << rv << endl;
            return rv;
        }
        rv = P11Functions->C_DestroyObject(m_hSession, publ_key);
        if (rv != CKR_OK)
        {
            cout << "C_DestroyObject public " << hex << rv << endl;
            delete[] decipher;
            return rv;
        }

        rv = P11Functions->C_DestroyObject(m_hSession, priv_key);
        if (rv != CKR_OK)
        {
            cout << "C_DestroyObject private " << hex << rv << endl;
            return rv;
        }
        if (orig_len == decipher_len)
        {
            unsigned int i = 0;
            for (i = 0; i < orig_len; i++)
            {
                if (original[i] != decipher[i])
                {
                    cout << "Error in decrypted data - data mismath !" << endl;
                    delete[] decipher;
                    return -1;
                }
            }
            cout << "Success in encrypt/decrypt !" << endl;
            delete[] decipher;
            return 0;
        }
        cout << "Error in decryption process - size mismatch !" << endl;
        delete[] decipher;
        return -1;
        // compare original and decipher
    }
    return rv;
}
#endif
