// Copyright © 2017-2020 Trust Wallet.

//

// This file is part of Trust. The full Trust copyright notice, including

// terms governing use, modification, and redistribution, is contained in the

// file LICENSE at the root of the source code distribution tree.

#include "HDWallet.h"

#include "Base58.h"

#include "BinaryCoding.h"

#include "Bitcoin/SegwitAddress.h"

#include "Bitcoin/CashAddress.h"

#include "Coin.h"

#include <TrezorCrypto/bip32.h>

#include <TrezorCrypto/bip39.h>

#include <TrezorCrypto/curves.h>

#include <TrustWalletCore/TWHRP.h>

#include <array>

using namespace TW;

namespace {

uint32_t fingerprint(HDNode *node, Hash::Hasher hasher);

std::string serialize(const HDNode *node, uint32_t fingerprint, uint32_t version, bool use_public, Hash::Hasher hasher);

bool deserialize(const std::string& extended, TWCurve curve, Hash::Hasher hasher, HDNode *node);

HDNode getNode(const HDWallet& wallet, TWCurve curve, const DerivationPath& derivationPath);

HDNode getMasterNode(const HDWallet& wallet, TWCurve curve);

const char* curveName(TWCurve curve);

} // namespace

bool HDWallet::isValid(const std::string& mnemonic) {

return mnemonic_check(mnemonic.c_str()) != 0;

}

HDWallet::HDWallet(int strength, const std::string& passphrase)

: seed(), mnemonic(), passphrase(passphrase) {

std::array<char, HDWallet::maxMnemomincSize> mnemonic_chars;

mnemonic_generate(strength, mnemonic_chars.data());

mnemonic_to_seed(mnemonic_chars.data(), passphrase.c_str(), seed.data(), nullptr);

mnemonic = mnemonic_chars.data();

updateEntropy();

}

HDWallet::HDWallet(const std::string& mnemonic, const std::string& passphrase)

: seed(), mnemonic(mnemonic), passphrase(passphrase) {

mnemonic_to_seed(mnemonic.c_str(), passphrase.c_str(), seed.data(), nullptr);

updateEntropy();

}

HDWallet::HDWallet(const Data& data, const std::string& passphrase)

: seed(), mnemonic(), passphrase(passphrase) {

std::array<char, HDWallet::maxMnemomincSize> mnemonic_chars;

mnemonic_from_data(data.data(), data.size(), mnemonic_chars.data());

mnemonic_to_seed(mnemonic_chars.data(), passphrase.c_str(), seed.data(), nullptr);

mnemonic = mnemonic_chars.data();

updateEntropy();

}

HDWallet::~HDWallet() {

std::fill(seed.begin(), seed.end(), 0);

std::fill(mnemonic.begin(), mnemonic.end(), 0);

std::fill(passphrase.begin(), passphrase.end(), 0);

}

void HDWallet::updateEntropy() {

// generate entropy (from mnemonic)

Data entropyRaw(32 + 1);

auto entropyBits = mnemonic_to_entropy(mnemonic.c_str(), entropyRaw.data());

// copy to truncate

entropy = data(entropyRaw.data(), entropyBits / 8);

}

PrivateKey HDWallet::getMasterKey(TWCurve curve) const {

auto node = getMasterNode(*this, curve);

auto data = Data(node.private_key, node.private_key + PrivateKey::size);

return PrivateKey(data);

}

PrivateKey HDWallet::getMasterKeyExtension(TWCurve curve) const {

auto node = getMasterNode(*this, curve);

auto data = Data(node.private_key_extension, node.private_key_extension + PrivateKey::size);

return PrivateKey(data);

}

PrivateKey HDWallet::getKey(const DerivationPath& derivationPath) const {

const auto curve = TWCoinTypeCurve(derivationPath.coin());

const auto privateKeyType = getPrivateKeyType(curve);

auto node = getNode(*this, curve, derivationPath);

switch (privateKeyType) {

case PrivateKeyTypeExtended96:

{

auto pkData = Data(node.private_key, node.private_key + PrivateKey::size);

auto extData = Data(node.private_key_extension, node.private_key_extension + PrivateKey::size);

auto chainCode = Data(node.chain_code, node.chain_code + PrivateKey::size);

return PrivateKey(pkData, extData, chainCode);

}

case PrivateKeyTypeDefault32:

default:

// default path

auto data = Data(node.private_key, node.private_key + PrivateKey::size);

return PrivateKey(data);

}

}

std::string HDWallet::deriveAddress(TWCoinType coin) const {

const auto derivationPath = TW::derivationPath(coin);

return TW::deriveAddress(coin, getKey(derivationPath));

}

std::string HDWallet::getExtendedPrivateKey(TWPurpose purpose, TWCoinType coin, TWHDVersion version) const {

if (version == TWHDVersionNone) {

return "";

}

const auto curve = TWCoinTypeCurve(coin);

auto derivationPath = TW::DerivationPath({DerivationPathIndex(purpose, true), DerivationPathIndex(coin, true)});

auto node = getNode(*this, curve, derivationPath);

auto fingerprintValue = fingerprint(&node, publicKeyHasher(coin));

hdnode_private_ckd(&node, 0x80000000);

return serialize(&node, fingerprintValue, version, false, base58Hasher(coin));

}

std::string HDWallet::getExtendedPublicKey(TWPurpose purpose, TWCoinType coin, TWHDVersion version) const {

if (version == TWHDVersionNone) {

return "";

}

const auto curve = TWCoinTypeCurve(coin);

auto derivationPath = TW::DerivationPath({DerivationPathIndex(purpose, true), DerivationPathIndex(coin, true)});

auto node = getNode(*this, curve, derivationPath);

auto fingerprintValue = fingerprint(&node, publicKeyHasher(coin));

hdnode_private_ckd(&node, 0x80000000);

hdnode_fill_public_key(&node);

return serialize(&node, fingerprintValue, version, true, base58Hasher(coin));

}

std::optional<PublicKey> HDWallet::getPublicKeyFromExtended(const std::string &extended, const DerivationPath& path) {

const auto coin = path.coin();

const auto curve = TW::curve(coin);

const auto hasher = TW::base58Hasher(coin);

auto node = HDNode{};

if (!deserialize(extended, curve, hasher, &node)) {

return {};

}

hdnode_public_ckd(&node, path.change());

hdnode_public_ckd(&node, path.address());

hdnode_fill_public_key(&node);

switch (curve) {

case TWCurveSECP256k1:

return PublicKey(Data(node.public_key, node.public_key + 33), TWPublicKeyTypeSECP256k1);

case TWCurveED25519:

return PublicKey(Data(node.public_key, node.public_key + 33), TWPublicKeyTypeED25519);

case TWCurveED25519Blake2bNano:

return PublicKey(Data(node.public_key, node.public_key + 33), TWPublicKeyTypeED25519Blake2b);

case TWCurveED25519Extended:

{

// concatenate public key and chain code (2x32 bytes)

Data concat(node.public_key, node.public_key + 32);

append(concat, Data(node.chain_code, node.chain_code + 32));

return PublicKey(concat, TWPublicKeyTypeED25519Extended);

}

case TWCurveCurve25519:

return PublicKey(Data(node.public_key, node.public_key + 32), TWPublicKeyTypeCURVE25519);

case TWCurveNIST256p1:

return PublicKey(Data(node.public_key, node.public_key + 33), TWPublicKeyTypeNIST256p1);

}

}

std::optional<PrivateKey> HDWallet::getPrivateKeyFromExtended(const std::string &extended, const DerivationPath& path) {

const auto coin = path.coin();

const auto curve = TW::curve(coin);

const auto hasher = TW::base58Hasher(coin);

auto node = HDNode{};

if (!deserialize(extended, curve, hasher, &node)) {

return {};

}

hdnode_private_ckd(&node, path.change());

hdnode_private_ckd(&node, path.address());

return PrivateKey(Data(node.private_key, node.private_key + 32));

}

HDWallet::PrivateKeyType HDWallet::getPrivateKeyType(TWCurve curve) {

if (curve == TWCurve::TWCurveED25519Extended) {

// used by Cardano

return PrivateKeyTypeExtended96;

}

// default

return PrivateKeyTypeDefault32;

}

namespace {

uint32_t fingerprint(HDNode *node, Hash::Hasher hasher) {

hdnode_fill_public_key(node);

auto digest = hasher(node->public_key, node->public_key + 33);

return ((uint32_t) digest[0] << 24) + (digest[1] << 16) + (digest[2] << 8) + digest[3];

}

std::string serialize(const HDNode *node, uint32_t fingerprint, uint32_t version, bool use_public, Hash::Hasher hasher) {

Data node_data;

node_data.reserve(78);

encode32BE(version, node_data);

node_data.push_back(static_cast<uint8_t>(node->depth));

encode32BE(fingerprint, node_data);

encode32BE(node->child_num, node_data);

node_data.insert(node_data.end(), node->chain_code, node->chain_code + 32);

if (use_public) {

node_data.insert(node_data.end(), node->public_key, node->public_key + 33);

} else {

node_data.push_back(0);

node_data.insert(node_data.end(), node->private_key, node->private_key + 32);

}

return Base58::bitcoin.encodeCheck(node_data, hasher);

}

bool deserialize(const std::string& extended, TWCurve curve, Hash::Hasher hasher, HDNode *node) {

memset(node, 0, sizeof(HDNode));

node->curve = get_curve_by_name(curveName(curve));

const auto node_data = Base58::bitcoin.decodeCheck(extended, hasher);

if (node_data.size() != 78) {

return false;

}

uint32_t version = decode32BE(node_data.data());

if (TWHDVersionIsPublic(static_cast<TWHDVersion>(version))) {

std::copy(node_data.begin() + 45, node_data.begin() + 45 + 33, node->public_key);

} else if (TWHDVersionIsPrivate(static_cast<TWHDVersion>(version))) {

if (node_data[45]) { // invalid data

return false;

}

std::copy(node_data.begin() + 46, node_data.begin() + 46 + 32, node->private_key);

} else {

return false; // invalid version

}

node->depth = node_data[4];

node->child_num = decode32BE(node_data.data() + 9);

std::copy(node_data.begin() + 13, node_data.begin() + 13 + 32, node->chain_code);

return true;

}

HDNode getNode(const HDWallet& wallet, TWCurve curve, const DerivationPath& derivationPath) {

const auto privateKeyType = HDWallet::getPrivateKeyType(curve);

auto node = getMasterNode(wallet, curve);

for (auto& index : derivationPath.indices) {

switch (privateKeyType) {

case HDWallet::PrivateKeyTypeExtended96:

// special handling for extended

hdnode_private_ckd_cardano(&node, index.derivationIndex());

break;

case HDWallet::PrivateKeyTypeDefault32:

default:

hdnode_private_ckd(&node, index.derivationIndex());

break;

}

}

return node;

}

HDNode getMasterNode(const HDWallet& wallet, TWCurve curve) {

const auto privateKeyType = HDWallet::getPrivateKeyType(curve);

auto node = HDNode();

switch (privateKeyType) {

case HDWallet::PrivateKeyTypeExtended96:

// special handling for extended, use entropy (not seed)

hdnode_from_seed_cardano((const uint8_t*)"", 0, wallet.entropy.data(), (int)wallet.entropy.size(), &node);

break;

case HDWallet::PrivateKeyTypeDefault32:

default:

hdnode_from_seed(wallet.seed.data(), HDWallet::seedSize, curveName(curve), &node);

break;

}

return node;

}

const char* curveName(TWCurve curve) {

switch (curve) {

case TWCurveSECP256k1:

return SECP256K1_NAME;

case TWCurveED25519:

return ED25519_NAME;

case TWCurveED25519Blake2bNano:

return ED25519_BLAKE2B_NANO_NAME;

case TWCurveNIST256p1:

return NIST256P1_NAME;

default:

return "";

}

}

} // namespace