uhash-cli 0.5.1

//! Wallet management for UniversalHash miner
//!
//! Handles mnemonic generation, import/export, and transaction signing.
//! Mnemonics are encrypted at rest using Argon2id + AES-256-GCM.

use aes_gcm::aead::Aead;
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use argon2::Argon2;
use bip32::secp256k1::ecdsa::SigningKey;
use bip32::{DerivationPath, XPrv};
use bip39::{Language, Mnemonic};
use cosmrs::crypto::secp256k1;
use cosmrs::AccountId;
use rand::RngCore;
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::PathBuf;
use thiserror::Error;

/// Default derivation path for Cosmos SDK chains
const DERIVATION_PATH: &str = "m/44'/118'/0'/0/0";

/// Bostrom address prefix
const BOSTROM_PREFIX: &str = "bostrom";

/// Current keystore file format version
const KEYSTORE_VERSION: u32 = 1;

#[derive(Error, Debug)]
pub enum WalletError {
    #[error("Failed to generate mnemonic: {0}")]
    MnemonicGeneration(String),

    #[error("Invalid mnemonic phrase: {0}")]
    InvalidMnemonic(String),

    #[error("Derivation error: {0}")]
    Derivation(String),

    #[error("File I/O error: {0}")]
    FileError(#[from] std::io::Error),

    #[error("Invalid wallet file format")]
    InvalidFormat,

    #[error("Decryption failed (wrong password or corrupted file)")]
    DecryptionFailed,
}

/// Argon2id KDF parameters stored alongside the ciphertext
#[derive(Serialize, Deserialize, Clone, Debug)]
struct KdfParams {
    m_cost: u32,
    t_cost: u32,
    p_cost: u32,
}

/// Encrypted wallet keystore file format (JSON)
#[derive(Serialize, Deserialize, Debug)]
struct KeystoreFile {
    version: u32,
    address: String,
    kdf: String,
    kdf_params: KdfParams,
    /// Hex-encoded 16-byte salt
    salt: String,
    /// Hex-encoded 12-byte nonce
    nonce: String,
    /// Hex-encoded ciphertext (encrypted mnemonic + GCM auth tag)
    ciphertext: String,
}

/// Derive a 32-byte AES key from password + salt using Argon2id
fn derive_key(password: &[u8], salt: &[u8], params: &KdfParams) -> Result<[u8; 32], WalletError> {
    let argon2 = Argon2::new(
        argon2::Algorithm::Argon2id,
        argon2::Version::V0x13,
        argon2::Params::new(params.m_cost, params.t_cost, params.p_cost, Some(32))
            .map_err(|e| WalletError::MnemonicGeneration(format!("argon2 params: {e}")))?,
    );
    let mut key = [0u8; 32];
    argon2
        .hash_password_into(password, salt, &mut key)
        .map_err(|e| WalletError::MnemonicGeneration(format!("argon2 hash: {e}")))?;
    Ok(key)
}

/// Encrypt plaintext mnemonic with the derived key
fn encrypt_mnemonic(
    mnemonic: &str,
    password: &str,
    address: &str,
) -> Result<KeystoreFile, WalletError> {
    let kdf_params = KdfParams {
        m_cost: 19456,
        t_cost: 2,
        p_cost: 1,
    };

    let mut salt = [0u8; 16];
    rand::thread_rng().fill_bytes(&mut salt);

    let mut nonce_bytes = [0u8; 12];
    rand::thread_rng().fill_bytes(&mut nonce_bytes);

    let key = derive_key(password.as_bytes(), &salt, &kdf_params)?;
    let cipher = Aes256Gcm::new_from_slice(&key)
        .map_err(|e| WalletError::MnemonicGeneration(format!("aes init: {e}")))?;

    let nonce = Nonce::from_slice(&nonce_bytes);
    let ciphertext = cipher
        .encrypt(nonce, mnemonic.as_bytes())
        .map_err(|e| WalletError::MnemonicGeneration(format!("encrypt: {e}")))?;

    Ok(KeystoreFile {
        version: KEYSTORE_VERSION,
        address: address.to_string(),
        kdf: "argon2id".to_string(),
        kdf_params,
        salt: hex::encode(salt),
        nonce: hex::encode(nonce_bytes),
        ciphertext: hex::encode(ciphertext),
    })
}

/// Decrypt a keystore file back to the mnemonic phrase
fn decrypt_mnemonic(keystore: &KeystoreFile, password: &str) -> Result<String, WalletError> {
    let salt = hex::decode(&keystore.salt).map_err(|_| WalletError::InvalidFormat)?;
    let nonce_bytes = hex::decode(&keystore.nonce).map_err(|_| WalletError::InvalidFormat)?;
    let ciphertext = hex::decode(&keystore.ciphertext).map_err(|_| WalletError::InvalidFormat)?;

    let key = derive_key(password.as_bytes(), &salt, &keystore.kdf_params)?;
    let cipher = Aes256Gcm::new_from_slice(&key)
        .map_err(|e| WalletError::MnemonicGeneration(format!("aes init: {e}")))?;

    let nonce = Nonce::from_slice(&nonce_bytes);
    let plaintext = cipher
        .decrypt(nonce, ciphertext.as_ref())
        .map_err(|_| WalletError::DecryptionFailed)?;

    String::from_utf8(plaintext).map_err(|_| WalletError::DecryptionFailed)
}

/// Get password from environment variable or interactive prompt.
///
/// Priority: `UHASH_PASSWORD` env var -> interactive `rpassword` prompt.
#[cfg(feature = "cli")]
pub fn get_password(prompt: &str) -> Result<String, WalletError> {
    if let Ok(pw) = std::env::var("UHASH_PASSWORD") {
        return Ok(pw);
    }
    rpassword::prompt_password(prompt).map_err(|e| WalletError::FileError(std::io::Error::other(e)))
}

/// Prompt for a new password with confirmation.
///
/// Returns the confirmed password or an error if they don't match.
#[cfg(feature = "cli")]
pub fn get_new_password() -> Result<String, WalletError> {
    let pw = get_password("Enter password: ")?;
    if pw.is_empty() {
        return Err(WalletError::FileError(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "password cannot be empty",
        )));
    }
    let confirm = get_password("Confirm password: ")?;
    if pw != confirm {
        return Err(WalletError::FileError(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "passwords do not match",
        )));
    }
    Ok(pw)
}

/// A wallet containing a mnemonic and derived keys
pub struct Wallet {
    mnemonic: Mnemonic,
    signing_key: SigningKey,
    address: AccountId,
}

impl Wallet {
    /// Create a new wallet with a random mnemonic
    pub fn new() -> Result<Self, WalletError> {
        let mut entropy = [0u8; 32];
        getrandom::getrandom(&mut entropy)
            .map_err(|e| WalletError::MnemonicGeneration(e.to_string()))?;

        let mnemonic = Mnemonic::from_entropy_in(Language::English, &entropy)
            .map_err(|e| WalletError::MnemonicGeneration(e.to_string()))?;

        Self::from_mnemonic(mnemonic)
    }

    /// Create a wallet from an existing mnemonic phrase
    pub fn from_phrase(phrase: &str) -> Result<Self, WalletError> {
        let mnemonic = Mnemonic::parse_in(Language::English, phrase)
            .map_err(|e| WalletError::InvalidMnemonic(e.to_string()))?;

        Self::from_mnemonic(mnemonic)
    }

    fn from_mnemonic(mnemonic: Mnemonic) -> Result<Self, WalletError> {
        let seed = mnemonic.to_seed("");

        let path: DerivationPath = DERIVATION_PATH
            .parse()
            .map_err(|e: bip32::Error| WalletError::Derivation(e.to_string()))?;

        let xprv = XPrv::derive_from_path(seed, &path)
            .map_err(|e| WalletError::Derivation(e.to_string()))?;

        let signing_key = xprv.private_key();

        let public_key = secp256k1::SigningKey::from_slice(&signing_key.to_bytes())
            .map_err(|e| WalletError::Derivation(e.to_string()))?
            .public_key();

        let address = public_key
            .account_id(BOSTROM_PREFIX)
            .map_err(|e| WalletError::Derivation(e.to_string()))?;

        Ok(Self {
            mnemonic,
            signing_key: signing_key.clone(),
            address,
        })
    }

    /// Get the mnemonic phrase
    pub fn mnemonic(&self) -> String {
        self.mnemonic.to_string()
    }

    /// Get the signing key for transaction signing
    pub fn signing_key(&self) -> &SigningKey {
        &self.signing_key
    }

    /// Get the address as a string
    pub fn address_str(&self) -> String {
        self.address.to_string()
    }

    /// Save wallet mnemonic to a file, encrypted with password (Argon2id + AES-256-GCM)
    pub fn save_to_file(&self, path: &PathBuf, password: &str) -> Result<(), WalletError> {
        let keystore = encrypt_mnemonic(&self.mnemonic(), password, &self.address_str())?;
        let json = serde_json::to_string_pretty(&keystore)
            .map_err(|e| WalletError::MnemonicGeneration(format!("json serialize: {e}")))?;
        fs::write(path, json)?;
        Ok(())
    }

    /// Load wallet from an encrypted keystore file
    pub fn load_from_file(path: &PathBuf, password: &str) -> Result<Self, WalletError> {
        let content = fs::read_to_string(path)?;
        let keystore: KeystoreFile =
            serde_json::from_str(&content).map_err(|_| WalletError::InvalidFormat)?;
        if keystore.version != KEYSTORE_VERSION {
            return Err(WalletError::InvalidFormat);
        }
        let phrase = decrypt_mnemonic(&keystore, password)?;
        Self::from_phrase(&phrase)
    }
}

impl Default for Wallet {
    fn default() -> Self {
        Self::new().expect("Failed to create wallet")
    }
}

/// Get the default wallet file path
#[cfg(feature = "cli")]
pub fn default_wallet_path() -> PathBuf {
    let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
    home.join(".uhash").join("wallet.json")
}

/// Ensure the wallet directory exists
#[cfg(feature = "cli")]
pub fn ensure_wallet_dir() -> Result<PathBuf, WalletError> {
    let wallet_path = default_wallet_path();
    if let Some(parent) = wallet_path.parent() {
        fs::create_dir_all(parent)?;
    }
    Ok(wallet_path)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::NamedTempFile;

    #[test]
    fn test_new_wallet() {
        let wallet = Wallet::new().unwrap();
        let phrase = wallet.mnemonic();

        assert_eq!(phrase.split_whitespace().count(), 24);
        assert!(wallet.address_str().starts_with("bostrom"));
    }

    #[test]
    fn test_wallet_from_phrase() {
        let wallet1 = Wallet::new().unwrap();
        let phrase = wallet1.mnemonic();

        let wallet2 = Wallet::from_phrase(&phrase).unwrap();
        assert_eq!(wallet1.address_str(), wallet2.address_str());
    }

    #[test]
    fn test_deterministic_derivation() {
        let phrase = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about";
        let wallet1 = Wallet::from_phrase(phrase).unwrap();
        let wallet2 = Wallet::from_phrase(phrase).unwrap();
        assert_eq!(wallet1.address_str(), wallet2.address_str());
        assert_eq!(
            wallet1.signing_key().to_bytes(),
            wallet2.signing_key().to_bytes()
        );
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let wallet = Wallet::new().unwrap();
        let password = "test-password-123";

        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();

        wallet.save_to_file(&path, password).unwrap();

        let content = fs::read_to_string(&path).unwrap();
        let keystore: KeystoreFile = serde_json::from_str(&content).unwrap();
        assert_eq!(keystore.version, KEYSTORE_VERSION);
        assert_eq!(keystore.kdf, "argon2id");
        assert_eq!(keystore.address, wallet.address_str());

        let loaded = Wallet::load_from_file(&path, password).unwrap();
        assert_eq!(loaded.address_str(), wallet.address_str());
        assert_eq!(loaded.mnemonic(), wallet.mnemonic());

        drop(tmp);
    }

    #[test]
    fn test_wrong_password_fails() {
        let wallet = Wallet::new().unwrap();

        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();

        wallet.save_to_file(&path, "correct-password").unwrap();

        let result = Wallet::load_from_file(&path, "wrong-password");
        assert!(result.is_err());
        match result.err().unwrap() {
            WalletError::DecryptionFailed => {}
            other => panic!("expected DecryptionFailed, got: {other}"),
        }

        drop(tmp);
    }

    #[test]
    fn test_corrupted_file_fails() {
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();

        std::fs::write(&path, b"not valid json").unwrap();

        let result = Wallet::load_from_file(&path, "any-password");
        assert!(matches!(result.err().unwrap(), WalletError::InvalidFormat));

        drop(tmp);
    }
}