ObsidianDragon/src/services/wallet_security_controller.cpp

#include "wallet_security_controller.h"
#include "../util/secure_vault.h"
#include "../util/address_validation.h"

#include <cctype>
#include <cstdint>
#include <cstdio>
#include <utility>
#include <vector>

namespace dragonx {
namespace services {

WalletSecurityController::~WalletSecurityController()
{
    clearDeferredEncryption();
}

void WalletSecurityController::beginDeferredEncryption(std::string passphrase, std::string pin)
{
    clearDeferredEncryption();
    deferred_.passphrase = std::move(passphrase);
    deferred_.pin = std::move(pin);
    deferred_.pending = true;
    deferred_.lastConnectAttempt = -10.0;
}

WalletSecurityController::DeferredEncryptionSnapshot WalletSecurityController::deferredEncryption() const
{
    return {deferred_.passphrase, deferred_.pin};
}

bool WalletSecurityController::shouldAttemptDeferredConnect(double nowSeconds, double minIntervalSeconds)
{
    if (!deferred_.pending) return false;
    if (nowSeconds - deferred_.lastConnectAttempt < minIntervalSeconds) return false;
    deferred_.lastConnectAttempt = nowSeconds;
    return true;
}

void WalletSecurityController::clearDeferredEncryption()
{
    secureClear(deferred_.passphrase);
    secureClear(deferred_.pin);
    deferred_.pending = false;
    deferred_.lastConnectAttempt = -10.0;
}

WalletSecurityController::DeferredEncryptionResult WalletSecurityController::runDeferredEncryption(
    DeferredEncryptionSnapshot request, RpcGateway& rpc, VaultGateway* vault)
{
    DeferredEncryptionResult result;
    result.pinProvided = !request.pin.empty();

    std::string error;
    if (!rpc.encryptWallet(request.passphrase, error)) {
        result.error = error.empty() ? "encryptwallet failed" : error;
        secureClear(request.passphrase);
        secureClear(request.pin);
        return result;
    }

    result.encrypted = true;
    result.restartRequired = true;
    if (result.pinProvided && vault) {
        result.pinStored = vault->storePin(request.pin, request.passphrase);
    }

    secureClear(request.passphrase);
    secureClear(request.pin);
    return result;
}

WalletSecurityController::PinValidationResult WalletSecurityController::validatePinSetup(
    const std::string& pin, const std::string& confirmation, bool allowEmpty, std::size_t minLength)
{
    if (pin.empty() && confirmation.empty()) {
        return allowEmpty
            ? PinValidationResult{true, PinValidationError::None, ""}
            : PinValidationResult{false, PinValidationError::Empty, "PIN is required"};
    }
    if (pin != confirmation) {
        return {false, PinValidationError::Mismatch, "PINs do not match"};
    }
    if (pin.size() < minLength) {
        return {false, PinValidationError::TooShort, "PIN is too short"};
    }
    for (unsigned char c : pin) {
        if (!std::isdigit(c)) {
            return {false, PinValidationError::NonDigit, "PIN must contain only digits"};
        }
    }
    return {true, PinValidationError::None, ""};
}

WalletSecurityController::KeyKind WalletSecurityController::classifyAddress(const std::string& address)
{
    return !address.empty() && address[0] == 'z' ? KeyKind::Shielded : KeyKind::Transparent;
}

WalletSecurityController::KeyKind WalletSecurityController::classifyPrivateKey(const std::string& key)
{
    // Shielded: Sapling extended spending key (secret-extended-key-...) or legacy Sprout (SK...).
    // (The old `key[0]=='s'` test misrouted an uppercase "SK..." shielded key to the transparent RPC.)
    if (key.rfind("secret-extended-key-", 0) == 0) return KeyKind::Shielded;
    if (key.size() >= 2 && key[0] == 'S' && key[1] == 'K') return KeyKind::Shielded;
    // A "z"-prefixed key is a shielded viewing key (zxview…); "s"-prefixed is a legacy Sprout key.
    if (!key.empty() && (key[0] == 's' || key[0] == 'z')) return KeyKind::Shielded;
    return KeyKind::Transparent;
}

bool WalletSecurityController::isViewingKey(const std::string& key)
{
    // DragonX's z_exportviewingkey returns a Sapling *incoming* viewing key (mainnet HRP "zivks");
    // z_importviewingkey only decodes that form. Recognize it structurally — a valid Bech32 checksum
    // plus a known HRP — instead of a bare prefix, and cover testnet/regtest too. (The old check
    // looked for Zcash's "zxview" extended-FVK HRP, which DragonX never emits, so every real viewing
    // key was rejected client-side.) Watch-only: reveals the address's funds but cannot spend them.
    const std::string hrp = util::bech32Hrp(key);
    return hrp == "zivks"                 // mainnet
        || hrp == "zivktestsapling"       // testnet
        || hrp == "zivkregtestsapling";   // regtest
}

bool WalletSecurityController::isRecognizedPrivateKey(const std::string& key)
{
    // Sapling z spending key (HRP "secret-extended-key-{main,test,regtest}"). These run ~300 chars,
    // past the Bech32 length cap, so match by HRP prefix and let the daemon vet the payload.
    if (key.rfind("secret-extended-key-", 0) == 0) return true;               // Sapling z spending key
    if (key.size() >= 2 && key[0] == 'S' && key[1] == 'K') return true;       // Sprout z spending key
    // Transparent WIF: decode Base58Check and confirm it is actually a secret key — version byte plus
    // a 32-byte key, optionally a compression flag (payload 33 or 34 bytes). This accepts BOTH the
    // compressed ("U…") and uncompressed ("7…") mainnet forms and the testnet form, and rejects
    // addresses / typos via the real checksum — the old length+first-char heuristic dropped the
    // uncompressed mainnet key (which starts with '7', not one of 5/K/L/U).
    std::vector<std::uint8_t> payload;
    if (util::decodeBase58Check(key, payload) &&
        (payload.size() == 33 || payload.size() == 34) &&
        (payload[0] == 188 /* DragonX main/regtest SECRET_KEY */ ||
         payload[0] == 128 /* DragonX testnet SECRET_KEY */)) {
        return true;
    }
    return false;
}

bool WalletSecurityController::isRecognizedImportKey(const std::string& key)
{
    return isRecognizedPrivateKey(key) || isViewingKey(key);
}

const char* WalletSecurityController::importSuccessMessage(KeyKind kind)
{
    return kind == KeyKind::Shielded
        ? "Z-address key imported successfully. Wallet is rescanning."
        : "T-address key imported successfully. Wallet is rescanning.";
}

std::string WalletSecurityController::decryptExportFileName(std::uint64_t timestampSeconds)
{
    char buffer[64];
    snprintf(buffer, sizeof(buffer), "obsidiandecryptexport%llu",
             static_cast<unsigned long long>(timestampSeconds));
    return std::string(buffer);
}

void WalletSecurityController::secureClear(std::string& value)
{
    if (!value.empty()) {
        util::SecureVault::secureZero(&value[0], value.size());
        value.clear();
    }
}

} // namespace services
} // namespace dragonx