// Copyright (c) 2009-2021 The Limenka developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <wallet/crypter.h>

#include <common/system.h>
#include <crypto/aes.h>
#include <crypto/hmac_sha256.h>
#include <crypto/hmac_sha512.h>
#include <crypto/sha256.h>
#include <crypto/sha512.h>
#include <random.h>
#include <support/allocators/secure.h>

#include <type_traits>
#include <vector>

namespace wallet {
int CCrypter::BytesToKeySHA512AES(const std::span<const unsigned char> salt, const SecureString& key_data, int count, unsigned char* key, unsigned char* iv) const
{
    if(!count || !key || !iv)
        return 0;

    unsigned char buf[CSHA512::OUTPUT_SIZE];
    CSHA512 di;

    di.Write(UCharCast(key_data.data()), key_data.size());
    di.Write(salt.data(), salt.size());
    di.Finalize(buf);

    for(int i = 0; i != count - 1; i++)
        di.Reset().Write(buf, sizeof(buf)).Finalize(buf);

    memcpy(key, buf, WALLET_CRYPTO_KEY_SIZE);
    memcpy(iv, buf + WALLET_CRYPTO_KEY_SIZE, WALLET_CRYPTO_IV_SIZE);
    memory_cleanse(buf, sizeof(buf));
    return WALLET_CRYPTO_KEY_SIZE;
}

int CCrypter::PBKDF2_SHA512_AES(const std::span<const unsigned char> salt, const SecureString& key_data, int count, unsigned char* key, unsigned char* iv, unsigned char* hmac_key) const
{
    if (!count || !key || !iv || !hmac_key) return 0;

    static constexpr unsigned int DK_LEN = WALLET_CRYPTO_KEY_SIZE + WALLET_CRYPTO_IV_SIZE + WALLET_CRYPTO_HMAC_SIZE;
    unsigned char dk[DK_LEN];
    const unsigned char* pass_bytes = UCharCast(key_data.data());
    const size_t pass_len = key_data.size();
    const size_t salt_len = salt.size();

    unsigned int blocks = (DK_LEN + CSHA512::OUTPUT_SIZE - 1) / CSHA512::OUTPUT_SIZE;
    for (unsigned int block = 1; block <= blocks; block++) {
        unsigned char u[CSHA512::OUTPUT_SIZE];
        unsigned char t[CSHA512::OUTPUT_SIZE];
        memset(t, 0, sizeof(t));

        // U_1 = PRF(Password, Salt || INT_32_BE(i))
        CHMAC_SHA512 hmac_u1(pass_bytes, pass_len);
        hmac_u1.Write(salt.data(), salt_len);
        unsigned char ibuf[4] = {
            static_cast<unsigned char>((block >> 24) & 0xff),
            static_cast<unsigned char>((block >> 16) & 0xff),
            static_cast<unsigned char>((block >> 8) & 0xff),
            static_cast<unsigned char>(block & 0xff)
        };
        hmac_u1.Write(ibuf, 4);
        hmac_u1.Finalize(u);
        for (unsigned int j = 0; j < CSHA512::OUTPUT_SIZE; j++) t[j] = u[j];

        for (int iter = 1; iter < count; iter++) {
            CHMAC_SHA512 hmac_iter(pass_bytes, pass_len);
            hmac_iter.Write(u, CSHA512::OUTPUT_SIZE);
            hmac_iter.Finalize(u);
            for (unsigned int j = 0; j < CSHA512::OUTPUT_SIZE; j++) t[j] ^= u[j];
        }

        unsigned int out_offset = (block - 1) * CSHA512::OUTPUT_SIZE;
        unsigned int copy_len = std::min<unsigned int>(CSHA512::OUTPUT_SIZE, DK_LEN - out_offset);
        memcpy(dk + out_offset, t, copy_len);
        memory_cleanse(t, sizeof(t));
        memory_cleanse(u, sizeof(u));
    }

    memcpy(key, dk, WALLET_CRYPTO_KEY_SIZE);
    memcpy(iv, dk + WALLET_CRYPTO_KEY_SIZE, WALLET_CRYPTO_IV_SIZE);
    memcpy(hmac_key, dk + WALLET_CRYPTO_KEY_SIZE + WALLET_CRYPTO_IV_SIZE, WALLET_CRYPTO_HMAC_SIZE);
    memory_cleanse(dk, sizeof(dk));
    return WALLET_CRYPTO_KEY_SIZE;
}

bool CCrypter::SetKeyFromPassphrase(const SecureString& key_data, const std::span<const unsigned char> salt, const unsigned int rounds, const unsigned int derivation_method)
{
    if (derivation_method == 2) {
        if (rounds < 1 || salt.size() != WALLET_CRYPTO_SALT_SIZE_V2) {
            return false;
        }
        int i = PBKDF2_SHA512_AES(salt, key_data, rounds, vchKey.data(), vchIV.data(), vchHmacKey.data());
        if (i != (int)WALLET_CRYPTO_KEY_SIZE) {
            memory_cleanse(vchKey.data(), vchKey.size());
            memory_cleanse(vchIV.data(), vchIV.size());
            memory_cleanse(vchHmacKey.data(), vchHmacKey.size());
            return false;
        }
        fKeySet = true;
        fUseAEAD = true;
        return true;
    }

    // Legacy method 0
    if (rounds < 1 || salt.size() != WALLET_CRYPTO_SALT_SIZE) {
        return false;
    }

    int i = 0;
    if (derivation_method == 0) {
        i = BytesToKeySHA512AES(salt, key_data, rounds, vchKey.data(), vchIV.data());
    }

    if (i != (int)WALLET_CRYPTO_KEY_SIZE)
    {
        memory_cleanse(vchKey.data(), vchKey.size());
        memory_cleanse(vchIV.data(), vchIV.size());
        return false;
    }

    fKeySet = true;
    fUseAEAD = false;
    return true;
}

bool CCrypter::SetKey(const CKeyingMaterial& new_key, const std::span<const unsigned char> new_iv)
{
    if (new_key.size() != WALLET_CRYPTO_KEY_SIZE || new_iv.size() != WALLET_CRYPTO_IV_SIZE) {
        return false;
    }

    memcpy(vchKey.data(), new_key.data(), new_key.size());
    memcpy(vchIV.data(), new_iv.data(), new_iv.size());
    fUseAEAD = false;

    fKeySet = true;
    return true;
}

bool CCrypter::Encrypt(const CKeyingMaterial& vchPlaintext, std::vector<unsigned char> &vchCiphertext) const
{
    if (!fKeySet)
        return false;

    // max ciphertext len for a n bytes of plaintext is
    // n + AES_BLOCKSIZE bytes
    std::vector<unsigned char> raw(vchPlaintext.size() + AES_BLOCKSIZE);

    AES256CBCEncrypt enc(vchKey.data(), vchIV.data(), true);
    size_t nLen = enc.Encrypt(vchPlaintext.data(), vchPlaintext.size(), raw.data());
    if(nLen < vchPlaintext.size())
        return false;
    raw.resize(nLen);

    if (fUseAEAD) {
        // Prepend HMAC-SHA256(ciphertext)
        CHMAC_SHA256 hmac(vchHmacKey.data(), vchHmacKey.size());
        hmac.Write(raw.data(), raw.size());
        unsigned char mac[CSHA256::OUTPUT_SIZE];
        hmac.Finalize(mac);

        vchCiphertext.resize(WALLET_CRYPTO_HMAC_SIZE + raw.size());
        memcpy(vchCiphertext.data(), mac, WALLET_CRYPTO_HMAC_SIZE);
        memcpy(vchCiphertext.data() + WALLET_CRYPTO_HMAC_SIZE, raw.data(), raw.size());
    } else {
        vchCiphertext = std::move(raw);
    }

    return true;
}

bool CCrypter::Decrypt(const std::span<const unsigned char> ciphertext, CKeyingMaterial& plaintext) const
{
    if (!fKeySet)
        return false;

    const unsigned char* ct_data;
    size_t ct_len;

    if (fUseAEAD) {
        if (ciphertext.size() < WALLET_CRYPTO_HMAC_SIZE) return false;

        // Verify HMAC
        const unsigned char* stored_mac = ciphertext.data();
        const unsigned char* ct_part = ciphertext.data() + WALLET_CRYPTO_HMAC_SIZE;
        size_t ct_part_len = ciphertext.size() - WALLET_CRYPTO_HMAC_SIZE;

        CHMAC_SHA256 hmac(vchHmacKey.data(), vchHmacKey.size());
        hmac.Write(ct_part, ct_part_len);
        unsigned char computed_mac[CSHA256::OUTPUT_SIZE];
        hmac.Finalize(computed_mac);

        if (memcmp(stored_mac, computed_mac, WALLET_CRYPTO_HMAC_SIZE) != 0) {
            memory_cleanse(computed_mac, sizeof(computed_mac));
            return false;
        }
        memory_cleanse(computed_mac, sizeof(computed_mac));

        ct_data = ct_part;
        ct_len = ct_part_len;
    } else {
        ct_data = ciphertext.data();
        ct_len = ciphertext.size();
    }

    // plaintext will always be equal to or lesser than length of ciphertext
    plaintext.resize(ct_len);

    AES256CBCDecrypt dec(vchKey.data(), vchIV.data(), true);
    int len = dec.Decrypt(ct_data, ct_len, plaintext.data());
    if (len == 0) {
        return false;
    }
    plaintext.resize(len);
    return true;
}

bool EncryptSecret(const CKeyingMaterial& vMasterKey, const CKeyingMaterial &vchPlaintext, const uint256& nIV, std::vector<unsigned char> &vchCiphertext)
{
    CCrypter cKeyCrypter;
    std::vector<unsigned char> chIV(WALLET_CRYPTO_IV_SIZE);
    memcpy(chIV.data(), &nIV, WALLET_CRYPTO_IV_SIZE);
    if(!cKeyCrypter.SetKey(vMasterKey, chIV))
        return false;
    return cKeyCrypter.Encrypt(vchPlaintext, vchCiphertext);
}

bool DecryptSecret(const CKeyingMaterial& master_key, const std::span<const unsigned char> ciphertext, const uint256& iv, CKeyingMaterial& plaintext)
{
    CCrypter key_crypter;
    static_assert(WALLET_CRYPTO_IV_SIZE <= std::remove_reference_t<decltype(iv)>::size());
    const std::span iv_prefix{iv.data(), WALLET_CRYPTO_IV_SIZE};
    if (!key_crypter.SetKey(master_key, iv_prefix)) {
        return false;
    }
    return key_crypter.Decrypt(ciphertext, plaintext);
}

bool DecryptKey(const CKeyingMaterial& master_key, const std::span<const unsigned char> crypted_secret, const CPubKey& pub_key, CKey& key)
{
    CKeyingMaterial secret;
    if (!DecryptSecret(master_key, crypted_secret, pub_key.GetHash(), secret)) {
        return false;
    }

    if (secret.size() != 32) {
        return false;
    }

    key.Set(secret.begin(), secret.end(), pub_key.IsCompressed());
    return key.VerifyPubKey(pub_key);
}

bool EncryptSecretV2(const CKeyingMaterial& master_key, const CKeyingMaterial& plaintext, std::vector<unsigned char>& ciphertext)
{
    // Use random IV for V2 key encryption
    uint256 random_iv = GetRandHash();
    std::vector<unsigned char> chIV(WALLET_CRYPTO_IV_SIZE);
    memcpy(chIV.data(), random_iv.begin(), WALLET_CRYPTO_IV_SIZE);

    CCrypter key_crypter;
    if (!key_crypter.SetKey(master_key, chIV)) return false;
    if (!key_crypter.Encrypt(plaintext, ciphertext)) return false;

    // Prepend IV so it can be recovered during decryption
    std::vector<unsigned char> result(WALLET_CRYPTO_IV_SIZE + ciphertext.size());
    memcpy(result.data(), chIV.data(), WALLET_CRYPTO_IV_SIZE);
    memcpy(result.data() + WALLET_CRYPTO_IV_SIZE, ciphertext.data(), ciphertext.size());
    ciphertext = std::move(result);
    return true;
}

bool DecryptSecretV2(const CKeyingMaterial& master_key, std::span<const unsigned char> ciphertext, uint256& iv_out, CKeyingMaterial& plaintext)
{
    if (ciphertext.size() < WALLET_CRYPTO_IV_SIZE) return false;

    const std::span iv_prefix{ciphertext.data(), WALLET_CRYPTO_IV_SIZE};
    const std::span ct_part{ciphertext.data() + WALLET_CRYPTO_IV_SIZE, ciphertext.size() - WALLET_CRYPTO_IV_SIZE};

    memcpy(iv_out.begin(), iv_prefix.data(), WALLET_CRYPTO_IV_SIZE);

    CCrypter key_crypter;
    if (!key_crypter.SetKey(master_key, iv_prefix)) return false;
    return key_crypter.Decrypt(ct_part, plaintext);
}
} // namespace wallet