Skip to main content

uhash/
wallet.rs

1//! Wallet management for UniversalHash miner
2//!
3//! Handles mnemonic generation, import/export, and transaction signing.
4//! Mnemonics are encrypted at rest using Argon2id + AES-256-GCM.
5
6use aes_gcm::aead::Aead;
7use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
8use argon2::Argon2;
9use bip32::secp256k1::ecdsa::SigningKey;
10use bip32::{DerivationPath, XPrv};
11use bip39::{Language, Mnemonic};
12use cosmrs::crypto::secp256k1;
13use cosmrs::AccountId;
14use rand::RngCore;
15use serde::{Deserialize, Serialize};
16use std::fs;
17use std::path::PathBuf;
18use thiserror::Error;
19
20/// Default derivation path for Cosmos SDK chains
21const DERIVATION_PATH: &str = "m/44'/118'/0'/0/0";
22
23/// Bostrom address prefix
24const BOSTROM_PREFIX: &str = "bostrom";
25
26/// Current keystore file format version
27const KEYSTORE_VERSION: u32 = 1;
28
29#[derive(Error, Debug)]
30pub enum WalletError {
31    #[error("Failed to generate mnemonic: {0}")]
32    MnemonicGeneration(String),
33
34    #[error("Invalid mnemonic phrase: {0}")]
35    InvalidMnemonic(String),
36
37    #[error("Derivation error: {0}")]
38    Derivation(String),
39
40    #[error("File I/O error: {0}")]
41    FileError(#[from] std::io::Error),
42
43    #[error("Invalid wallet file format")]
44    InvalidFormat,
45
46    #[error("Decryption failed (wrong password or corrupted file)")]
47    DecryptionFailed,
48}
49
50/// Argon2id KDF parameters stored alongside the ciphertext
51#[derive(Serialize, Deserialize, Clone, Debug)]
52struct KdfParams {
53    m_cost: u32,
54    t_cost: u32,
55    p_cost: u32,
56}
57
58/// Encrypted wallet keystore file format (JSON)
59#[derive(Serialize, Deserialize, Debug)]
60struct KeystoreFile {
61    version: u32,
62    address: String,
63    kdf: String,
64    kdf_params: KdfParams,
65    /// Hex-encoded 16-byte salt
66    salt: String,
67    /// Hex-encoded 12-byte nonce
68    nonce: String,
69    /// Hex-encoded ciphertext (encrypted mnemonic + GCM auth tag)
70    ciphertext: String,
71}
72
73/// Derive a 32-byte AES key from password + salt using Argon2id
74fn derive_key(password: &[u8], salt: &[u8], params: &KdfParams) -> Result<[u8; 32], WalletError> {
75    let argon2 = Argon2::new(
76        argon2::Algorithm::Argon2id,
77        argon2::Version::V0x13,
78        argon2::Params::new(params.m_cost, params.t_cost, params.p_cost, Some(32))
79            .map_err(|e| WalletError::MnemonicGeneration(format!("argon2 params: {e}")))?,
80    );
81    let mut key = [0u8; 32];
82    argon2
83        .hash_password_into(password, salt, &mut key)
84        .map_err(|e| WalletError::MnemonicGeneration(format!("argon2 hash: {e}")))?;
85    Ok(key)
86}
87
88/// Encrypt plaintext mnemonic with the derived key
89fn encrypt_mnemonic(
90    mnemonic: &str,
91    password: &str,
92    address: &str,
93) -> Result<KeystoreFile, WalletError> {
94    let kdf_params = KdfParams {
95        m_cost: 19456,
96        t_cost: 2,
97        p_cost: 1,
98    };
99
100    let mut salt = [0u8; 16];
101    rand::thread_rng().fill_bytes(&mut salt);
102
103    let mut nonce_bytes = [0u8; 12];
104    rand::thread_rng().fill_bytes(&mut nonce_bytes);
105
106    let key = derive_key(password.as_bytes(), &salt, &kdf_params)?;
107    let cipher = Aes256Gcm::new_from_slice(&key)
108        .map_err(|e| WalletError::MnemonicGeneration(format!("aes init: {e}")))?;
109
110    let nonce = Nonce::from_slice(&nonce_bytes);
111    let ciphertext = cipher
112        .encrypt(nonce, mnemonic.as_bytes())
113        .map_err(|e| WalletError::MnemonicGeneration(format!("encrypt: {e}")))?;
114
115    Ok(KeystoreFile {
116        version: KEYSTORE_VERSION,
117        address: address.to_string(),
118        kdf: "argon2id".to_string(),
119        kdf_params,
120        salt: hex::encode(salt),
121        nonce: hex::encode(nonce_bytes),
122        ciphertext: hex::encode(ciphertext),
123    })
124}
125
126/// Decrypt a keystore file back to the mnemonic phrase
127fn decrypt_mnemonic(keystore: &KeystoreFile, password: &str) -> Result<String, WalletError> {
128    let salt = hex::decode(&keystore.salt).map_err(|_| WalletError::InvalidFormat)?;
129    let nonce_bytes = hex::decode(&keystore.nonce).map_err(|_| WalletError::InvalidFormat)?;
130    let ciphertext = hex::decode(&keystore.ciphertext).map_err(|_| WalletError::InvalidFormat)?;
131
132    let key = derive_key(password.as_bytes(), &salt, &keystore.kdf_params)?;
133    let cipher = Aes256Gcm::new_from_slice(&key)
134        .map_err(|e| WalletError::MnemonicGeneration(format!("aes init: {e}")))?;
135
136    let nonce = Nonce::from_slice(&nonce_bytes);
137    let plaintext = cipher
138        .decrypt(nonce, ciphertext.as_ref())
139        .map_err(|_| WalletError::DecryptionFailed)?;
140
141    String::from_utf8(plaintext).map_err(|_| WalletError::DecryptionFailed)
142}
143
144/// Get password from environment variable or interactive prompt.
145///
146/// Priority: `UHASH_PASSWORD` env var -> interactive `rpassword` prompt.
147#[cfg(feature = "cli")]
148pub fn get_password(prompt: &str) -> Result<String, WalletError> {
149    if let Ok(pw) = std::env::var("UHASH_PASSWORD") {
150        return Ok(pw);
151    }
152    rpassword::prompt_password(prompt).map_err(|e| WalletError::FileError(std::io::Error::other(e)))
153}
154
155/// Prompt for a new password with confirmation.
156///
157/// Returns the confirmed password or an error if they don't match.
158#[cfg(feature = "cli")]
159pub fn get_new_password() -> Result<String, WalletError> {
160    let pw = get_password("Enter password: ")?;
161    if pw.is_empty() {
162        return Err(WalletError::FileError(std::io::Error::new(
163            std::io::ErrorKind::InvalidInput,
164            "password cannot be empty",
165        )));
166    }
167    let confirm = get_password("Confirm password: ")?;
168    if pw != confirm {
169        return Err(WalletError::FileError(std::io::Error::new(
170            std::io::ErrorKind::InvalidInput,
171            "passwords do not match",
172        )));
173    }
174    Ok(pw)
175}
176
177/// A wallet containing a mnemonic and derived keys
178pub struct Wallet {
179    mnemonic: Mnemonic,
180    signing_key: SigningKey,
181    address: AccountId,
182}
183
184impl Wallet {
185    /// Create a new wallet with a random mnemonic
186    pub fn new() -> Result<Self, WalletError> {
187        let mut entropy = [0u8; 32];
188        getrandom::getrandom(&mut entropy)
189            .map_err(|e| WalletError::MnemonicGeneration(e.to_string()))?;
190
191        let mnemonic = Mnemonic::from_entropy_in(Language::English, &entropy)
192            .map_err(|e| WalletError::MnemonicGeneration(e.to_string()))?;
193
194        Self::from_mnemonic(mnemonic)
195    }
196
197    /// Create a wallet from an existing mnemonic phrase
198    pub fn from_phrase(phrase: &str) -> Result<Self, WalletError> {
199        let mnemonic = Mnemonic::parse_in(Language::English, phrase)
200            .map_err(|e| WalletError::InvalidMnemonic(e.to_string()))?;
201
202        Self::from_mnemonic(mnemonic)
203    }
204
205    fn from_mnemonic(mnemonic: Mnemonic) -> Result<Self, WalletError> {
206        let seed = mnemonic.to_seed("");
207
208        let path: DerivationPath = DERIVATION_PATH
209            .parse()
210            .map_err(|e: bip32::Error| WalletError::Derivation(e.to_string()))?;
211
212        let xprv = XPrv::derive_from_path(seed, &path)
213            .map_err(|e| WalletError::Derivation(e.to_string()))?;
214
215        let signing_key = xprv.private_key();
216
217        let public_key = secp256k1::SigningKey::from_slice(&signing_key.to_bytes())
218            .map_err(|e| WalletError::Derivation(e.to_string()))?
219            .public_key();
220
221        let address = public_key
222            .account_id(BOSTROM_PREFIX)
223            .map_err(|e| WalletError::Derivation(e.to_string()))?;
224
225        Ok(Self {
226            mnemonic,
227            signing_key: signing_key.clone(),
228            address,
229        })
230    }
231
232    /// Get the mnemonic phrase
233    pub fn mnemonic(&self) -> String {
234        self.mnemonic.to_string()
235    }
236
237    /// Get the signing key for transaction signing
238    pub fn signing_key(&self) -> &SigningKey {
239        &self.signing_key
240    }
241
242    /// Get the address as a string
243    pub fn address_str(&self) -> String {
244        self.address.to_string()
245    }
246
247    /// Save wallet mnemonic to a file, encrypted with password (Argon2id + AES-256-GCM)
248    pub fn save_to_file(&self, path: &PathBuf, password: &str) -> Result<(), WalletError> {
249        let keystore = encrypt_mnemonic(&self.mnemonic(), password, &self.address_str())?;
250        let json = serde_json::to_string_pretty(&keystore)
251            .map_err(|e| WalletError::MnemonicGeneration(format!("json serialize: {e}")))?;
252        fs::write(path, json)?;
253        Ok(())
254    }
255
256    /// Load wallet from an encrypted keystore file
257    pub fn load_from_file(path: &PathBuf, password: &str) -> Result<Self, WalletError> {
258        let content = fs::read_to_string(path)?;
259        let keystore: KeystoreFile =
260            serde_json::from_str(&content).map_err(|_| WalletError::InvalidFormat)?;
261        if keystore.version != KEYSTORE_VERSION {
262            return Err(WalletError::InvalidFormat);
263        }
264        let phrase = decrypt_mnemonic(&keystore, password)?;
265        Self::from_phrase(&phrase)
266    }
267}
268
269impl Default for Wallet {
270    fn default() -> Self {
271        Self::new().expect("Failed to create wallet")
272    }
273}
274
275/// Get the default wallet file path
276#[cfg(feature = "cli")]
277pub fn default_wallet_path() -> PathBuf {
278    let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
279    home.join(".uhash").join("wallet.json")
280}
281
282/// Ensure the wallet directory exists
283#[cfg(feature = "cli")]
284pub fn ensure_wallet_dir() -> Result<PathBuf, WalletError> {
285    let wallet_path = default_wallet_path();
286    if let Some(parent) = wallet_path.parent() {
287        fs::create_dir_all(parent)?;
288    }
289    Ok(wallet_path)
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295    use tempfile::NamedTempFile;
296
297    #[test]
298    fn test_new_wallet() {
299        let wallet = Wallet::new().unwrap();
300        let phrase = wallet.mnemonic();
301
302        assert_eq!(phrase.split_whitespace().count(), 24);
303        assert!(wallet.address_str().starts_with("bostrom"));
304    }
305
306    #[test]
307    fn test_wallet_from_phrase() {
308        let wallet1 = Wallet::new().unwrap();
309        let phrase = wallet1.mnemonic();
310
311        let wallet2 = Wallet::from_phrase(&phrase).unwrap();
312        assert_eq!(wallet1.address_str(), wallet2.address_str());
313    }
314
315    #[test]
316    fn test_deterministic_derivation() {
317        let phrase = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about";
318        let wallet1 = Wallet::from_phrase(phrase).unwrap();
319        let wallet2 = Wallet::from_phrase(phrase).unwrap();
320        assert_eq!(wallet1.address_str(), wallet2.address_str());
321        assert_eq!(
322            wallet1.signing_key().to_bytes(),
323            wallet2.signing_key().to_bytes()
324        );
325    }
326
327    #[test]
328    fn test_encrypt_decrypt_roundtrip() {
329        let wallet = Wallet::new().unwrap();
330        let password = "test-password-123";
331
332        let tmp = NamedTempFile::new().unwrap();
333        let path = tmp.path().to_path_buf();
334
335        wallet.save_to_file(&path, password).unwrap();
336
337        let content = fs::read_to_string(&path).unwrap();
338        let keystore: KeystoreFile = serde_json::from_str(&content).unwrap();
339        assert_eq!(keystore.version, KEYSTORE_VERSION);
340        assert_eq!(keystore.kdf, "argon2id");
341        assert_eq!(keystore.address, wallet.address_str());
342
343        let loaded = Wallet::load_from_file(&path, password).unwrap();
344        assert_eq!(loaded.address_str(), wallet.address_str());
345        assert_eq!(loaded.mnemonic(), wallet.mnemonic());
346
347        drop(tmp);
348    }
349
350    #[test]
351    fn test_wrong_password_fails() {
352        let wallet = Wallet::new().unwrap();
353
354        let tmp = NamedTempFile::new().unwrap();
355        let path = tmp.path().to_path_buf();
356
357        wallet.save_to_file(&path, "correct-password").unwrap();
358
359        let result = Wallet::load_from_file(&path, "wrong-password");
360        assert!(result.is_err());
361        match result.err().unwrap() {
362            WalletError::DecryptionFailed => {}
363            other => panic!("expected DecryptionFailed, got: {other}"),
364        }
365
366        drop(tmp);
367    }
368
369    #[test]
370    fn test_corrupted_file_fails() {
371        let tmp = NamedTempFile::new().unwrap();
372        let path = tmp.path().to_path_buf();
373
374        std::fs::write(&path, b"not valid json").unwrap();
375
376        let result = Wallet::load_from_file(&path, "any-password");
377        assert!(matches!(result.err().unwrap(), WalletError::InvalidFormat));
378
379        drop(tmp);
380    }
381}