1use 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
20const DERIVATION_PATH: &str = "m/44'/118'/0'/0/0";
22
23const BOSTROM_PREFIX: &str = "bostrom";
25
26const 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#[derive(Serialize, Deserialize, Clone, Debug)]
52struct KdfParams {
53 m_cost: u32,
54 t_cost: u32,
55 p_cost: u32,
56}
57
58#[derive(Serialize, Deserialize, Debug)]
60struct KeystoreFile {
61 version: u32,
62 address: String,
63 kdf: String,
64 kdf_params: KdfParams,
65 salt: String,
67 nonce: String,
69 ciphertext: String,
71}
72
73fn 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
88fn 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
126fn 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#[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#[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
177pub struct Wallet {
179 mnemonic: Mnemonic,
180 signing_key: SigningKey,
181 address: AccountId,
182}
183
184impl Wallet {
185 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 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 pub fn mnemonic(&self) -> String {
234 self.mnemonic.to_string()
235 }
236
237 pub fn signing_key(&self) -> &SigningKey {
239 &self.signing_key
240 }
241
242 pub fn address_str(&self) -> String {
244 self.address.to_string()
245 }
246
247 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 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#[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#[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}