/*
 * Copyright (C) Jan 2013 Mellanox Technologies Ltd. All rights reserved.
 * Copyright (c) 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 NO_OPEN_SSL
#include <openssl/sha.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/bio.h>
#include <openssl/hmac.h>
#include <openssl/engine.h>
#include <openssl/conf.h>
#include <openssl/err.h>
#include <openssl/x509v3.h>
#include "mlxsign_lib.h"

using namespace std;
using namespace MlxSign;

#define CHECK_RC(rc, expRc, errCode) \
    do                               \
    {                                \
        if ((rc) != (expRc))         \
        {                            \
            return (errCode);        \
        }                            \
    } while (0)

MlxSignSHA::MlxSignSHA(u_int32_t digestLength)
{
    _digestLength = digestLength;
}

int MlxSignSHA::getDigest(std::string& digest)
{
    int rc = 0;
    std::vector<u_int8_t> digestVec;
    char* digestStr = new char[_digestLength * 2 + 1];
    digestStr = (char*)memset(digestStr, 0, (_digestLength * 2 + 1));

    rc = getDigest(digestVec);
    if (rc != MlxSign::MLX_SIGN_SUCCESS)
    {
        delete[] digestStr;
        return rc;
    }
    // transform to string
    for (u_int64_t i = 0; i < digestVec.size(); i++)
    {
        snprintf(digestStr + i * 2, 3, "%02x", digestVec[i]);
    }
    digest = digestStr;
    delete[] digestStr;
    return MlxSign::MLX_SIGN_SUCCESS;
}

void MlxSignSHA::reset()
{
    _buff.resize(0);
}

MlxSignSHA& operator<<(MlxSignSHA& lhs, u_int8_t data)
{
    lhs._buff.push_back(data);
    return lhs;
}

MlxSignSHA& operator<<(MlxSignSHA& lhs, const std::vector<u_int8_t>& buff)
{
    for (std::vector<u_int8_t>::const_iterator it = buff.begin(); it != buff.end(); it++)
    {
        lhs._buff.push_back(*it);
    }
    return lhs;
}

/*
 * MlxSignSHA256
 *
 */

MlxSignSHA256::MlxSignSHA256() : MlxSignSHA(SHA256_DIGEST_LENGTH) {}

int MlxSignSHA256::getDigest(std::vector<u_int8_t>& digest)
{
    int rc;
    SHA256_CTX ctx;
    digest.resize(_digestLength);
    memset(&digest[0], 0, digest.size());
    rc = SHA256_Init(&ctx);
    CHECK_RC(rc, 1, MlxSign::MLX_SIGN_SHA_INIT_ERROR);
    rc = SHA256_Update(&ctx, (void*)(&_buff[0]), _buff.size());
    CHECK_RC(rc, 1, MlxSign::MLX_SIGN_SHA_CALCULATION_ERROR);
    rc = SHA256_Final(&digest[0], &ctx);
    CHECK_RC(rc, 1, MlxSign::MLX_SIGN_SHA_CALCULATION_ERROR);
    return MlxSign::MLX_SIGN_SUCCESS;
}

/*
 * MlxSignSHA512
 *
 */
MlxSignSHA512::MlxSignSHA512() : MlxSignSHA(SHA512_DIGEST_LENGTH) {}

int MlxSignSHA512::getDigest(std::vector<u_int8_t>& digest)
{
    int rc;
    SHA512_CTX ctx;
    digest.resize(_digestLength);
    memset(&digest[0], 0, digest.size());
    rc = SHA512_Init(&ctx);
    CHECK_RC(rc, 1, MlxSign::MLX_SIGN_SHA_INIT_ERROR);
    rc = SHA512_Update(&ctx, (void*)(&_buff[0]), _buff.size());
    CHECK_RC(rc, 1, MlxSign::MLX_SIGN_SHA_CALCULATION_ERROR);
    rc = SHA512_Final(&digest[0], &ctx);
    CHECK_RC(rc, 1, MlxSign::MLX_SIGN_SHA_CALCULATION_ERROR);
    return MlxSign::MLX_SIGN_SUCCESS;
}

MlxSignRSA::MlxSignRSA() : _privCtx(NULL), _pubCtx(NULL) {}

MlxSignRSA::~MlxSignRSA()
{
    if (_privCtx)
    {
        RSA_free((RSA*)_privCtx);
    }

    if (_pubCtx)
    {
        RSA_free((RSA*)_pubCtx);
    }
}
int MlxSignRSA::setPrivKeyFromFile(const std::string& pemKeyFilePath)
{
    return createRSAFromPEMFileName(pemKeyFilePath, true);
}

int MlxSignRSA::setPrivKey(const std::string& pemKey)
{
    return createRSAFromPEMKeyString(pemKey, true);
}

int MlxSignRSA::getPrivKeyLength() const
{
    if (_privCtx)
    {
        return RSA_size((RSA*)_privCtx);
    }
    else
    {
        return 0x0;
    }
}

int MlxSignRSA::setPubKeyFromFile(const std::string& pemKeyFilePath)
{
    return createRSAFromPEMFileName(pemKeyFilePath, false);
}

int MlxSignRSA::setPubKey(const std::string& pemKey)
{
    return createRSAFromPEMKeyString(pemKey, false);
}

int MlxSignRSA::sign(MlxSign::SHAType shaType, const std::vector<u_int8_t>& msg, std::vector<u_int8_t>& encryptedMsg)
{
    unsigned int maxMsgSize, signLen;
    std::vector<u_int8_t> encryptedMsgTemp;

    if (!_privCtx)
    {
        return MlxSign::MLX_SIGN_RSA_NO_PRIV_KEY_ERROR;
    }

    // size check
    maxMsgSize = RSA_size((RSA*)_privCtx);
    if (msg.size() > maxMsgSize)
    {
        return MlxSign::MLX_SIGN_RSA_MESSAGE_TOO_LONG_ERROR;
    }
    // do the job
    encryptedMsgTemp.resize(maxMsgSize, 0);

    int type;

    if (shaType == MlxSign::SHA256)
    {
        type = NID_sha256;
    }
    else if (shaType == MlxSign::SHA512)
    {
        type = NID_sha512;
    }
    else
    {
        return MlxSign::MLX_SIGN_UNSUPPORTED_SHA_TYPE;
    }

    if (!RSA_sign(type, msg.data(), msg.size(), encryptedMsgTemp.data(), &signLen, (RSA*)_privCtx))
    {
        return MlxSign::MLX_SIGN_RSA_CALCULATION_ERROR;
    }
    encryptedMsg.resize(signLen, 0);
    memcpy(encryptedMsg.data(), encryptedMsgTemp.data(), signLen);

    return MlxSign::MLX_SIGN_SUCCESS;
}

int MlxSignRSA::verify(MlxSign::SHAType shaType,
                       const std::vector<u_int8_t>& digest,
                       const std::vector<u_int8_t>& sig,
                       bool& result)
{
    if (!_pubCtx)
    {
        return MlxSign::MLX_SIGN_RSA_NO_PUB_KEY_ERROR;
    }

    int type;

    if (shaType == MlxSign::SHA256)
    {
        type = NID_sha256;
    }
    else if (shaType == MlxSign::SHA512)
    {
        type = NID_sha512;
    }
    else
    {
        return MlxSign::MLX_SIGN_UNSUPPORTED_SHA_TYPE;
    }

    result = RSA_verify(type, digest.data(), digest.size(), sig.data(), sig.size(), (RSA*)_pubCtx);

    return MlxSign::MLX_SIGN_SUCCESS;
}

std::string MlxSignRSA::str(const std::vector<u_int8_t>& msg)
{
    char* digestStr = NULL;
    try
    {
        digestStr = new char[msg.size() * 2 + 1];
        memset(digestStr, 0, sizeof(char) * (msg.size() * 2 + 1));
    }
    catch (...)
    {
        return "";
    }

    for (u_int64_t i = 0; i < msg.size(); i++)
    {
        snprintf(digestStr + i * 2, 3, "%02x", msg[i]);
    }
    std::string result = digestStr;
    delete[] digestStr;
    return result;
}

#define REPLACE_RSA_CTX(ctx, newCtx)   \
    do                                 \
    {                                  \
        if ((newCtx))                  \
        {                              \
            if ((ctx))                 \
            {                          \
                RSA_free((RSA*)(ctx)); \
            }                          \
            (ctx) = (void*)(newCtx);   \
        }                              \
    } while (0)

int MlxSignRSA::createRSAFromPEMFileName(const std::string& fname, bool isPrivateKey)
{
    FILE* fp = fopen(fname.c_str(), "rb");
    RSA* rsa = NULL;
    if (!fp)
    {
        return MlxSign::MLX_SIGN_RSA_FILE_OPEN_ERROR;
    }
    if (isPrivateKey)
    {
        rsa = PEM_read_RSAPrivateKey(fp, NULL, NULL, NULL);
        REPLACE_RSA_CTX(_privCtx, rsa);
    }
    else
    {
        rsa = PEM_read_RSA_PUBKEY(fp, NULL, NULL, NULL);
        REPLACE_RSA_CTX(_pubCtx, rsa);
    }

    fclose(fp);
    if (rsa == NULL)
    {
        return MlxSign::MLX_SIGN_RSA_INIT_CTX_ERROR;
    }
    return MlxSign::MLX_SIGN_SUCCESS;
}

int MlxSignRSA::createRSAFromPEMKeyString(const std::string& pemKey, bool isPrivateKey)
{
    RSA* rsa = NULL;
    BIO* keybio;
    // TODO: check if this may leak
    keybio = BIO_new_mem_buf((void*)pemKey.c_str(), -1);
    if (keybio == NULL)
    {
        return MlxSign::MLX_SIGN_RSA_KEY_BIO_ERROR;
    }
    if (isPrivateKey)
    {
        rsa = PEM_read_bio_RSAPrivateKey(keybio, &rsa, NULL, NULL);
        REPLACE_RSA_CTX(_privCtx, rsa);
    }
    else
    {
        rsa = PEM_read_bio_RSA_PUBKEY(keybio, &rsa, NULL, NULL);
        REPLACE_RSA_CTX(_pubCtx, rsa);
    }
    BIO_free_all(keybio);
    if (rsa == NULL)
    {
        return MlxSign::MLX_SIGN_RSA_INIT_CTX_ERROR;
    }
    return MlxSign::MLX_SIGN_SUCCESS;
}

MlxSignHMAC::MlxSignHMAC()
{
#if OPENSSL_VERSION_NUMBER < 0x10100000L
    ctx = malloc(sizeof(HMAC_CTX));
    HMAC_CTX_init((HMAC_CTX*)ctx);
#else
    ctx = HMAC_CTX_new();
#endif
}

int MlxSignHMAC::setKey(const std::vector<u_int8_t>& key)
{
    if (HMAC_Init_ex((HMAC_CTX*)ctx, (char*)key.data(), key.size(), EVP_sha512(), NULL) == 0)
    {
        return MlxSign::MLX_SIGN_HMAC_ERROR;
    }

    return MlxSign::MLX_SIGN_SUCCESS;
}

MlxSignHMAC& operator<<(MlxSignHMAC& lhs, const std::vector<u_int8_t>& buff)
{
    for (std::vector<u_int8_t>::const_iterator it = buff.begin(); it != buff.end(); it++)
    {
        lhs.data.push_back(*it);
    }

    return lhs;
}

int MlxSignHMAC::getDigest(std::vector<u_int8_t>& digest)
{
    unsigned int len = 64; // 512 bits

    if (HMAC_Update((HMAC_CTX*)ctx, (unsigned char*)data.data(), data.size()) == 0)
    {
        return MlxSign::MLX_SIGN_HMAC_ERROR;
    }

    digest.resize(len);

    // TODO why len is passed by reference?
    // TODO HMAC_Final must return 1 on success and 0 on failure but that is not happening!
    HMAC_Final((HMAC_CTX*)ctx, digest.data(), &len);
    return MlxSign::MLX_SIGN_SUCCESS;
}

MlxSignHMAC::~MlxSignHMAC()
{
#if OPENSSL_VERSION_NUMBER < 0x10100000L
    HMAC_CTX_cleanup((HMAC_CTX*)ctx);
    free(ctx);
#else
    HMAC_CTX_free((HMAC_CTX*)ctx);
#endif
}

#endif // ENABLE_OPENSSL
