#include "vnc_apple_dh_auth.h"

#include <openssl/bn.h>
#include <openssl/core_names.h>
#include <openssl/evp.h>
#include <openssl/param_build.h>

#include <cstring>
#include <memory>

namespace VncAppleDhAuth {

namespace {

struct EvpPkeyDeleter {
    void operator()(EVP_PKEY* key) const { EVP_PKEY_free(key); }
};
struct EvpPkeyCtxDeleter {
    void operator()(EVP_PKEY_CTX* ctx) const { EVP_PKEY_CTX_free(ctx); }
};
struct BnDeleter {
    void operator()(BIGNUM* bn) const { BN_free(bn); }
};
struct ParamBldDeleter {
    void operator()(OSSL_PARAM_BLD* bld) const { OSSL_PARAM_BLD_free(bld); }
};
struct ParamDeleter {
    void operator()(OSSL_PARAM* params) const { OSSL_PARAM_free(params); }
};
struct CipherCtxDeleter {
    void operator()(EVP_CIPHER_CTX* ctx) const { EVP_CIPHER_CTX_free(ctx); }
};

using EvpPkeyPtr = std::unique_ptr<EVP_PKEY, EvpPkeyDeleter>;
using EvpPkeyCtxPtr = std::unique_ptr<EVP_PKEY_CTX, EvpPkeyCtxDeleter>;
using BnPtr = std::unique_ptr<BIGNUM, BnDeleter>;
using ParamBldPtr = std::unique_ptr<OSSL_PARAM_BLD, ParamBldDeleter>;
using ParamPtr = std::unique_ptr<OSSL_PARAM, ParamDeleter>;
using CipherCtxPtr = std::unique_ptr<EVP_CIPHER_CTX, CipherCtxDeleter>;

BnPtr bnFromBytes(const QByteArray& bytes)
{
    return BnPtr(
        BN_bin2bn(reinterpret_cast<const unsigned char*>(bytes.constData()), bytes.size(), nullptr));
}

// Builds an EVP_PKEY holding just the DH domain parameters (generator,
// prime) -- used both as the basis for our own keygen and, with a public
// value added, to represent the server's public key for derive().
EvpPkeyPtr buildDomainParams(const BIGNUM* g, const BIGNUM* p)
{
    ParamBldPtr bld(OSSL_PARAM_BLD_new());
    if (!bld) {
        return nullptr;
    }
    if (OSSL_PARAM_BLD_push_BN(bld.get(), OSSL_PKEY_PARAM_FFC_P, p) <= 0
        || OSSL_PARAM_BLD_push_BN(bld.get(), OSSL_PKEY_PARAM_FFC_G, g) <= 0) {
        return nullptr;
    }
    ParamPtr params(OSSL_PARAM_BLD_to_param(bld.get()));
    if (!params) {
        return nullptr;
    }

    EvpPkeyCtxPtr ctx(EVP_PKEY_CTX_new_from_name(nullptr, "DH", nullptr));
    if (!ctx || EVP_PKEY_fromdata_init(ctx.get()) <= 0) {
        return nullptr;
    }
    EVP_PKEY* rawKey = nullptr;
    if (EVP_PKEY_fromdata(ctx.get(), &rawKey, EVP_PKEY_KEY_PARAMETERS, params.get()) <= 0) {
        return nullptr;
    }
    return EvpPkeyPtr(rawKey);
}

EvpPkeyPtr buildPeerPublicKey(const BIGNUM* g, const BIGNUM* p, const BIGNUM* pub)
{
    ParamBldPtr bld(OSSL_PARAM_BLD_new());
    if (!bld) {
        return nullptr;
    }
    if (OSSL_PARAM_BLD_push_BN(bld.get(), OSSL_PKEY_PARAM_FFC_P, p) <= 0
        || OSSL_PARAM_BLD_push_BN(bld.get(), OSSL_PKEY_PARAM_FFC_G, g) <= 0
        || OSSL_PARAM_BLD_push_BN(bld.get(), OSSL_PKEY_PARAM_PUB_KEY, pub) <= 0) {
        return nullptr;
    }
    ParamPtr params(OSSL_PARAM_BLD_to_param(bld.get()));
    if (!params) {
        return nullptr;
    }

    EvpPkeyCtxPtr ctx(EVP_PKEY_CTX_new_from_name(nullptr, "DH", nullptr));
    if (!ctx || EVP_PKEY_fromdata_init(ctx.get()) <= 0) {
        return nullptr;
    }
    EVP_PKEY* rawKey = nullptr;
    if (EVP_PKEY_fromdata(ctx.get(), &rawKey, EVP_PKEY_PUBLIC_KEY, params.get()) <= 0) {
        return nullptr;
    }
    return EvpPkeyPtr(rawKey);
}

}

Response computeResponse(const QByteArray& generator, const QByteArray& prime,
                         const QByteArray& serverPublicKey, const QString& username,
                         const QString& password)
{
    Response response;
    if (generator.isEmpty() || prime.isEmpty() || serverPublicKey.isEmpty()) {
        return response;
    }

    BnPtr g = bnFromBytes(generator);
    BnPtr p = bnFromBytes(prime);
    BnPtr serverPub = bnFromBytes(serverPublicKey);
    if (!g || !p || !serverPub) {
        return response;
    }

    EvpPkeyPtr domainParams = buildDomainParams(g.get(), p.get());
    if (!domainParams) {
        return response;
    }

    // Generate our own ephemeral DH keypair against the same domain
    // parameters the server offered.
    EvpPkeyCtxPtr keygenCtx(EVP_PKEY_CTX_new_from_pkey(nullptr, domainParams.get(), nullptr));
    if (!keygenCtx || EVP_PKEY_keygen_init(keygenCtx.get()) <= 0) {
        return response;
    }
    EVP_PKEY* rawOurKey = nullptr;
    if (EVP_PKEY_keygen(keygenCtx.get(), &rawOurKey) <= 0) {
        return response;
    }
    EvpPkeyPtr ourKey(rawOurKey);

    BIGNUM* ourPubRaw = nullptr;
    if (EVP_PKEY_get_bn_param(ourKey.get(), OSSL_PKEY_PARAM_PUB_KEY, &ourPubRaw) <= 0
        || ourPubRaw == nullptr) {
        return response;
    }
    BnPtr ourPub(ourPubRaw);

    QByteArray clientPublicKey(prime.size(), char(0));
    if (BN_bn2binpad(ourPub.get(), reinterpret_cast<unsigned char*>(clientPublicKey.data()),
                     prime.size())
        < 0) {
        return response;
    }

    EvpPkeyPtr peerKey = buildPeerPublicKey(g.get(), p.get(), serverPub.get());
    if (!peerKey) {
        return response;
    }

    EvpPkeyCtxPtr deriveCtx(EVP_PKEY_CTX_new_from_pkey(nullptr, ourKey.get(), nullptr));
    if (!deriveCtx || EVP_PKEY_derive_init(deriveCtx.get()) <= 0
        || EVP_PKEY_derive_set_peer(deriveCtx.get(), peerKey.get()) <= 0) {
        return response;
    }

    size_t secretLen = 0;
    if (EVP_PKEY_derive(deriveCtx.get(), nullptr, &secretLen) <= 0 || secretLen == 0) {
        return response;
    }
    QByteArray secret(static_cast<int>(secretLen), char(0));
    if (EVP_PKEY_derive(deriveCtx.get(), reinterpret_cast<unsigned char*>(secret.data()),
                        &secretLen)
        <= 0) {
        return response;
    }
    secret.resize(static_cast<int>(secretLen));

    unsigned char aesKey[16];
    if (EVP_Digest(secret.constData(), static_cast<size_t>(secret.size()), aesKey, nullptr,
                   EVP_md5(), nullptr)
        <= 0) {
        return response;
    }

    // 64 bytes username + 64 bytes password, NUL-padded/truncated.
    QByteArray credentials(128, char(0));
    const QByteArray userBytes = username.toLatin1().left(64);
    const QByteArray passBytes = password.toLatin1().left(64);
    std::memcpy(credentials.data(), userBytes.constData(),
               static_cast<size_t>(userBytes.size()));
    std::memcpy(credentials.data() + 64, passBytes.constData(),
               static_cast<size_t>(passBytes.size()));

    CipherCtxPtr cipherCtx(EVP_CIPHER_CTX_new());
    if (!cipherCtx
        || EVP_EncryptInit_ex(cipherCtx.get(), EVP_aes_128_ecb(), nullptr, aesKey, nullptr) <= 0) {
        return response;
    }
    EVP_CIPHER_CTX_set_padding(cipherCtx.get(), 0);

    QByteArray ciphertext(credentials.size() + EVP_MAX_BLOCK_LENGTH, char(0));
    int outLen1 = 0;
    if (EVP_EncryptUpdate(cipherCtx.get(), reinterpret_cast<unsigned char*>(ciphertext.data()),
                          &outLen1, reinterpret_cast<const unsigned char*>(credentials.constData()),
                          credentials.size())
        <= 0) {
        return response;
    }
    int outLen2 = 0;
    if (EVP_EncryptFinal_ex(cipherCtx.get(),
                            reinterpret_cast<unsigned char*>(ciphertext.data()) + outLen1, &outLen2)
        <= 0) {
        return response;
    }
    ciphertext.resize(outLen1 + outLen2);

    response.clientPublicKey = clientPublicKey;
    response.encryptedCredentials = ciphertext;
    return response;
}

}