1pub mod keystore;
9pub mod password;
10
11use crate::error::{Result, WalletError};
12pub use keystore::{Keystore, QuantumKeyPair, WalletData};
13use qp_rusty_crystals_hdwallet::{generate_mnemonic, HDLattice};
14use rand::{rng, RngCore};
15use serde::{Deserialize, Serialize};
16use sp_core::crypto::Ss58Codec;
17use sp_runtime::traits::IdentifyAccount;
18#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct WalletInfo {
21 pub name: String,
22 pub address: String,
23 pub created_at: chrono::DateTime<chrono::Utc>,
24 pub key_type: String,
25}
26
27pub struct WalletManager {
29 wallets_dir: std::path::PathBuf,
30}
31
32impl WalletManager {
33 pub fn new() -> Result<Self> {
35 let wallets_dir = dirs::home_dir()
36 .ok_or(WalletError::KeyGeneration)?
37 .join(".quantus")
38 .join("wallets");
39
40 std::fs::create_dir_all(&wallets_dir)?;
42
43 Ok(Self { wallets_dir })
44 }
45
46 pub async fn create_wallet(&self, name: &str, password: Option<&str>) -> Result<WalletInfo> {
48 let keystore = Keystore::new(&self.wallets_dir);
50 if keystore.load_wallet(name)?.is_some() {
51 return Err(WalletError::AlreadyExists.into());
52 }
53
54 let mut seed = [0u8; 32];
56 rng().fill_bytes(&mut seed);
57 let mnemonic = generate_mnemonic(24, seed).map_err(|_| WalletError::KeyGeneration)?;
58 let lattice =
59 HDLattice::from_mnemonic(&mnemonic, None).expect("Failed to generate lattice");
60 let dilithium_keypair = lattice.generate_keys();
61 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
62
63 let mut metadata = std::collections::HashMap::new();
65 metadata.insert("version".to_string(), "1.0.0".to_string());
66 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
67
68 let address = quantum_keypair.to_account_id_ss58check();
70
71 let wallet_data = WalletData {
72 name: name.to_string(),
73 keypair: quantum_keypair,
74 mnemonic: Some(mnemonic.clone()),
75 metadata,
76 };
77
78 let password = password.unwrap_or(""); let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, password)?;
81 keystore.save_wallet(&encrypted_wallet)?;
82
83 Ok(WalletInfo {
84 name: name.to_string(),
85 address,
86 created_at: encrypted_wallet.created_at,
87 key_type: "Dilithium ML-DSA-87".to_string(),
88 })
89 }
90
91 pub async fn create_developer_wallet(&self, name: &str) -> Result<WalletInfo> {
93 let keystore = Keystore::new(&self.wallets_dir);
95
96 let resonance_pair = match name {
98 "crystal_alice" => qp_dilithium_crypto::crystal_alice(),
99 "crystal_bob" => qp_dilithium_crypto::dilithium_bob(),
100 "crystal_charlie" => qp_dilithium_crypto::crystal_charlie(),
101 _ => return Err(WalletError::KeyGeneration.into()),
102 };
103
104 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
105
106 println!("🔑 Resonance pair: {:?}", resonance_pair.public().into_account().to_ss58check());
107 println!("🔑 Quantum keypair: {:?}", quantum_keypair.to_account_id_ss58check());
108
109 let mut metadata = std::collections::HashMap::new();
111 metadata.insert("version".to_string(), "1.0.0".to_string());
112 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
113 metadata.insert("test_wallet".to_string(), "true".to_string());
114
115 let address = quantum_keypair.to_account_id_ss58check();
117
118 let wallet_data = WalletData {
119 name: name.to_string(),
120 keypair: quantum_keypair,
121 mnemonic: None, metadata,
123 };
124
125 let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, "")?;
127 keystore.save_wallet(&encrypted_wallet)?;
128
129 Ok(WalletInfo {
130 name: name.to_string(),
131 address,
132 created_at: encrypted_wallet.created_at,
133 key_type: "Dilithium ML-DSA-87".to_string(),
134 })
135 }
136
137 pub fn export_mnemonic(&self, name: &str, password: Option<&str>) -> Result<String> {
139 let final_password = password::get_wallet_password(name, password.map(String::from), None)?;
140
141 let wallet_data = self.load_wallet(name, &final_password)?;
142
143 wallet_data.mnemonic.ok_or_else(|| WalletError::MnemonicNotAvailable.into())
144 }
145
146 pub fn list_wallets(&self) -> Result<Vec<WalletInfo>> {
148 let keystore = Keystore::new(&self.wallets_dir);
149 let wallet_names = keystore.list_wallets()?;
150
151 let mut wallets = Vec::new();
152 for name in wallet_names {
153 if let Some(encrypted_wallet) = keystore.load_wallet(&name)? {
154 let wallet_info = WalletInfo {
156 name: encrypted_wallet.name,
157 address: encrypted_wallet.address, created_at: encrypted_wallet.created_at,
159 key_type: "Dilithium ML-DSA-87".to_string(),
160 };
161 wallets.push(wallet_info);
162 }
163 }
164
165 wallets.sort_by(|a, b| b.created_at.cmp(&a.created_at));
167 Ok(wallets)
168 }
169
170 pub async fn import_wallet(
172 &self,
173 name: &str,
174 mnemonic: &str,
175 password: Option<&str>,
176 ) -> Result<WalletInfo> {
177 let keystore = Keystore::new(&self.wallets_dir);
179 if keystore.load_wallet(name)?.is_some() {
180 return Err(WalletError::AlreadyExists.into());
181 }
182
183 let lattice =
185 HDLattice::from_mnemonic(mnemonic, None).map_err(|_| WalletError::InvalidMnemonic)?;
186 let dilithium_keypair = lattice.generate_keys();
187 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
188
189 let mut metadata = std::collections::HashMap::new();
191 metadata.insert("version".to_string(), "1.0.0".to_string());
192 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
193 metadata.insert("imported".to_string(), "true".to_string());
194
195 let address = quantum_keypair.to_account_id_ss58check();
197
198 let wallet_data = WalletData {
199 name: name.to_string(),
200 keypair: quantum_keypair,
201 mnemonic: Some(mnemonic.to_string()),
202 metadata,
203 };
204
205 let password = password.unwrap_or(""); let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, password)?;
208 keystore.save_wallet(&encrypted_wallet)?;
209
210 Ok(WalletInfo {
211 name: name.to_string(),
212 address,
213 created_at: encrypted_wallet.created_at,
214 key_type: "Dilithium ML-DSA-87".to_string(),
215 })
216 }
217
218 pub fn get_wallet(&self, name: &str, password: Option<&str>) -> Result<Option<WalletInfo>> {
220 let keystore = Keystore::new(&self.wallets_dir);
221
222 if let Some(encrypted_wallet) = keystore.load_wallet(name)? {
223 if let Some(pwd) = password {
224 match keystore.decrypt_wallet_data(&encrypted_wallet, pwd) {
226 Ok(wallet_data) => {
227 let address = wallet_data.keypair.to_account_id_ss58check();
228 Ok(Some(WalletInfo {
229 name: wallet_data.name,
230 address,
231 created_at: encrypted_wallet.created_at,
232 key_type: "Dilithium ML-DSA-87".to_string(),
233 }))
234 },
235 Err(_) => {
236 Ok(Some(WalletInfo {
238 name: encrypted_wallet.name,
239 address: "[Wrong password]".to_string(),
240 created_at: encrypted_wallet.created_at,
241 key_type: "Dilithium ML-DSA-87".to_string(),
242 }))
243 },
244 }
245 } else {
246 Ok(Some(WalletInfo {
248 name: encrypted_wallet.name,
249 address: encrypted_wallet.address, created_at: encrypted_wallet.created_at,
251 key_type: "Dilithium ML-DSA-87".to_string(),
252 }))
253 }
254 } else {
255 Ok(None)
256 }
257 }
258
259 pub fn load_wallet(&self, name: &str, password: &str) -> Result<WalletData> {
261 let keystore = Keystore::new(&self.wallets_dir);
262
263 let encrypted_wallet = keystore.load_wallet(name)?.ok_or(WalletError::NotFound)?;
265
266 let wallet_data = keystore.decrypt_wallet_data(&encrypted_wallet, password)?;
268
269 Ok(wallet_data)
270 }
271
272 pub fn delete_wallet(&self, name: &str) -> Result<bool> {
274 let keystore = Keystore::new(&self.wallets_dir);
275 keystore.delete_wallet(name)
276 }
277
278 pub fn find_wallet_address(&self, name: &str) -> Result<Option<String>> {
280 let keystore = Keystore::new(&self.wallets_dir);
281
282 if let Some(encrypted_wallet) = keystore.load_wallet(name)? {
283 Ok(Some(encrypted_wallet.address))
285 } else {
286 Ok(None)
287 }
288 }
289}
290
291pub fn load_keypair_from_wallet(
292 wallet_name: &str,
293 password: Option<String>,
294 password_file: Option<String>,
295) -> Result<QuantumKeyPair> {
296 let wallet_manager = WalletManager::new()?;
297 let wallet_password = password::get_wallet_password(wallet_name, password, password_file)?;
298 let wallet_data = wallet_manager.load_wallet(wallet_name, &wallet_password)?;
299 let keypair = wallet_data.keypair;
300 Ok(keypair)
301}
302
303#[cfg(test)]
304mod tests {
305 use super::*;
306 use std::fs;
307 use tempfile::TempDir;
308
309 async fn create_test_wallet_manager() -> (WalletManager, TempDir) {
310 let temp_dir = TempDir::new().expect("Failed to create temp directory");
311 let wallets_dir = temp_dir.path().join("wallets");
312 fs::create_dir_all(&wallets_dir).expect("Failed to create wallets directory");
313
314 let wallet_manager = WalletManager { wallets_dir };
315
316 (wallet_manager, temp_dir)
317 }
318
319 #[tokio::test]
320 async fn test_wallet_creation() {
321 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
322
323 let wallet_info = wallet_manager
325 .create_wallet("test-wallet", Some("test-password"))
326 .await
327 .expect("Failed to create wallet");
328
329 assert_eq!(wallet_info.name, "test-wallet");
331 assert!(wallet_info.address.starts_with("qz")); assert_eq!(wallet_info.key_type, "Dilithium ML-DSA-87");
333 assert!(wallet_info.created_at <= chrono::Utc::now());
334 }
335
336 #[tokio::test]
337 async fn test_wallet_already_exists() {
338 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
339
340 wallet_manager
342 .create_wallet("duplicate-wallet", None)
343 .await
344 .expect("Failed to create first wallet");
345
346 let result = wallet_manager.create_wallet("duplicate-wallet", None).await;
348
349 assert!(result.is_err());
350 match result.unwrap_err() {
351 crate::error::QuantusError::Wallet(WalletError::AlreadyExists) => {},
352 _ => panic!("Expected AlreadyExists error"),
353 }
354 }
355
356 #[tokio::test]
357 async fn test_wallet_file_creation() {
358 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
359
360 let _ = wallet_manager
362 .create_wallet("file-test-wallet", Some("password123"))
363 .await
364 .expect("Failed to create wallet");
365
366 let wallet_file = wallet_manager.wallets_dir.join("file-test-wallet.json");
368 assert!(wallet_file.exists(), "Wallet file should exist");
369
370 let file_size = fs::metadata(&wallet_file).expect("Failed to get file metadata").len();
372 assert!(file_size > 0, "Wallet file should not be empty");
373 }
374
375 #[tokio::test]
376 async fn test_keystore_encryption_decryption() {
377 let temp_dir = TempDir::new().expect("Failed to create temp directory");
378 let keystore = keystore::Keystore::new(temp_dir.path());
379
380 let entropy = [1u8; 32]; let dilithium_keypair =
383 qp_rusty_crystals_dilithium::ml_dsa_87::Keypair::generate(Some(&entropy));
384 let quantum_keypair = keystore::QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
385
386 let mut metadata = std::collections::HashMap::new();
387 metadata.insert("test_key".to_string(), "test_value".to_string());
388
389 let original_wallet_data = keystore::WalletData {
390 name: "test-wallet".to_string(),
391 keypair: quantum_keypair,
392 mnemonic: Some(
393 "test mnemonic phrase with twenty four words here for testing purposes only"
394 .to_string(),
395 ),
396 metadata,
397 };
398
399 let encrypted_wallet = keystore
401 .encrypt_wallet_data(&original_wallet_data, "test-password")
402 .expect("Failed to encrypt wallet data");
403
404 assert_eq!(encrypted_wallet.name, "test-wallet");
405 assert!(!encrypted_wallet.encrypted_data.is_empty());
406 assert!(!encrypted_wallet.argon2_salt.is_empty());
407 assert!(!encrypted_wallet.aes_nonce.is_empty());
408
409 let decrypted_wallet_data = keystore
411 .decrypt_wallet_data(&encrypted_wallet, "test-password")
412 .expect("Failed to decrypt wallet data");
413
414 assert_eq!(decrypted_wallet_data.name, original_wallet_data.name);
416 assert_eq!(decrypted_wallet_data.mnemonic, original_wallet_data.mnemonic);
417 assert_eq!(decrypted_wallet_data.metadata, original_wallet_data.metadata);
418 assert_eq!(
419 decrypted_wallet_data.keypair.public_key,
420 original_wallet_data.keypair.public_key
421 );
422 assert_eq!(
423 decrypted_wallet_data.keypair.private_key,
424 original_wallet_data.keypair.private_key
425 );
426 }
427
428 #[tokio::test]
429 async fn test_quantum_keypair_address_generation() {
430 let entropy = [2u8; 32]; let dilithium_keypair =
433 qp_rusty_crystals_dilithium::ml_dsa_87::Keypair::generate(Some(&entropy));
434 let quantum_keypair = keystore::QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
435
436 let account_id = quantum_keypair.to_account_id_32();
438 let ss58_address = quantum_keypair.to_account_id_ss58check();
439
440 assert!(ss58_address.starts_with("qz"), "SS58 address should start with 5");
442 assert!(ss58_address.len() >= 47, "SS58 address should be at least 47 characters");
443
444 let converted_account_bytes = keystore::QuantumKeyPair::ss58_to_account_id(&ss58_address);
446 let account_bytes: &[u8] = account_id.as_ref();
447 assert_eq!(converted_account_bytes, account_bytes);
448 }
449
450 #[tokio::test]
451 async fn test_keystore_save_and_load() {
452 let temp_dir = TempDir::new().expect("Failed to create temp directory");
453 let keystore = keystore::Keystore::new(temp_dir.path());
454
455 let entropy = [3u8; 32]; let dilithium_keypair =
458 qp_rusty_crystals_dilithium::ml_dsa_87::Keypair::generate(Some(&entropy));
459 let quantum_keypair = keystore::QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
460
461 let wallet_data = keystore::WalletData {
462 name: "save-load-test".to_string(),
463 keypair: quantum_keypair,
464 mnemonic: Some("save load test mnemonic phrase".to_string()),
465 metadata: std::collections::HashMap::new(),
466 };
467
468 let encrypted_wallet = keystore
469 .encrypt_wallet_data(&wallet_data, "save-load-password")
470 .expect("Failed to encrypt wallet");
471
472 keystore.save_wallet(&encrypted_wallet).expect("Failed to save wallet");
474
475 let loaded_wallet = keystore
477 .load_wallet("save-load-test")
478 .expect("Failed to load wallet")
479 .expect("Wallet should exist");
480
481 assert_eq!(loaded_wallet.name, encrypted_wallet.name);
483 assert_eq!(loaded_wallet.encrypted_data, encrypted_wallet.encrypted_data);
484 assert_eq!(loaded_wallet.argon2_salt, encrypted_wallet.argon2_salt);
485 assert_eq!(loaded_wallet.aes_nonce, encrypted_wallet.aes_nonce);
486
487 let non_existent = keystore
489 .load_wallet("non-existent-wallet")
490 .expect("Load should succeed but return None");
491 assert!(non_existent.is_none());
492 }
493
494 #[tokio::test]
495 async fn test_mnemonic_generation_and_key_derivation() {
496 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
497
498 let wallet1 = wallet_manager
500 .create_wallet("mnemonic-test-1", None)
501 .await
502 .expect("Failed to create wallet 1");
503
504 let wallet2 = wallet_manager
505 .create_wallet("mnemonic-test-2", None)
506 .await
507 .expect("Failed to create wallet 2");
508
509 assert_ne!(wallet1.address, wallet2.address);
511
512 assert!(wallet1.address.starts_with("qz"));
514 assert!(wallet2.address.starts_with("qz"));
515 }
516
517 #[tokio::test]
518 async fn test_wallet_import() {
519 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
520
521 let test_mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon art";
523
524 let imported_wallet = wallet_manager
526 .import_wallet("imported-test-wallet", test_mnemonic, Some("import-password"))
527 .await
528 .expect("Failed to import wallet");
529
530 assert_eq!(imported_wallet.name, "imported-test-wallet");
532 assert!(imported_wallet.address.starts_with("qz"));
533 assert_eq!(imported_wallet.key_type, "Dilithium ML-DSA-87");
534
535 let imported_wallet2 = wallet_manager
537 .import_wallet("imported-test-wallet-2", test_mnemonic, None)
538 .await
539 .expect("Failed to import wallet again");
540
541 assert_eq!(imported_wallet.address, imported_wallet2.address);
542 }
543
544 #[tokio::test]
545 async fn test_known_values() {
546 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
547
548 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
549 let test_mnemonic = "orchard answer curve patient visual flower maze noise retreat penalty cage small earth domain scan pitch bottom crunch theme club client swap slice raven";
550 let expected_address = "qzpKnCCUvfXQdanRBkoPVDxcXbLja9JkYzv26hTQwP9C5mZWP";
551
552 let imported_wallet = wallet_manager
553 .import_wallet("imported-test-wallet", test_mnemonic, Some("import-password"))
554 .await
555 .expect("Failed to import wallet");
556
557 assert_eq!(imported_wallet.address, expected_address);
559 }
560
561 #[tokio::test]
562 async fn test_wallet_import_invalid_mnemonic() {
563 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
564
565 let invalid_mnemonic = "invalid mnemonic phrase that should not work";
567
568 let result = wallet_manager.import_wallet("invalid-wallet", invalid_mnemonic, None).await;
569
570 assert!(result.is_err());
571 match result.unwrap_err() {
572 crate::error::QuantusError::Wallet(WalletError::InvalidMnemonic) => {},
573 _ => panic!("Expected InvalidMnemonic error"),
574 }
575 }
576
577 #[tokio::test]
578 async fn test_wallet_import_already_exists() {
579 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
580
581 let test_mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon art";
582
583 wallet_manager
585 .import_wallet("duplicate-import-wallet", test_mnemonic, None)
586 .await
587 .expect("Failed to import first wallet");
588
589 let result = wallet_manager
591 .import_wallet("duplicate-import-wallet", test_mnemonic, None)
592 .await;
593
594 assert!(result.is_err());
595 match result.unwrap_err() {
596 crate::error::QuantusError::Wallet(WalletError::AlreadyExists) => {},
597 _ => panic!("Expected AlreadyExists error"),
598 }
599 }
600
601 #[tokio::test]
602 async fn test_list_wallets() {
603 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
604
605 let wallets = wallet_manager.list_wallets().expect("Failed to list wallets");
607 assert_eq!(wallets.len(), 0);
608
609 wallet_manager
611 .create_wallet("wallet-1", Some("password1"))
612 .await
613 .expect("Failed to create wallet 1");
614
615 wallet_manager
616 .create_wallet("wallet-2", None)
617 .await
618 .expect("Failed to create wallet 2");
619
620 let test_mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon art";
621 wallet_manager
622 .import_wallet("imported-wallet", test_mnemonic, Some("password3"))
623 .await
624 .expect("Failed to import wallet");
625
626 let wallets = wallet_manager.list_wallets().expect("Failed to list wallets");
628
629 assert_eq!(wallets.len(), 3);
630
631 let wallet_names: Vec<&String> = wallets.iter().map(|w| &w.name).collect();
633 assert!(wallet_names.contains(&&"wallet-1".to_string()));
634 assert!(wallet_names.contains(&&"wallet-2".to_string()));
635 assert!(wallet_names.contains(&&"imported-wallet".to_string()));
636
637 for wallet in &wallets {
639 assert!(wallet.address.starts_with("qz")); assert_eq!(wallet.key_type, "Dilithium ML-DSA-87");
641 }
642
643 assert!(wallets[0].created_at >= wallets[1].created_at);
645 assert!(wallets[1].created_at >= wallets[2].created_at);
646 }
647
648 #[tokio::test]
649 async fn test_get_wallet() {
650 let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
651
652 let created_wallet = wallet_manager
654 .create_wallet("test-get-wallet", Some("test-password"))
655 .await
656 .expect("Failed to create wallet");
657
658 let wallet_info = wallet_manager
660 .get_wallet("test-get-wallet", None)
661 .expect("Failed to get wallet")
662 .expect("Wallet should exist");
663
664 assert_eq!(wallet_info.name, "test-get-wallet");
665 assert_eq!(wallet_info.address, created_wallet.address); let wallet_info = wallet_manager
670 .get_wallet("test-get-wallet", Some("wrong-password"))
671 .expect("Failed to get wallet")
672 .expect("Wallet should exist");
673
674 assert_eq!(wallet_info.name, "test-get-wallet");
675 assert_eq!(wallet_info.address, "[Wrong password]");
677
678 let wallet_info = wallet_manager
680 .get_wallet("test-get-wallet", Some("test-password"))
681 .expect("Failed to get wallet")
682 .expect("Wallet should exist");
683
684 assert_eq!(wallet_info.name, "test-get-wallet");
685 assert_eq!(wallet_info.address, created_wallet.address);
686 assert!(wallet_info.address.starts_with("qz"));
687
688 let result = wallet_manager
690 .get_wallet("non-existent-wallet", None)
691 .expect("Should not error on non-existent wallet");
692
693 assert!(result.is_none());
694 }
695}