Skip to main content

quantus_cli/wallet/
mod.rs

1/// Wallet management module
2///
3/// This module provides functionality for:
4/// - Creating quantum-safe wallets using Dilithium keys
5/// - Importing/exporting wallets with mnemonic phrases
6/// - Encrypting/decrypting wallet data
7/// - Managing multiple wallets
8pub mod keystore;
9pub mod password;
10
11use crate::error::{Result, WalletError};
12pub use keystore::{Keystore, QuantumKeyPair, WalletData};
13use qp_rusty_crystals_hdwallet::{derive_key_from_mnemonic, generate_mnemonic, SensitiveBytes32};
14use rand::{rng, RngCore};
15use serde::{Deserialize, Serialize};
16use sp_runtime::traits::IdentifyAccount;
17
18/// Default derivation path for Quantus wallets: m/44'/189189'/0'/0'/0'
19pub const DEFAULT_DERIVATION_PATH: &str = "m/44'/189189'/0'/0'/0'";
20
21/// Wallet information structure
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct WalletInfo {
24	pub name: String,
25	pub address: String,
26	pub created_at: chrono::DateTime<chrono::Utc>,
27	pub key_type: String,
28	pub derivation_path: String,
29}
30
31/// Main wallet manager
32pub struct WalletManager {
33	wallets_dir: std::path::PathBuf,
34}
35
36impl WalletManager {
37	/// Create a new wallet manager
38	pub fn new() -> Result<Self> {
39		let wallets_dir = dirs::home_dir()
40			.ok_or(WalletError::KeyGeneration)?
41			.join(".quantus")
42			.join("wallets");
43
44		// Create directory if it doesn't exist
45		std::fs::create_dir_all(&wallets_dir)?;
46
47		Ok(Self { wallets_dir })
48	}
49
50	/// Create a new wallet
51	pub async fn create_wallet(&self, name: &str, password: Option<&str>) -> Result<WalletInfo> {
52		self.create_wallet_with_derivation_path(name, password, DEFAULT_DERIVATION_PATH)
53			.await
54	}
55
56	/// Create a new wallet with custom derivation path
57	pub async fn create_wallet_with_derivation_path(
58		&self,
59		name: &str,
60		password: Option<&str>,
61		derivation_path: &str,
62	) -> Result<WalletInfo> {
63		// Check if wallet already exists
64		let keystore = Keystore::new(&self.wallets_dir);
65		if keystore.load_wallet(name)?.is_some() {
66			return Err(WalletError::AlreadyExists.into());
67		}
68
69		// Generate a new Dilithium keypair using derivation path
70		let mut seed = [0u8; 32];
71		rng().fill_bytes(&mut seed);
72		let sensitive_seed = SensitiveBytes32::from(&mut seed);
73		let mnemonic = generate_mnemonic(sensitive_seed).map_err(|_| WalletError::KeyGeneration)?;
74		let dilithium_keypair = derive_key_from_mnemonic(&mnemonic, None, derivation_path)
75			.map_err(|_| WalletError::KeyGeneration)?;
76		let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
77
78		// Create wallet data
79		let mut metadata = std::collections::HashMap::new();
80		metadata.insert("version".to_string(), "1.0.0".to_string());
81		metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
82		metadata.insert("derivation_path".to_string(), derivation_path.to_string());
83
84		// Generate address from public key (simplified version)
85		let address = quantum_keypair.to_account_id_ss58check();
86
87		let wallet_data = WalletData {
88			name: name.to_string(),
89			keypair: quantum_keypair,
90			mnemonic: Some(mnemonic.clone()),
91			derivation_path: derivation_path.to_string(),
92			metadata,
93		};
94
95		// Encrypt and save the wallet
96		let password = password.unwrap_or(""); // Use empty password if none provided
97		let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, password)?;
98		keystore.save_wallet(&encrypted_wallet)?;
99
100		Ok(WalletInfo {
101			name: name.to_string(),
102			address,
103			created_at: encrypted_wallet.created_at,
104			key_type: "Dilithium ML-DSA-87".to_string(),
105			derivation_path: derivation_path.to_string(),
106		})
107	}
108
109	/// Create a new developer wallet
110	pub async fn create_developer_wallet(&self, name: &str) -> Result<WalletInfo> {
111		// Check if wallet already exists
112		let keystore = Keystore::new(&self.wallets_dir);
113
114		// Generate the appropriate test keypair
115		let resonance_pair = match name {
116			"crystal_alice" => qp_dilithium_crypto::crystal_alice(),
117			"crystal_bob" => qp_dilithium_crypto::dilithium_bob(),
118			"crystal_charlie" => qp_dilithium_crypto::crystal_charlie(),
119			_ => return Err(WalletError::KeyGeneration.into()),
120		};
121
122		let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
123
124		// Format addresses with SS58 version 189 (Quantus format)
125		use sp_core::crypto::Ss58Codec;
126		let resonance_addr = resonance_pair
127			.public()
128			.into_account()
129			.to_ss58check_with_version(sp_core::crypto::Ss58AddressFormat::custom(189));
130		println!("🔑 Resonance pair: {:?}", resonance_addr);
131		println!("🔑 Quantum keypair: {:?}", quantum_keypair.to_account_id_ss58check());
132
133		// Create wallet data
134		let mut metadata = std::collections::HashMap::new();
135		metadata.insert("version".to_string(), "1.0.0".to_string());
136		metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
137		metadata.insert("test_wallet".to_string(), "true".to_string());
138
139		// Generate address from public key
140		let address = quantum_keypair.to_account_id_ss58check();
141
142		let wallet_data = WalletData {
143			name: name.to_string(),
144			keypair: quantum_keypair,
145			mnemonic: None,                    // Test wallets don't have mnemonics
146			derivation_path: "m/".to_string(), // Developer wallets use root path
147			metadata,
148		};
149
150		// Encrypt and save the wallet with empty password for test wallets
151		let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, "")?;
152		keystore.save_wallet(&encrypted_wallet)?;
153
154		Ok(WalletInfo {
155			name: name.to_string(),
156			address,
157			created_at: encrypted_wallet.created_at,
158			key_type: "Dilithium ML-DSA-87".to_string(),
159			derivation_path: "m/".to_string(),
160		})
161	}
162
163	/// Export a wallet's mnemonic phrase
164	pub fn export_mnemonic(&self, name: &str, password: Option<&str>) -> Result<String> {
165		let final_password = password::get_wallet_password(name, password.map(String::from), None)?;
166
167		let wallet_data = self.load_wallet(name, &final_password)?;
168
169		wallet_data.mnemonic.ok_or_else(|| WalletError::MnemonicNotAvailable.into())
170	}
171
172	/// List all wallets
173	pub fn list_wallets(&self) -> Result<Vec<WalletInfo>> {
174		let keystore = Keystore::new(&self.wallets_dir);
175		let wallet_names = keystore.list_wallets()?;
176
177		let mut wallets = Vec::new();
178		for name in wallet_names {
179			if let Some(encrypted_wallet) = keystore.load_wallet(&name)? {
180				// Create wallet info using stored public address
181				let wallet_info = WalletInfo {
182					name: encrypted_wallet.name,
183					address: encrypted_wallet.address, // Address is stored unencrypted
184					created_at: encrypted_wallet.created_at,
185					key_type: "Dilithium ML-DSA-87".to_string(),
186					derivation_path: "[Encrypted]".to_string(), // Derivation path is encrypted
187				};
188				wallets.push(wallet_info);
189			}
190		}
191
192		// Sort by creation date (newest first)
193		wallets.sort_by(|a, b| b.created_at.cmp(&a.created_at));
194		Ok(wallets)
195	}
196
197	/// Import wallet from mnemonic phrase
198	pub async fn import_wallet(
199		&self,
200		name: &str,
201		mnemonic: &str,
202		password: Option<&str>,
203	) -> Result<WalletInfo> {
204		self.import_wallet_with_derivation_path(name, mnemonic, password, DEFAULT_DERIVATION_PATH)
205			.await
206	}
207
208	/// Create wallet from mnemonic without derivation (master seed)
209	pub async fn create_wallet_no_derivation(
210		&self,
211		name: &str,
212		password: Option<&str>,
213	) -> Result<WalletInfo> {
214		// Check if wallet already exists
215		let keystore = Keystore::new(&self.wallets_dir);
216		if keystore.load_wallet(name)?.is_some() {
217			return Err(WalletError::AlreadyExists.into());
218		}
219
220		// Generate new mnemonic and use master seed directly (no derivation path)
221		let mut seed = [0u8; 32];
222		rng().fill_bytes(&mut seed);
223		let sensitive_seed = SensitiveBytes32::from(&mut seed);
224		let mnemonic = generate_mnemonic(sensitive_seed).map_err(|_| WalletError::KeyGeneration)?;
225		// For "no derivation" mode, we use the root path m/
226		let dilithium_keypair = derive_key_from_mnemonic(&mnemonic, None, "m/44'/189189'/0'")
227			.map_err(|_| WalletError::KeyGeneration)?;
228		let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
229
230		// Create wallet data
231		let mut metadata = std::collections::HashMap::new();
232		metadata.insert("version".to_string(), "1.0.0".to_string());
233		metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
234		metadata.insert("no_derivation".to_string(), "true".to_string());
235
236		// Generate address from public key
237		let address = quantum_keypair.to_account_id_ss58check();
238
239		let wallet_data = WalletData {
240			name: name.to_string(),
241			keypair: quantum_keypair,
242			mnemonic: Some(mnemonic),
243			derivation_path: "master".to_string(),
244			metadata,
245		};
246
247		// Encrypt and save the wallet
248		let password = password.unwrap_or(""); // Use empty password if none provided
249		let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, password)?;
250		keystore.save_wallet(&encrypted_wallet)?;
251
252		Ok(WalletInfo {
253			name: name.to_string(),
254			address,
255			created_at: chrono::Utc::now(),
256			key_type: "Dilithium ML-DSA-87".to_string(),
257			derivation_path: "master".to_string(),
258		})
259	}
260
261	/// Import wallet from mnemonic without derivation (master seed)
262	pub async fn import_wallet_no_derivation(
263		&self,
264		name: &str,
265		mnemonic: &str,
266		password: Option<&str>,
267	) -> Result<WalletInfo> {
268		// Check if wallet already exists
269		let keystore = Keystore::new(&self.wallets_dir);
270		if keystore.load_wallet(name)?.is_some() {
271			return Err(WalletError::AlreadyExists.into());
272		}
273
274		// Use mnemonic to generate keys directly (no derivation path)
275		let dilithium_keypair = derive_key_from_mnemonic(mnemonic, None, "m/44'/189189'/0'")
276			.map_err(|_| WalletError::InvalidMnemonic)?;
277		let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
278
279		// Create wallet data
280		let mut metadata = std::collections::HashMap::new();
281		metadata.insert("version".to_string(), "1.0.0".to_string());
282		metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
283		metadata.insert("imported".to_string(), "true".to_string());
284		metadata.insert("no_derivation".to_string(), "true".to_string());
285
286		// Generate address from public key
287		let address = quantum_keypair.to_account_id_ss58check();
288
289		let wallet_data = WalletData {
290			name: name.to_string(),
291			keypair: quantum_keypair,
292			mnemonic: Some(mnemonic.to_string()),
293			derivation_path: "master".to_string(),
294			metadata,
295		};
296
297		// Encrypt and save the wallet
298		let password = password.unwrap_or(""); // Use empty password if none provided
299		let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, password)?;
300		keystore.save_wallet(&encrypted_wallet)?;
301
302		Ok(WalletInfo {
303			name: name.to_string(),
304			address,
305			created_at: chrono::Utc::now(),
306			key_type: "Dilithium ML-DSA-87".to_string(),
307			derivation_path: "master".to_string(),
308		})
309	}
310
311	/// Import wallet from mnemonic phrase with custom derivation path
312	pub async fn import_wallet_with_derivation_path(
313		&self,
314		name: &str,
315		mnemonic: &str,
316		password: Option<&str>,
317		derivation_path: &str,
318	) -> Result<WalletInfo> {
319		// Check if wallet already exists
320		let keystore = Keystore::new(&self.wallets_dir);
321		if keystore.load_wallet(name)?.is_some() {
322			return Err(WalletError::AlreadyExists.into());
323		}
324
325		// Validate and import from mnemonic using derivation path
326		let dilithium_keypair = derive_key_from_mnemonic(mnemonic, None, derivation_path)
327			.map_err(|_| WalletError::InvalidMnemonic)?;
328		let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
329
330		// Create wallet data
331		let mut metadata = std::collections::HashMap::new();
332		metadata.insert("version".to_string(), "1.0.0".to_string());
333		metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
334		metadata.insert("imported".to_string(), "true".to_string());
335		metadata.insert("derivation_path".to_string(), derivation_path.to_string());
336
337		// Generate address from public key
338		let address = quantum_keypair.to_account_id_ss58check();
339
340		let wallet_data = WalletData {
341			name: name.to_string(),
342			keypair: quantum_keypair,
343			mnemonic: Some(mnemonic.to_string()),
344			derivation_path: derivation_path.to_string(),
345			metadata,
346		};
347
348		// Encrypt and save the wallet
349		let password = password.unwrap_or(""); // Use empty password if none provided
350		let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, password)?;
351		keystore.save_wallet(&encrypted_wallet)?;
352
353		Ok(WalletInfo {
354			name: name.to_string(),
355			address,
356			created_at: encrypted_wallet.created_at,
357			key_type: "Dilithium ML-DSA-87".to_string(),
358			derivation_path: derivation_path.to_string(),
359		})
360	}
361
362	/// Create wallet from 32-byte seed
363	pub async fn create_wallet_from_seed(
364		&self,
365		name: &str,
366		seed_hex: &str,
367		password: Option<&str>,
368	) -> Result<WalletInfo> {
369		// Check if wallet already exists
370		let keystore = Keystore::new(&self.wallets_dir);
371		if keystore.load_wallet(name)?.is_some() {
372			return Err(WalletError::AlreadyExists.into());
373		}
374
375		// Validate seed hex format (should be 64 hex characters for 32 bytes)
376		if seed_hex.len() != 64 {
377			return Err(WalletError::InvalidMnemonic.into()); // Reusing error type
378		}
379
380		// Convert hex to bytes
381		let seed_bytes = hex::decode(seed_hex).map_err(|_| WalletError::InvalidMnemonic)?;
382		if seed_bytes.len() != 32 {
383			return Err(WalletError::InvalidMnemonic.into());
384		}
385
386		// Create DilithiumPair from seed
387		let seed_array: [u8; 32] =
388			seed_bytes.try_into().map_err(|_| WalletError::InvalidMnemonic)?;
389
390		println!("Debug: seed_array length: {}", seed_array.len());
391		println!("Debug: seed_hex: {}", seed_hex);
392		println!("Debug: calling DilithiumPair::from_seed");
393
394		let dilithium_pair = qp_dilithium_crypto::types::DilithiumPair::from_seed(&seed_array)
395			.map_err(|e| {
396				println!("Debug: DilithiumPair::from_seed failed with error: {:?}", e);
397				WalletError::InvalidMnemonic
398			})?;
399
400		println!("Debug: DilithiumPair created successfully");
401
402		// Convert to QuantumKeyPair
403		let quantum_keypair = QuantumKeyPair::from_resonance_pair(&dilithium_pair);
404
405		// Create wallet data
406		let mut metadata = std::collections::HashMap::new();
407		metadata.insert("version".to_string(), "1.0.0".to_string());
408		metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
409		metadata.insert("from_seed".to_string(), "true".to_string());
410
411		// Generate address from public key
412		let address = quantum_keypair.to_account_id_ss58check();
413
414		let wallet_data = WalletData {
415			name: name.to_string(),
416			keypair: quantum_keypair,
417			mnemonic: None,                    // No mnemonic for seed-based wallets
418			derivation_path: "m/".to_string(), // Seed-based wallets use root path
419			metadata,
420		};
421
422		// Encrypt and save the wallet
423		let password = password.unwrap_or(""); // Use empty password if none provided
424		let encrypted_wallet = keystore.encrypt_wallet_data(&wallet_data, password)?;
425		keystore.save_wallet(&encrypted_wallet)?;
426
427		Ok(WalletInfo {
428			name: name.to_string(),
429			address,
430			created_at: encrypted_wallet.created_at,
431			key_type: "Dilithium ML-DSA-87".to_string(),
432			derivation_path: "m/".to_string(),
433		})
434	}
435
436	/// Get wallet by name with password for decryption
437	pub fn get_wallet(&self, name: &str, password: Option<&str>) -> Result<Option<WalletInfo>> {
438		let keystore = Keystore::new(&self.wallets_dir);
439
440		if let Some(encrypted_wallet) = keystore.load_wallet(name)? {
441			if let Some(pwd) = password {
442				// Decrypt and show full details
443				match keystore.decrypt_wallet_data(&encrypted_wallet, pwd) {
444					Ok(wallet_data) => {
445						let address = wallet_data.keypair.to_account_id_ss58check();
446						Ok(Some(WalletInfo {
447							name: wallet_data.name,
448							address,
449							created_at: encrypted_wallet.created_at,
450							key_type: "Dilithium ML-DSA-87".to_string(),
451							derivation_path: wallet_data.derivation_path,
452						}))
453					},
454					Err(_) => {
455						// Wrong password, return basic info
456						Ok(Some(WalletInfo {
457							name: encrypted_wallet.name,
458							address: "[Wrong password]".to_string(),
459							created_at: encrypted_wallet.created_at,
460							key_type: "Dilithium ML-DSA-87".to_string(),
461							derivation_path: "[Wrong password]".to_string(),
462						}))
463					},
464				}
465			} else {
466				// No password provided, return basic info with public address
467				Ok(Some(WalletInfo {
468					name: encrypted_wallet.name,
469					address: encrypted_wallet.address, // Address is public
470					created_at: encrypted_wallet.created_at,
471					key_type: "Dilithium ML-DSA-87".to_string(),
472					derivation_path: "[Encrypted]".to_string(), // Derivation path is encrypted
473				}))
474			}
475		} else {
476			Ok(None)
477		}
478	}
479
480	/// Load a wallet from disk and decrypt it with the provided password
481	pub fn load_wallet(&self, name: &str, password: &str) -> Result<WalletData> {
482		let keystore = Keystore::new(&self.wallets_dir);
483
484		// Load the encrypted wallet
485		let encrypted_wallet = keystore.load_wallet(name)?.ok_or(WalletError::NotFound)?;
486
487		// Decrypt the wallet data using the provided password
488		let wallet_data = keystore.decrypt_wallet_data(&encrypted_wallet, password)?;
489
490		Ok(wallet_data)
491	}
492
493	/// Delete a wallet
494	pub fn delete_wallet(&self, name: &str) -> Result<bool> {
495		let keystore = Keystore::new(&self.wallets_dir);
496		keystore.delete_wallet(name)
497	}
498
499	/// Find wallet by name and return its address
500	pub fn find_wallet_address(&self, name: &str) -> Result<Option<String>> {
501		let keystore = Keystore::new(&self.wallets_dir);
502
503		if let Some(encrypted_wallet) = keystore.load_wallet(name)? {
504			// Return the stored address (it's stored unencrypted)
505			Ok(Some(encrypted_wallet.address))
506		} else {
507			Ok(None)
508		}
509	}
510}
511
512pub fn load_keypair_from_wallet(
513	wallet_name: &str,
514	password: Option<String>,
515	password_file: Option<String>,
516) -> Result<QuantumKeyPair> {
517	let wallet_manager = WalletManager::new()?;
518	let wallet_password = password::get_wallet_password(wallet_name, password, password_file)?;
519	let wallet_data = wallet_manager.load_wallet(wallet_name, &wallet_password)?;
520	let keypair = wallet_data.keypair;
521	Ok(keypair)
522}
523
524#[cfg(test)]
525mod tests {
526	use super::*;
527	use std::fs;
528	use tempfile::TempDir;
529
530	async fn create_test_wallet_manager() -> (WalletManager, TempDir) {
531		let temp_dir = TempDir::new().expect("Failed to create temp directory");
532		let wallets_dir = temp_dir.path().join("wallets");
533		fs::create_dir_all(&wallets_dir).expect("Failed to create wallets directory");
534
535		let wallet_manager = WalletManager { wallets_dir };
536
537		(wallet_manager, temp_dir)
538	}
539
540	#[tokio::test]
541	async fn test_wallet_creation() {
542		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
543
544		// Test wallet creation
545		let wallet_info = wallet_manager
546			.create_wallet("test-wallet", Some("test-password"))
547			.await
548			.expect("Failed to create wallet");
549
550		// Verify wallet info
551		assert_eq!(wallet_info.name, "test-wallet");
552		assert!(wallet_info.address.starts_with("qz")); // SS58 addresses start with 5
553		assert_eq!(wallet_info.key_type, "Dilithium ML-DSA-87");
554		assert!(wallet_info.created_at <= chrono::Utc::now());
555	}
556
557	#[tokio::test]
558	async fn test_wallet_already_exists() {
559		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
560
561		// Create first wallet
562		wallet_manager
563			.create_wallet("duplicate-wallet", None)
564			.await
565			.expect("Failed to create first wallet");
566
567		// Try to create wallet with same name
568		let result = wallet_manager.create_wallet("duplicate-wallet", None).await;
569
570		assert!(result.is_err());
571		match result.unwrap_err() {
572			crate::error::QuantusError::Wallet(WalletError::AlreadyExists) => {},
573			_ => panic!("Expected AlreadyExists error"),
574		}
575	}
576
577	#[tokio::test]
578	async fn test_wallet_file_creation() {
579		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
580
581		// Create wallet
582		let _ = wallet_manager
583			.create_wallet("file-test-wallet", Some("password123"))
584			.await
585			.expect("Failed to create wallet");
586
587		// Check if wallet file exists
588		let wallet_file = wallet_manager.wallets_dir.join("file-test-wallet.json");
589		assert!(wallet_file.exists(), "Wallet file should exist");
590
591		// Verify file is not empty
592		let file_size = fs::metadata(&wallet_file).expect("Failed to get file metadata").len();
593		assert!(file_size > 0, "Wallet file should not be empty");
594	}
595
596	#[tokio::test]
597	async fn test_keystore_encryption_decryption() {
598		let temp_dir = TempDir::new().expect("Failed to create temp directory");
599		let keystore = keystore::Keystore::new(temp_dir.path());
600
601		// Create test wallet data
602		let mut entropy = [1u8; 32]; // Use fixed entropy for deterministic tests
603		let dilithium_keypair = qp_rusty_crystals_dilithium::ml_dsa_87::Keypair::generate(
604			SensitiveBytes32::from(&mut entropy),
605		);
606		let quantum_keypair = keystore::QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
607
608		let mut metadata = std::collections::HashMap::new();
609		metadata.insert("test_key".to_string(), "test_value".to_string());
610
611		let original_wallet_data = keystore::WalletData {
612			name: "test-wallet".to_string(),
613			keypair: quantum_keypair,
614			mnemonic: Some(
615				"test mnemonic phrase with twenty four words here for testing purposes only"
616					.to_string(),
617			),
618			derivation_path: DEFAULT_DERIVATION_PATH.to_string(),
619			metadata,
620		};
621
622		// Test encryption
623		let encrypted_wallet = keystore
624			.encrypt_wallet_data(&original_wallet_data, "test-password")
625			.expect("Failed to encrypt wallet data");
626
627		assert_eq!(encrypted_wallet.name, "test-wallet");
628		assert!(!encrypted_wallet.encrypted_data.is_empty());
629		assert!(!encrypted_wallet.argon2_salt.is_empty());
630		assert!(!encrypted_wallet.aes_nonce.is_empty());
631
632		// Test decryption
633		let decrypted_wallet_data = keystore
634			.decrypt_wallet_data(&encrypted_wallet, "test-password")
635			.expect("Failed to decrypt wallet data");
636
637		// Verify decrypted data matches original
638		assert_eq!(decrypted_wallet_data.name, original_wallet_data.name);
639		assert_eq!(decrypted_wallet_data.mnemonic, original_wallet_data.mnemonic);
640		assert_eq!(decrypted_wallet_data.metadata, original_wallet_data.metadata);
641		assert_eq!(
642			decrypted_wallet_data.keypair.public_key,
643			original_wallet_data.keypair.public_key
644		);
645		assert_eq!(
646			decrypted_wallet_data.keypair.private_key,
647			original_wallet_data.keypair.private_key
648		);
649	}
650
651	#[tokio::test]
652	async fn test_quantum_keypair_address_generation() {
653		// Generate keypair
654		let mut entropy = [2u8; 32]; // Use different entropy for variety
655		let dilithium_keypair = qp_rusty_crystals_dilithium::ml_dsa_87::Keypair::generate(
656			SensitiveBytes32::from(&mut entropy),
657		);
658		let quantum_keypair = keystore::QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
659
660		// Test address generation
661		let account_id = quantum_keypair.to_account_id_32();
662		let ss58_address = quantum_keypair.to_account_id_ss58check();
663
664		// Verify SS58 address format
665		assert!(ss58_address.starts_with("qz"), "SS58 address should start with 5");
666		assert!(ss58_address.len() >= 47, "SS58 address should be at least 47 characters");
667
668		// Test round-trip conversion
669		let converted_account_bytes = keystore::QuantumKeyPair::ss58_to_account_id(&ss58_address);
670		let account_bytes: &[u8] = account_id.as_ref();
671		assert_eq!(converted_account_bytes, account_bytes);
672	}
673
674	#[tokio::test]
675	async fn test_keystore_save_and_load() {
676		let temp_dir = TempDir::new().expect("Failed to create temp directory");
677		let keystore = keystore::Keystore::new(temp_dir.path());
678
679		// Create and encrypt wallet data
680		let mut entropy = [3u8; 32]; // Use different entropy for each test
681		let dilithium_keypair = qp_rusty_crystals_dilithium::ml_dsa_87::Keypair::generate(
682			SensitiveBytes32::from(&mut entropy),
683		);
684		let quantum_keypair = keystore::QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
685
686		let wallet_data = keystore::WalletData {
687			name: "save-load-test".to_string(),
688			keypair: quantum_keypair,
689			mnemonic: Some("save load test mnemonic phrase".to_string()),
690			derivation_path: DEFAULT_DERIVATION_PATH.to_string(),
691			metadata: std::collections::HashMap::new(),
692		};
693
694		let encrypted_wallet = keystore
695			.encrypt_wallet_data(&wallet_data, "save-load-password")
696			.expect("Failed to encrypt wallet");
697
698		// Save wallet
699		keystore.save_wallet(&encrypted_wallet).expect("Failed to save wallet");
700
701		// Load wallet
702		let loaded_wallet = keystore
703			.load_wallet("save-load-test")
704			.expect("Failed to load wallet")
705			.expect("Wallet should exist");
706
707		// Verify loaded wallet matches saved wallet
708		assert_eq!(loaded_wallet.name, encrypted_wallet.name);
709		assert_eq!(loaded_wallet.encrypted_data, encrypted_wallet.encrypted_data);
710		assert_eq!(loaded_wallet.argon2_salt, encrypted_wallet.argon2_salt);
711		assert_eq!(loaded_wallet.aes_nonce, encrypted_wallet.aes_nonce);
712
713		// Test loading non-existent wallet
714		let non_existent = keystore
715			.load_wallet("non-existent-wallet")
716			.expect("Load should succeed but return None");
717		assert!(non_existent.is_none());
718	}
719
720	#[tokio::test]
721	async fn test_mnemonic_generation_and_key_derivation() {
722		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
723
724		// Create multiple wallets to test mnemonic uniqueness
725		let wallet1 = wallet_manager
726			.create_wallet("mnemonic-test-1", None)
727			.await
728			.expect("Failed to create wallet 1");
729
730		let wallet2 = wallet_manager
731			.create_wallet("mnemonic-test-2", None)
732			.await
733			.expect("Failed to create wallet 2");
734
735		// Addresses should be different (extremely unlikely to be the same)
736		assert_ne!(wallet1.address, wallet2.address);
737
738		// Both should be valid SS58 addresses
739		assert!(wallet1.address.starts_with("qz"));
740		assert!(wallet2.address.starts_with("qz"));
741	}
742
743	#[tokio::test]
744	async fn test_wallet_import() {
745		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
746
747		// Test mnemonic phrase (24 words)
748		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";
749
750		// Import wallet
751		let imported_wallet = wallet_manager
752			.import_wallet("imported-test-wallet", test_mnemonic, Some("import-password"))
753			.await
754			.expect("Failed to import wallet");
755
756		// Verify wallet info
757		assert_eq!(imported_wallet.name, "imported-test-wallet");
758		assert!(imported_wallet.address.starts_with("qz"));
759		assert_eq!(imported_wallet.key_type, "Dilithium ML-DSA-87");
760
761		// Import the same mnemonic again should create the same address
762		let imported_wallet2 = wallet_manager
763			.import_wallet("imported-test-wallet-2", test_mnemonic, None)
764			.await
765			.expect("Failed to import wallet again");
766
767		assert_eq!(imported_wallet.address, imported_wallet2.address);
768	}
769
770	#[tokio::test]
771	async fn test_known_values() {
772		sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
773
774		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
775		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";
776		let expected_address = "qzoog56PJKvDwqo9GwkzRN74kxEgDEspxu5zVA62y18ttt3tG"; // default derivation path index 0
777		let expected_address_no_derive = "qzofkFbmnEYLX6iHwqJ9uKYXFi7ypQwcBBMxcYYLVD17vGpsm";
778
779		let imported_wallet = wallet_manager
780			.import_wallet("imported-test-wallet", test_mnemonic, Some("import-password"))
781			.await
782			.expect("Failed to import wallet");
783
784		let imported_wallet_no_derive = wallet_manager
785			.import_wallet_no_derivation(
786				"imported-test-wallet_no_derive",
787				test_mnemonic,
788				Some("import-password"),
789			)
790			.await
791			.expect("Failed to import wallet");
792
793		assert_eq!(imported_wallet.address, expected_address, "address at index 0 is wrong");
794		assert_eq!(
795			imported_wallet_no_derive.address, expected_address_no_derive,
796			"no-derivation address is wrong"
797		);
798	}
799
800	#[tokio::test]
801	async fn test_wallet_import_invalid_mnemonic() {
802		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
803
804		// Test with invalid mnemonic
805		let invalid_mnemonic = "invalid mnemonic phrase that should not work";
806
807		let result = wallet_manager.import_wallet("invalid-wallet", invalid_mnemonic, None).await;
808
809		assert!(result.is_err());
810		match result.unwrap_err() {
811			crate::error::QuantusError::Wallet(WalletError::InvalidMnemonic) => {},
812			_ => panic!("Expected InvalidMnemonic error"),
813		}
814	}
815
816	#[tokio::test]
817	async fn test_wallet_import_already_exists() {
818		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
819
820		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";
821
822		// Import first wallet
823		wallet_manager
824			.import_wallet("duplicate-import-wallet", test_mnemonic, None)
825			.await
826			.expect("Failed to import first wallet");
827
828		// Try to import with same name
829		let result = wallet_manager
830			.import_wallet("duplicate-import-wallet", test_mnemonic, None)
831			.await;
832
833		assert!(result.is_err());
834		match result.unwrap_err() {
835			crate::error::QuantusError::Wallet(WalletError::AlreadyExists) => {},
836			_ => panic!("Expected AlreadyExists error"),
837		}
838	}
839
840	#[tokio::test]
841	async fn test_list_wallets() {
842		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
843
844		// Initially should be empty
845		let wallets = wallet_manager.list_wallets().expect("Failed to list wallets");
846		assert_eq!(wallets.len(), 0);
847
848		// Create some wallets
849		wallet_manager
850			.create_wallet("wallet-1", Some("password1"))
851			.await
852			.expect("Failed to create wallet 1");
853
854		wallet_manager
855			.create_wallet("wallet-2", None)
856			.await
857			.expect("Failed to create wallet 2");
858
859		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";
860		wallet_manager
861			.import_wallet("imported-wallet", test_mnemonic, Some("password3"))
862			.await
863			.expect("Failed to import wallet");
864
865		// List wallets
866		let wallets = wallet_manager.list_wallets().expect("Failed to list wallets");
867
868		assert_eq!(wallets.len(), 3);
869
870		// Check that all wallet names are present
871		let wallet_names: Vec<&String> = wallets.iter().map(|w| &w.name).collect();
872		assert!(wallet_names.contains(&&"wallet-1".to_string()));
873		assert!(wallet_names.contains(&&"wallet-2".to_string()));
874		assert!(wallet_names.contains(&&"imported-wallet".to_string()));
875
876		// Check that addresses are real addresses (now stored unencrypted)
877		for wallet in &wallets {
878			assert!(wallet.address.starts_with("qz")); // Real SS58 addresses start with 5
879			assert_eq!(wallet.key_type, "Dilithium ML-DSA-87");
880		}
881
882		// Check sorting (newest first)
883		assert!(wallets[0].created_at >= wallets[1].created_at);
884		assert!(wallets[1].created_at >= wallets[2].created_at);
885	}
886
887	#[tokio::test]
888	async fn test_get_wallet() {
889		let (wallet_manager, _temp_dir) = create_test_wallet_manager().await;
890
891		// Create a wallet
892		let created_wallet = wallet_manager
893			.create_wallet("test-get-wallet", Some("test-password"))
894			.await
895			.expect("Failed to create wallet");
896
897		// Test getting wallet without password
898		let wallet_info = wallet_manager
899			.get_wallet("test-get-wallet", None)
900			.expect("Failed to get wallet")
901			.expect("Wallet should exist");
902
903		assert_eq!(wallet_info.name, "test-get-wallet");
904		assert_eq!(wallet_info.address, created_wallet.address); // Now returns real address
905
906		// Test getting wallet with wrong password
907		// Now with real quantum-safe encryption, wrong password should be detected
908		let wallet_info = wallet_manager
909			.get_wallet("test-get-wallet", Some("wrong-password"))
910			.expect("Failed to get wallet")
911			.expect("Wallet should exist");
912
913		assert_eq!(wallet_info.name, "test-get-wallet");
914		// With real encryption, wrong password returns placeholder text
915		assert_eq!(wallet_info.address, "[Wrong password]");
916
917		// Test getting wallet with correct password
918		let wallet_info = wallet_manager
919			.get_wallet("test-get-wallet", Some("test-password"))
920			.expect("Failed to get wallet")
921			.expect("Wallet should exist");
922
923		assert_eq!(wallet_info.name, "test-get-wallet");
924		assert_eq!(wallet_info.address, created_wallet.address);
925		assert!(wallet_info.address.starts_with("qz"));
926
927		// Test getting non-existent wallet
928		let result = wallet_manager
929			.get_wallet("non-existent-wallet", None)
930			.expect("Should not error on non-existent wallet");
931
932		assert!(result.is_none());
933	}
934}