#include "vnc_apple_rsa_auth.h"

#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/rsa.h>
#include <openssl/x509.h>

#include <memory>

namespace VncAppleRsaAuth {

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 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 CipherCtxPtr = std::unique_ptr<EVP_CIPHER_CTX, CipherCtxDeleter>;

}

QByteArray packCredential(const QString& text)
{
    QByteArray data = text.toUtf8();
    data.append(char(0));
    if (data.size() < 64) {
        QByteArray padding(64 - data.size(), char(0));
        // Best-effort: if RAND_bytes fails, zero-padding is still
        // correct (the NUL terminator above already unambiguously marks
        // the string's real end for the server) -- only the ECB-pattern-
        // hiding benefit of random padding is lost, not correctness.
        RAND_bytes(reinterpret_cast<unsigned char*>(padding.data()), padding.size());
        data += padding;
    } else {
        data = data.left(64);
    }
    return data;
}

Response computeResponse(const QByteArray& hostKeyDer, const QString& username,
                         const QString& password)
{
    Response response;
    if (hostKeyDer.isEmpty()) {
        return response;
    }

    const auto* derPtr = reinterpret_cast<const unsigned char*>(hostKeyDer.constData());
    EvpPkeyPtr hostKey(d2i_PUBKEY(nullptr, &derPtr, hostKeyDer.size()));
    if (!hostKey) {
        return response;
    }

    unsigned char aesKeyBytes[16];
    if (RAND_bytes(aesKeyBytes, sizeof(aesKeyBytes)) != 1) {
        return response;
    }

    const QByteArray credentials = packCredential(username) + packCredential(password);

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

    QByteArray encryptedCredentials(credentials.size() + EVP_MAX_BLOCK_LENGTH, char(0));
    int outLen1 = 0;
    if (EVP_EncryptUpdate(cipherCtx.get(),
                          reinterpret_cast<unsigned char*>(encryptedCredentials.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*>(encryptedCredentials.data()) + outLen1,
                            &outLen2)
        <= 0) {
        return response;
    }
    encryptedCredentials.resize(outLen1 + outLen2);

    EvpPkeyCtxPtr rsaCtx(EVP_PKEY_CTX_new(hostKey.get(), nullptr));
    if (!rsaCtx || EVP_PKEY_encrypt_init(rsaCtx.get()) <= 0
        || EVP_PKEY_CTX_set_rsa_padding(rsaCtx.get(), RSA_PKCS1_PADDING) <= 0) {
        return response;
    }
    size_t encryptedKeyLen = 0;
    if (EVP_PKEY_encrypt(rsaCtx.get(), nullptr, &encryptedKeyLen, aesKeyBytes, sizeof(aesKeyBytes))
        <= 0) {
        return response;
    }
    QByteArray encryptedAesKey(static_cast<int>(encryptedKeyLen), char(0));
    if (EVP_PKEY_encrypt(rsaCtx.get(), reinterpret_cast<unsigned char*>(encryptedAesKey.data()),
                         &encryptedKeyLen, aesKeyBytes, sizeof(aesKeyBytes))
        <= 0) {
        return response;
    }
    encryptedAesKey.resize(static_cast<int>(encryptedKeyLen));

    response.encryptedCredentials = encryptedCredentials;
    response.encryptedAesKey = encryptedAesKey;
    return response;
}

}