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