#include "qdhmgr.h"

QECDHMgr::QECDHMgr()
{

}

QString QECDHMgr::GetCurveName()
{
    return this->curveName_;
}

bool QECDHMgr::SetCurveName(QString curveName)
{
    EC_builtin_curve * buffer = nullptr;

    size_t count = EC_get_builtin_curves(nullptr, 0);
    if (count <= 0) {
        return false;
    }

    buffer = new EC_builtin_curve[count];
    if (EC_get_builtin_curves(buffer, count) != count) {
        return false;
    }

    // Check if the curve is in the list.
    for (int i = 0; i < count; ++i) {
        const char* shortName = OBJ_nid2sn(buffer[i].nid);
        const char* comment = buffer[i].comment;
        if (shortName && strlen(shortName) > 0) {
//            qDebug() << buffer[i].nid << "|"<< shortName<<"|" << comment;
            if(curveName == QString(shortName))
            {
                curveName_ = curveName;
                return true;
            }
        }
    }
    return false;
}

bool QECDHMgr::GenerateKeys()
{
    ERR_clear_error();

    // First, create an OSSL_PARAM_BLD.
    QScopedPointer <OSSL_PARAM_BLD , ScopedPointerParamBldDeleter> paramBuild(OSSL_PARAM_BLD_new());
    if (!paramBuild.data())
    {
        return false;
    }

    // Push the curve name to the OSSL_PARAM_BLD.
    if (!OSSL_PARAM_BLD_push_utf8_string(paramBuild.data(), OSSL_PKEY_PARAM_GROUP_NAME, curveName_.toStdString().c_str(), 0)) {
        return false;
    }

    // Convert OSSL_PARAM_BLD to OSSL_PARAM.
    QScopedPointer <OSSL_PARAM , ScopedPointerParamDeleter> params(OSSL_PARAM_BLD_to_param(paramBuild.data()));
    if (!params.data()) {
        return false;
    }

    // Create the EC key generation context.
    QScopedPointer <EVP_PKEY_CTX , ScopedPointerKeyCtxDeleter> ctx(EVP_PKEY_CTX_new_from_name(nullptr, "EC", nullptr));
    if (!ctx.data())
    {
        return false;
    }

    // Initialize the key generation context.
    if (EVP_PKEY_keygen_init(ctx.data()) <= 0) {
        return false;
    }

    // Set the parameters which include the curve name.
    if (!EVP_PKEY_CTX_set_params(ctx.data(), params.data())) {
        return false;
    }

    // Generate a key pair.
    EVP_PKEY* keyPair = nullptr;
    if (EVP_PKEY_generate(ctx.data(), &keyPair) <= 0) {
        return false;
    }
    keyPair_.reset(keyPair);
    return true;
}

QString QECDHMgr::GetPrivateKey()
{
    ERR_clear_error();

    if(keyPair_.data())
    {
        // The private key is stored as a BIGNUM object.
        BIGNUM* privateKey = nullptr;
        if (!EVP_PKEY_get_bn_param(keyPair_.data(), OSSL_PKEY_PARAM_PRIV_KEY, &privateKey)) {
            return "";
        }

        // Convert the BIGNUM to a hex string.
        QString hexPrivateKey = ConvertBigNumToHex(privateKey);
        BN_free(privateKey);
        return hexPrivateKey;
    }else
    {
        return "";
    }
}

QString QECDHMgr::GetPublicKey()
{
    ERR_clear_error();
    if(keyPair_.data())
    {
        // The public key is stored as a byte array.
        // Get the array size.
        size_t keyLength = 0;
        if (!EVP_PKEY_get_octet_string_param(keyPair_.data(), OSSL_PKEY_PARAM_PUB_KEY, nullptr, 0, &keyLength)) {
            return "";
        }

        // Get the key.
        unsigned char * publicKey = new unsigned char [keyLength];
        if (!EVP_PKEY_get_octet_string_param(keyPair_.data(), OSSL_PKEY_PARAM_PUB_KEY, publicKey, keyLength, &keyLength)) {
            return "";
        }

        // Convert the byte array key to a hex string.
        QString hexPublicKey = ConvertDataToHex(QByteArray((char *)publicKey, keyLength));
        delete [] publicKey;
        publicKey = nullptr;
        return hexPublicKey;
    }else
    {
        return "";
    }
}

bool QECDHMgr::SetPublicKey(QString hexPublicKey)
{
    QScopedPointer <EVP_PKEY , ScopedPointerKeyDeleter>  peerPublicKey;
    if(!CreatePeerPublicKey(hexPublicKey, &peerPublicKey))
    {
        return false;
    }

    keyPair_.reset(peerPublicKey.take());
    return true;
}

bool QECDHMgr::SetPrivateKey(QString hexPrivateKey)
{
    QScopedPointer <EVP_PKEY , ScopedPointerKeyDeleter>  peerPrivateKey;
    if(!CreatePeerPrivateKey(hexPrivateKey, &peerPrivateKey))
    {
        return false;
    }
    keyPair_.reset(peerPrivateKey.take());
    return true;
}

bool QECDHMgr::SetKeyPair(QString hexPublicKey, QString hexPrivateKey)
{
    QScopedPointer <EVP_PKEY , ScopedPointerKeyDeleter>  peerKeyPair;
    if(!CreatePeerKeyPair(hexPublicKey, hexPrivateKey, &peerKeyPair))
    {
        return false;
    }
    keyPair_.reset(peerKeyPair.take());
    return true;
}

QString QECDHMgr::DeriveSharedSecret(QString hexPeerPublicKey)
{
    ERR_clear_error();

    // First, you have to create the peer public key object.
    // It takes several calls, so it is done in a separate function.
    QScopedPointer <EVP_PKEY , ScopedPointerKeyDeleter>  peerPublicKey;
    if(!CreatePeerPublicKey(hexPeerPublicKey, &peerPublicKey))
    {
        return "";
    }

    // Create the derivation context.
    QScopedPointer <EVP_PKEY_CTX , ScopedPointerKeyCtxDeleter>  derivationCtx(EVP_PKEY_CTX_new(keyPair_.data(), nullptr));
    if (!derivationCtx.data()) {
        return "";
    }

    // Initialize the derivation context.
    if (EVP_PKEY_derive_init(derivationCtx.data()) <= 0) {
        return "";
    }

    // Set the peer public key object.
    if (EVP_PKEY_derive_set_peer(derivationCtx.data(), peerPublicKey.data()) <= 0) {
        return "";
    }

    // Get the shared secret length.
    size_t sharedSecretLength = 0;
    if (EVP_PKEY_derive(derivationCtx.data(), nullptr, &sharedSecretLength) <= 0) {
        return "";
    }

    if (sharedSecretLength == 0) {
        return "";
    }

    unsigned char * sharedSecret = new unsigned char [sharedSecretLength];
    // Derive the shared secret.
    if (EVP_PKEY_derive(derivationCtx.data(), sharedSecret, &sharedSecretLength) <= 0) {
        return "";
    }
    // Convert to a hex string.
    QString hexSharedSecret = ConvertDataToHex(QByteArray((char *)sharedSecret, sharedSecretLength));
    delete [] sharedSecret;

    return hexSharedSecret;
}

QString QECDHMgr::GetLastError()
{
    unsigned long lastError = ERR_peek_last_error();
    if (lastError == 0) {
        return "";
    }
    char errorString[256];
    ERR_error_string_n(lastError, errorString, sizeof(errorString));
    return errorString;
}
// ***************************** 以下是私有函数 **********************************

QString QECDHMgr::ConvertBigNumToHex(const BIGNUM *bigNum)
{
    char* tmpHexBigNum = BN_bn2hex(bigNum);
    if (!tmpHexBigNum) {
        return "";
    }
    QString hexBigNum(tmpHexBigNum);
    OPENSSL_free(tmpHexBigNum);
    return hexBigNum;
}

BIGNUM*  QECDHMgr::ConvertHexToBigNum(QString hexBigNum)
{
    BIGNUM* bn = nullptr;
    if (!BN_hex2bn(&bn, hexBigNum.toStdString().c_str())) {
        return nullptr;
    }
    return bn;
}

QString QECDHMgr::ConvertDataToHex(QByteArray data)
{
    static const char hexDigits[] = "0123456789ABCDEF";

    QString hex;
    hex.reserve(data.size() * 2);

    for (unsigned char c : data) {
        hex.push_back(hexDigits[c >> 4]);
        hex.push_back(hexDigits[c & 15]);
    }

    return hex;
}

// Converts a hex string to binary data.
QByteArray QECDHMgr::ConvertHexToData(QString hex)
{
    QByteArray data;
    // Must be an even number!
    if (hex.size() & 1) {
        return data;
    }

    auto it = hex.begin();
    while (it != hex.end()) {
        int hi = GetHexValue((*it++).toLatin1());
        int lo = GetHexValue((*it++).toLatin1());
        if (hi == -1 || lo == -1) {
            data.clear();
            return data;
        }
        data.push_back(hi << 4 | lo);
    }

    return data;
}

bool QECDHMgr::CreatePeerPublicKey(QString hexPeerPublicKey, QScopedPointer <EVP_PKEY , ScopedPointerKeyDeleter>*  peerPublicKey)
{
    ERR_clear_error();

    // First, we sould create an OSSL_PARAM_BLD with the curve name
    // and the raw peer public key.
    QScopedPointer <OSSL_PARAM_BLD, ScopedPointerParamBldDeleter> paramBuild(OSSL_PARAM_BLD_new());
    if (!paramBuild.data()) {
        return false;
    }

    // Set the curve name.
    if (!OSSL_PARAM_BLD_push_utf8_string(paramBuild.data(), OSSL_PKEY_PARAM_GROUP_NAME,curveName_.toStdString().c_str(), 0)) {
        return false;
    }

    // Convert the peer hex public key to raw data.
    QByteArray binPeerPublicKey = ConvertHexToData(hexPeerPublicKey);
    if (binPeerPublicKey.size() <= 0) {
        return false;
    }

    // Set the raw peer public key.
    if (!OSSL_PARAM_BLD_push_octet_string(paramBuild.data(), OSSL_PKEY_PARAM_PUB_KEY, binPeerPublicKey.data(), binPeerPublicKey.size())) {
        return false;
    }

    // Convert the OSSL_PARAM_BLD to OSSL_PARAM.
    QScopedPointer <OSSL_PARAM , ScopedPointerParamDeleter> params(OSSL_PARAM_BLD_to_param(paramBuild.data()));
    if (!params.data()) {
        return false;
    }

    // Create a EVP_PKEY context.
    QScopedPointer <EVP_PKEY_CTX , ScopedPointerKeyCtxDeleter> peerPublicKeyCtx(EVP_PKEY_CTX_new_from_name(nullptr, "EC", nullptr));
    if (!peerPublicKeyCtx.data()) {
        return false;
    }

    // Initialize the context.
    if (EVP_PKEY_fromdata_init(peerPublicKeyCtx.data()) <= 0) {
        return false;
    }

    // Create the peer public key object.
    EVP_PKEY* tmp = nullptr;
    if (EVP_PKEY_fromdata(peerPublicKeyCtx.data(), &tmp, EVP_PKEY_PUBLIC_KEY, params.data()) <= 0) {
        return false;
    }
    peerPublicKey->reset(tmp);
    return true;
}

bool QECDHMgr::CreatePeerPrivateKey(QString hexPeerPrivateKey, QScopedPointer <EVP_PKEY , ScopedPointerKeyDeleter>*  peerPrivateKey)
{
    ERR_clear_error();

    // First, we sould create an OSSL_PARAM_BLD with the curve name
    // and the raw peer public key.
    QScopedPointer <OSSL_PARAM_BLD, ScopedPointerParamBldDeleter> paramBuild(OSSL_PARAM_BLD_new());
    if (!paramBuild.data()) {
        return false;
    }

    // Set the curve name.
    if (!OSSL_PARAM_BLD_push_utf8_string(paramBuild.data(), OSSL_PKEY_PARAM_GROUP_NAME,curveName_.toStdString().c_str(), 0)) {
        return false;
    }

    // Convert the peer hex private key to raw data.
    BIGNUM * binPeerPrivateKey = ConvertHexToBigNum(hexPeerPrivateKey);
    if(!binPeerPrivateKey)
    {
        return false;
    }

    // Set the raw peer private key.
    if (!OSSL_PARAM_BLD_push_BN(paramBuild.data(), OSSL_PKEY_PARAM_PRIV_KEY, binPeerPrivateKey)) {
        return false;
    }

    // Convert the OSSL_PARAM_BLD to OSSL_PARAM.
    QScopedPointer <OSSL_PARAM , ScopedPointerParamDeleter> params(OSSL_PARAM_BLD_to_param(paramBuild.data()));
    if (!params.data()) {
        return false;
    }

    // Create a EVP_PKEY context.
    QScopedPointer <EVP_PKEY_CTX , ScopedPointerKeyCtxDeleter> peerPrivateKeyCtx(EVP_PKEY_CTX_new_from_name(nullptr, "EC", nullptr));
    if (!peerPrivateKeyCtx.data()) {
        return false;
    }

    // Initialize the context.
    if (EVP_PKEY_fromdata_init(peerPrivateKeyCtx.data()) <= 0) {
        return false;
    }

    // Create the peer private key object.
    EVP_PKEY* tmp = nullptr;
    if (EVP_PKEY_fromdata(peerPrivateKeyCtx.data(), &tmp, EVP_PKEY_KEYPAIR, params.data()) <= 0) {
        return false;
    }
    peerPrivateKey->reset(tmp);
    return true;
}

bool QECDHMgr::CreatePeerKeyPair(QString hexPeerPublicKey, QString hexPeerPrivateKey, QScopedPointer<EVP_PKEY, ScopedPointerKeyDeleter> *peerPrivateKey)
{
    ERR_clear_error();

    // First, we sould create an OSSL_PARAM_BLD with the curve name
    // and the raw peer public key.
    QScopedPointer <OSSL_PARAM_BLD, ScopedPointerParamBldDeleter> paramBuild(OSSL_PARAM_BLD_new());
    if (!paramBuild.data()) {
        return false;
    }

    // Set the curve name.
    if (!OSSL_PARAM_BLD_push_utf8_string(paramBuild.data(), OSSL_PKEY_PARAM_GROUP_NAME,curveName_.toStdString().c_str(), 0)) {
        return false;
    }

    // Convert the peer hex private key to raw data.
    BIGNUM * binPeerPrivateKey = ConvertHexToBigNum(hexPeerPrivateKey);
    if(!binPeerPrivateKey)
    {
        return false;
    }

    // Set the raw peer private key.
    if (!OSSL_PARAM_BLD_push_BN(paramBuild.data(), OSSL_PKEY_PARAM_PRIV_KEY, binPeerPrivateKey)) {
        return false;
    }

    // Convert the peer hex public key to raw data.
    QByteArray binPeerPublicKey = ConvertHexToData(hexPeerPublicKey);
    if (binPeerPublicKey.size() <= 0) {
        return false;
    }

    // Set the raw peer public key.
    if (!OSSL_PARAM_BLD_push_octet_string(paramBuild.data(), OSSL_PKEY_PARAM_PUB_KEY, binPeerPublicKey.data(), binPeerPublicKey.size())) {
        return false;
    }

    // Convert the OSSL_PARAM_BLD to OSSL_PARAM.
    QScopedPointer <OSSL_PARAM , ScopedPointerParamDeleter> params(OSSL_PARAM_BLD_to_param(paramBuild.data()));
    if (!params.data()) {
        return false;
    }

    // Create a EVP_PKEY context.
    QScopedPointer <EVP_PKEY_CTX , ScopedPointerKeyCtxDeleter> peerPrivateKeyCtx(EVP_PKEY_CTX_new_from_name(nullptr, "EC", nullptr));
    if (!peerPrivateKeyCtx.data()) {
        return false;
    }

    // Initialize the context.
    if (EVP_PKEY_fromdata_init(peerPrivateKeyCtx.data()) <= 0) {
        return false;
    }

    // Create the peer private key object.
    EVP_PKEY* tmp = nullptr;
    if (EVP_PKEY_fromdata(peerPrivateKeyCtx.data(), &tmp, EVP_PKEY_KEYPAIR, params.data()) <= 0) {
        return false;
    }
    peerPrivateKey->reset(tmp);
    return true;
}