1use crate::error::{Result, WalletError};
8use qp_rusty_crystals_dilithium::ml_dsa_87::{Keypair, PublicKey, SecretKey};
9#[cfg(test)]
10use qp_rusty_crystals_hdwallet::SensitiveBytes32;
11use serde::{Deserialize, Serialize};
12#[cfg(test)]
13use sp_core::crypto::Ss58AddressFormat;
14use sp_core::{
15 crypto::{AccountId32, Ss58Codec},
16 ByteArray,
17};
18use aes_gcm::{
20 aead::{Aead, AeadCore, KeyInit, OsRng as AesOsRng},
21 Aes256Gcm, Key, Nonce,
22};
23use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
24use rand::{rng, RngCore};
25
26use std::path::Path;
27
28use qp_dilithium_crypto::types::{DilithiumPair, DilithiumPublic};
29use sp_runtime::traits::IdentifyAccount;
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct QuantumKeyPair {
34 pub public_key: Vec<u8>,
35 pub private_key: Vec<u8>,
36}
37
38impl QuantumKeyPair {
39 pub fn from_dilithium_keypair(keypair: &Keypair) -> Self {
41 Self {
42 public_key: keypair.public.to_bytes().to_vec(),
43 private_key: keypair.secret.to_bytes().to_vec(),
44 }
45 }
46
47 #[allow(dead_code)]
49 pub fn to_dilithium_keypair(&self) -> Result<Keypair> {
50 Ok(Keypair {
54 public: PublicKey::from_bytes(&self.public_key).expect("Failed to parse public key"),
55 secret: SecretKey::from_bytes(&self.private_key).expect("Failed to parse private key"),
56 })
57 }
58
59 pub fn to_resonance_pair(&self) -> Result<DilithiumPair> {
61 Ok(DilithiumPair {
64 public: self.public_key.as_slice().try_into().unwrap(),
65 secret: self.private_key.as_slice().try_into().unwrap(),
66 })
67 }
68
69 #[allow(dead_code)]
70 pub fn from_resonance_pair(keypair: &DilithiumPair) -> Self {
71 Self {
72 public_key: keypair.public.as_ref().to_vec(),
73 private_key: keypair.secret.as_ref().to_vec(),
74 }
75 }
76
77 pub fn to_account_id_32(&self) -> AccountId32 {
78 let resonance_public =
80 DilithiumPublic::from_slice(&self.public_key).expect("Invalid public key");
81 resonance_public.into_account()
82 }
83
84 pub fn to_account_id_ss58check(&self) -> String {
85 use crate::cli::address_format::quantus_ss58_format;
86 let account = self.to_account_id_32();
87 account.to_ss58check_with_version(quantus_ss58_format())
88 }
89
90 pub fn to_subxt_signer(&self) -> Result<qp_dilithium_crypto::types::DilithiumPair> {
92 let resonance_pair = self.to_resonance_pair()?;
94
95 Ok(resonance_pair)
96 }
97
98 #[allow(dead_code)]
99 pub fn ss58_to_account_id(s: &str) -> Vec<u8> {
100 AsRef::<[u8]>::as_ref(&AccountId32::from_ss58check_with_version(s).unwrap().0).to_vec()
104 }
105}
106
107#[derive(Debug, Serialize, Deserialize)]
109pub struct EncryptedWallet {
110 pub name: String,
111 pub address: String, pub encrypted_data: Vec<u8>,
113 pub kyber_ciphertext: Vec<u8>, pub kyber_public_key: Vec<u8>, pub argon2_salt: Vec<u8>, pub argon2_params: String, pub aes_nonce: Vec<u8>, pub encryption_version: u32, pub created_at: chrono::DateTime<chrono::Utc>,
120}
121
122#[derive(Debug, Serialize, Deserialize)]
124pub struct WalletData {
125 pub name: String,
126 pub keypair: QuantumKeyPair,
127 pub mnemonic: Option<String>,
128 pub derivation_path: String,
129 pub metadata: std::collections::HashMap<String, String>,
130}
131
132pub struct Keystore {
134 storage_path: std::path::PathBuf,
135}
136
137impl Keystore {
138 pub fn new<P: AsRef<Path>>(storage_path: P) -> Self {
140 Self { storage_path: storage_path.as_ref().to_path_buf() }
141 }
142
143 pub fn save_wallet(&self, wallet: &EncryptedWallet) -> Result<()> {
145 let wallet_file = self.storage_path.join(format!("{}.json", wallet.name));
146 let wallet_json = serde_json::to_string_pretty(wallet)?;
147 std::fs::write(wallet_file, wallet_json)?;
148 Ok(())
149 }
150
151 pub fn load_wallet(&self, name: &str) -> Result<Option<EncryptedWallet>> {
153 let wallet_file = self.storage_path.join(format!("{name}.json"));
154
155 if !wallet_file.exists() {
156 return Ok(None);
157 }
158
159 let wallet_json = std::fs::read_to_string(wallet_file)?;
160 let wallet: EncryptedWallet = serde_json::from_str(&wallet_json)?;
161 Ok(Some(wallet))
162 }
163
164 pub fn list_wallets(&self) -> Result<Vec<String>> {
166 let mut wallets = Vec::new();
167
168 if !self.storage_path.exists() {
169 return Ok(wallets);
170 }
171
172 for entry in std::fs::read_dir(&self.storage_path)? {
173 let entry = entry?;
174 let path = entry.path();
175
176 if path.extension().and_then(|s| s.to_str()) == Some("json") {
177 if let Some(name) = path.file_stem().and_then(|s| s.to_str()) {
178 wallets.push(name.to_string());
179 }
180 }
181 }
182
183 Ok(wallets)
184 }
185
186 pub fn delete_wallet(&self, name: &str) -> Result<bool> {
188 let wallet_file = self.storage_path.join(format!("{name}.json"));
189
190 if wallet_file.exists() {
191 std::fs::remove_file(wallet_file)?;
192 Ok(true)
193 } else {
194 Ok(false)
195 }
196 }
197
198 pub fn encrypt_wallet_data(
201 &self,
202 data: &WalletData,
203 password: &str,
204 ) -> Result<EncryptedWallet> {
205 let mut argon2_salt = [0u8; 16];
207 rng().fill_bytes(&mut argon2_salt);
208
209 let argon2 = Argon2::default();
211 let salt_string = argon2::password_hash::SaltString::encode_b64(&argon2_salt)
212 .map_err(|e| WalletError::Encryption(e.to_string()))?;
213 let password_hash = argon2
214 .hash_password(password.as_bytes(), &salt_string)
215 .map_err(|e| WalletError::Encryption(e.to_string()))?;
216
217 let hash_bytes = password_hash.hash.as_ref().unwrap().as_bytes();
219 let aes_key = Key::<Aes256Gcm>::from(<[u8; 32]>::try_from(&hash_bytes[..32]).unwrap());
220 let cipher = Aes256Gcm::new(&aes_key);
221
222 let nonce = Aes256Gcm::generate_nonce(&mut AesOsRng);
224 let serialized_data = serde_json::to_vec(data)?;
225 let encrypted_data = cipher
226 .encrypt(&nonce, serialized_data.as_ref())
227 .map_err(|e| WalletError::Encryption(e.to_string()))?;
228
229 Ok(EncryptedWallet {
230 name: data.name.clone(),
231 address: data.keypair.to_account_id_ss58check(), encrypted_data,
233 kyber_ciphertext: vec![], kyber_public_key: vec![], argon2_salt: argon2_salt.to_vec(),
236 argon2_params: password_hash.to_string(),
237 aes_nonce: nonce.to_vec(),
238 encryption_version: 1, created_at: chrono::Utc::now(),
240 })
241 }
242
243 pub fn decrypt_wallet_data(
245 &self,
246 encrypted: &EncryptedWallet,
247 password: &str,
248 ) -> Result<WalletData> {
249 let argon2 = Argon2::default();
251 let password_hash = PasswordHash::new(&encrypted.argon2_params)
252 .map_err(|_| WalletError::InvalidPassword)?;
253
254 argon2
255 .verify_password(password.as_bytes(), &password_hash)
256 .map_err(|_| WalletError::InvalidPassword)?;
257
258 let hash_bytes = password_hash.hash.as_ref().unwrap().as_bytes();
260 let aes_key = Key::<Aes256Gcm>::from(<[u8; 32]>::try_from(&hash_bytes[..32]).unwrap());
261 let cipher = Aes256Gcm::new(&aes_key);
262
263 let nonce = Nonce::from(<[u8; 12]>::try_from(&encrypted.aes_nonce[..]).unwrap());
265 let decrypted_data = cipher
266 .decrypt(&nonce, encrypted.encrypted_data.as_ref())
267 .map_err(|_| WalletError::Decryption)?;
268
269 let wallet_data: WalletData = serde_json::from_slice(&decrypted_data)?;
271
272 Ok(wallet_data)
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 use super::*;
279 use qp_dilithium_crypto::{crystal_alice, crystal_charlie, dilithium_bob};
280 use qp_rusty_crystals_dilithium::ml_dsa_87::Keypair;
281 use tempfile::TempDir;
282
283 #[test]
284 fn test_quantum_keypair_from_dilithium_keypair() {
285 let mut entropy = [1u8; 32];
287 let dilithium_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
288
289 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
291
292 assert_eq!(quantum_keypair.public_key, dilithium_keypair.public.to_bytes().to_vec());
294 assert_eq!(quantum_keypair.private_key, dilithium_keypair.secret.to_bytes().to_vec());
295 }
296
297 #[test]
298 fn test_quantum_keypair_to_dilithium_keypair_roundtrip() {
299 let mut entropy = [2u8; 32];
301 let original_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
302
303 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&original_keypair);
305 let converted_keypair =
306 quantum_keypair.to_dilithium_keypair().expect("Conversion should succeed");
307
308 assert_eq!(original_keypair.public.to_bytes(), converted_keypair.public.to_bytes());
310 assert_eq!(original_keypair.secret.to_bytes(), converted_keypair.secret.to_bytes());
311 }
312
313 #[test]
314 fn test_quantum_keypair_from_resonance_pair() {
315 let resonance_pair = crystal_alice();
317 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
318
319 assert_eq!(quantum_keypair.public_key, resonance_pair.public.as_ref().to_vec());
321 assert_eq!(quantum_keypair.private_key, resonance_pair.secret.as_ref().to_vec());
322 }
323
324 #[test]
325 fn test_quantum_keypair_to_resonance_pair_roundtrip() {
326 let original_pair = dilithium_bob();
328 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&original_pair);
329 let converted_pair =
330 quantum_keypair.to_resonance_pair().expect("Conversion should succeed");
331
332 assert_eq!(original_pair.public.as_ref(), converted_pair.public.as_ref());
334 assert_eq!(original_pair.secret.as_ref(), converted_pair.secret.as_ref());
335 }
336
337 #[test]
338 fn test_quantum_keypair_address_generation() {
339 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
340 let test_pairs = vec![
342 ("crystal_alice", crystal_alice()),
343 ("crystal_bob", dilithium_bob()),
344 ("crystal_charlie", crystal_charlie()),
345 ];
346
347 for (name, resonance_pair) in test_pairs {
348 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
349
350 let account_id = quantum_keypair.to_account_id_32();
352 let ss58_address = quantum_keypair.to_account_id_ss58check();
353
354 assert!(ss58_address.starts_with("qz"), "SS58 address for {name} should start with 5");
356 assert!(
357 ss58_address.len() >= 47,
358 "SS58 address for {name} should be at least 47 characters"
359 );
360
361 use crate::cli::address_format::quantus_ss58_format;
363 assert_eq!(
364 account_id.to_ss58check_with_version(quantus_ss58_format()),
365 ss58_address,
366 "Address methods should be consistent for {name}"
367 );
368
369 let expected_address = resonance_pair
371 .public()
372 .into_account()
373 .to_ss58check_with_version(quantus_ss58_format());
374 assert_eq!(
375 ss58_address, expected_address,
376 "Address should match DilithiumPair method for {name}"
377 );
378 }
379 }
380
381 #[test]
382 fn test_ss58_to_account_id_conversion() {
383 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
384 use crate::cli::address_format::quantus_ss58_format;
386 let test_cases = vec![
387 crystal_alice()
388 .public()
389 .into_account()
390 .to_ss58check_with_version(quantus_ss58_format()),
391 dilithium_bob()
392 .public()
393 .into_account()
394 .to_ss58check_with_version(quantus_ss58_format()),
395 crystal_charlie()
396 .public()
397 .into_account()
398 .to_ss58check_with_version(quantus_ss58_format()),
399 ];
400
401 for ss58_address in test_cases {
402 let account_bytes = QuantumKeyPair::ss58_to_account_id(&ss58_address);
404
405 assert_eq!(account_bytes.len(), 32, "Account ID should be 32 bytes");
407
408 let account_id =
410 AccountId32::from_slice(&account_bytes).expect("Should create valid AccountId32");
411 let round_trip_address =
412 account_id.to_ss58check_with_version(Ss58AddressFormat::custom(189));
413 assert_eq!(
414 ss58_address, round_trip_address,
415 "Round-trip conversion should preserve address"
416 );
417 }
418 }
419
420 #[test]
421 fn test_address_consistency_across_conversions() {
422 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
424
425 let mut entropy = [3u8; 32];
426 let dilithium_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
427
428 let quantum_from_dilithium = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
430 let resonance_from_quantum =
431 quantum_from_dilithium.to_resonance_pair().expect("Should convert");
432 let quantum_from_resonance = QuantumKeyPair::from_resonance_pair(&resonance_from_quantum);
433
434 let addr1 = quantum_from_dilithium.to_account_id_ss58check();
436 let addr2 = quantum_from_resonance.to_account_id_ss58check();
437 let addr3 = resonance_from_quantum
438 .public()
439 .into_account()
440 .to_ss58check_with_version(Ss58AddressFormat::custom(189));
441
442 assert_eq!(addr1, addr2, "Addresses should be consistent across conversion paths");
443 assert_eq!(addr2, addr3, "Address should match direct DilithiumPair calculation");
444 }
445
446 #[test]
447 fn test_known_test_wallet_addresses() {
448 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
450 let alice_pair = crystal_alice();
451 let bob_pair = dilithium_bob();
452 let charlie_pair = crystal_charlie();
453
454 let alice_quantum = QuantumKeyPair::from_resonance_pair(&alice_pair);
455 let bob_quantum = QuantumKeyPair::from_resonance_pair(&bob_pair);
456 let charlie_quantum = QuantumKeyPair::from_resonance_pair(&charlie_pair);
457
458 let alice_addr = alice_quantum.to_account_id_ss58check();
459 let bob_addr = bob_quantum.to_account_id_ss58check();
460 let charlie_addr = charlie_quantum.to_account_id_ss58check();
461
462 assert_ne!(alice_addr, bob_addr, "Alice and Bob should have different addresses");
464 assert_ne!(bob_addr, charlie_addr, "Bob and Charlie should have different addresses");
465 assert_ne!(alice_addr, charlie_addr, "Alice and Charlie should have different addresses");
466
467 assert!(alice_addr.starts_with("qz"), "Alice address should be valid SS58");
469 assert!(bob_addr.starts_with("qz"), "Bob address should be valid SS58");
470 assert!(charlie_addr.starts_with("qz"), "Charlie address should be valid SS58");
471
472 println!("Test wallet addresses:");
473 println!(" Alice: {alice_addr}");
474 println!(" Bob: {bob_addr}");
475 println!(" Charlie: {charlie_addr}");
476 }
477
478 #[test]
479 fn test_invalid_ss58_address_handling() {
480 let invalid_addresses = vec![
482 "invalid",
483 "5", "1234567890", "", ];
487
488 for invalid_addr in invalid_addresses {
489 let result =
490 std::panic::catch_unwind(|| QuantumKeyPair::ss58_to_account_id(invalid_addr));
491 assert!(result.is_err(), "Should panic on invalid address: {invalid_addr}");
492 }
493 }
494
495 #[test]
496 fn test_stored_wallet_address_generation() {
497 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
498
499 let alice_pair = crystal_alice();
504 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
505
506 let mut metadata = std::collections::HashMap::new();
508 metadata.insert("version".to_string(), "1.0.0".to_string());
509 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
510 metadata.insert("test_wallet".to_string(), "true".to_string());
511
512 let wallet_data = WalletData {
513 name: "test_crystal_alice".to_string(),
514 keypair: quantum_keypair.clone(),
515 mnemonic: None,
516 derivation_path: "m/".to_string(),
517 metadata,
518 };
519
520 let result = std::panic::catch_unwind(|| wallet_data.keypair.to_account_id_ss58check());
522
523 match result {
524 Ok(address) => {
525 println!("✅ Address generation successful: {address}");
526 let expected = alice_pair
528 .public()
529 .into_account()
530 .to_ss58check_with_version(Ss58AddressFormat::custom(189));
531 assert_eq!(address, expected, "Stored wallet should generate correct address");
532 },
533 Err(_) => {
534 panic!("❌ Address generation failed - this is the bug we need to fix!");
535 },
536 }
537 }
538
539 #[test]
540 fn test_encrypted_wallet_address_generation() {
541 let temp_dir = tempfile::TempDir::new().expect("Failed to create temp directory");
545 let keystore = Keystore::new(temp_dir.path());
546
547 let alice_pair = crystal_alice();
549 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
550
551 let mut metadata = std::collections::HashMap::new();
552 metadata.insert("version".to_string(), "1.0.0".to_string());
553 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
554 metadata.insert("test_wallet".to_string(), "true".to_string());
555
556 let wallet_data = WalletData {
557 name: "test_crystal_alice".to_string(),
558 keypair: quantum_keypair,
559 mnemonic: None,
560 derivation_path: "m/".to_string(),
561 metadata,
562 };
563
564 let encrypted_wallet = keystore
566 .encrypt_wallet_data(&wallet_data, "")
567 .expect("Encryption should succeed");
568
569 keystore.save_wallet(&encrypted_wallet).expect("Save should succeed");
571 let loaded_wallet = keystore
572 .load_wallet("test_crystal_alice")
573 .expect("Load should succeed")
574 .expect("Wallet should exist");
575
576 let decrypted_data = keystore
578 .decrypt_wallet_data(&loaded_wallet, "")
579 .expect("Decryption should succeed");
580
581 let result = std::panic::catch_unwind(|| decrypted_data.keypair.to_account_id_ss58check());
583
584 match result {
585 Ok(address) => {
586 println!("✅ Encrypted wallet address generation successful: {address}");
587 let expected = alice_pair
589 .public()
590 .into_account()
591 .to_ss58check_with_version(Ss58AddressFormat::custom(189));
592 assert_eq!(address, expected, "Decrypted wallet should generate correct address");
593 },
594 Err(_) => {
595 panic!("❌ Encrypted wallet address generation failed - this reproduces the send command bug!");
596 },
597 }
598 }
599
600 #[test]
601 fn test_send_command_wallet_loading_flow() {
602 let temp_dir = TempDir::new().expect("Failed to create temp directory");
607 let keystore = Keystore::new(temp_dir.path());
608
609 let alice_pair = crystal_alice();
611 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
612
613 let mut metadata = std::collections::HashMap::new();
614 metadata.insert("version".to_string(), "1.0.0".to_string());
615 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
616 metadata.insert("test_wallet".to_string(), "true".to_string());
617
618 let wallet_data = WalletData {
619 name: "crystal_alice".to_string(),
620 keypair: quantum_keypair,
621 mnemonic: None,
622 derivation_path: "m/".to_string(),
623 metadata,
624 };
625
626 let encrypted_wallet = keystore
628 .encrypt_wallet_data(&wallet_data, "")
629 .expect("Encryption should succeed");
630 keystore.save_wallet(&encrypted_wallet).expect("Save should succeed");
631
632 use crate::wallet::WalletManager;
635 let wallet_manager = WalletManager { wallets_dir: temp_dir.path().to_path_buf() };
636 let loaded_wallet_data =
637 wallet_manager.load_wallet("crystal_alice", "").expect("Should load wallet");
638
639 let result = std::panic::catch_unwind(|| {
641 loaded_wallet_data.keypair.to_account_id_ss58check()
643 });
644
645 match result {
646 Ok(address) => {
647 println!("✅ Send command flow works: {address}");
648 let expected = alice_pair
650 .public()
651 .into_account()
652 .to_ss58check_with_version(Ss58AddressFormat::custom(189));
653 assert_eq!(address, expected, "Loaded wallet should generate correct address");
654 },
655 Err(_) => {
656 println!("❌ Send command flow failed - this reproduces the bug!");
657 panic!(
659 "This test reproduces the send command bug - load_wallet returns dummy data!"
660 );
661 },
662 }
663 }
664
665 #[test]
666 fn test_keypair_data_integrity() {
667 for i in 0..5 {
669 let mut entropy = [i as u8; 32];
670 let dilithium_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
671 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
672
673 if i == 0 {
675 println!("Actual public key size: {}", quantum_keypair.public_key.len());
676 println!("Actual private key size: {}", quantum_keypair.private_key.len());
677 }
678
679 assert!(
681 quantum_keypair.public_key.len() > 1000,
682 "Public key should be reasonably large (actual: {})",
683 quantum_keypair.public_key.len()
684 );
685 assert!(
686 quantum_keypair.private_key.len() > 2000,
687 "Private key should be reasonably large (actual: {})",
688 quantum_keypair.private_key.len()
689 );
690
691 assert!(
693 quantum_keypair.public_key.iter().any(|&b| b != 0),
694 "Public key should not be all zeros"
695 );
696 assert!(
697 quantum_keypair.private_key.iter().any(|&b| b != 0),
698 "Private key should not be all zeros"
699 );
700 }
701 }
702}