1use crate::error::{Result, WalletError};
8use qp_rusty_crystals_dilithium::ml_dsa_87::{Keypair, PublicKey, SecretKey};
9use serde::{Deserialize, Serialize};
10use sp_core::{
11 crypto::{AccountId32, Ss58AddressFormat, Ss58Codec},
12 ByteArray,
13};
14use aes_gcm::{
16 aead::{Aead, AeadCore, KeyInit, OsRng as AesOsRng},
17 Aes256Gcm, Key, Nonce,
18};
19use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
20use rand::{rng, RngCore};
21
22use std::path::Path;
23
24use qp_dilithium_crypto::types::{DilithiumPair, DilithiumPublic};
25use sp_runtime::traits::IdentifyAccount;
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct QuantumKeyPair {
30 pub public_key: Vec<u8>,
31 pub private_key: Vec<u8>,
32}
33
34impl QuantumKeyPair {
35 pub fn from_dilithium_keypair(keypair: &Keypair) -> Self {
37 Self {
38 public_key: keypair.public.to_bytes().to_vec(),
39 private_key: keypair.secret.to_bytes().to_vec(),
40 }
41 }
42
43 #[allow(dead_code)]
45 pub fn to_dilithium_keypair(&self) -> Result<Keypair> {
46 Ok(Keypair {
50 public: PublicKey::from_bytes(&self.public_key).expect("Failed to parse public key"),
51 secret: SecretKey::from_bytes(&self.private_key).expect("Failed to parse private key"),
52 })
53 }
54
55 pub fn to_resonance_pair(&self) -> Result<DilithiumPair> {
57 Ok(DilithiumPair {
60 public: self.public_key.as_slice().try_into().unwrap(),
61 secret: self.private_key.as_slice().try_into().unwrap(),
62 })
63 }
64
65 #[allow(dead_code)]
66 pub fn from_resonance_pair(keypair: &DilithiumPair) -> Self {
67 Self {
68 public_key: keypair.public.as_ref().to_vec(),
69 private_key: keypair.secret.as_ref().to_vec(),
70 }
71 }
72
73 pub fn to_account_id_32(&self) -> AccountId32 {
74 let resonance_public =
76 DilithiumPublic::from_slice(&self.public_key).expect("Invalid public key");
77 resonance_public.into_account()
78 }
79
80 pub fn to_account_id_ss58check(&self) -> String {
81 let account = self.to_account_id_32();
82 account.to_ss58check_with_version(Ss58AddressFormat::custom(189))
83 }
84
85 pub fn to_subxt_signer(&self) -> Result<qp_dilithium_crypto::types::DilithiumPair> {
87 let resonance_pair = self.to_resonance_pair()?;
89
90 Ok(resonance_pair)
91 }
92
93 #[allow(dead_code)]
94 pub fn ss58_to_account_id(s: &str) -> Vec<u8> {
95 AsRef::<[u8]>::as_ref(&AccountId32::from_ss58check_with_version(s).unwrap().0).to_vec()
99 }
100}
101
102#[derive(Debug, Serialize, Deserialize)]
104pub struct EncryptedWallet {
105 pub name: String,
106 pub address: String, pub encrypted_data: Vec<u8>,
108 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>,
115}
116
117#[derive(Debug, Serialize, Deserialize)]
119pub struct WalletData {
120 pub name: String,
121 pub keypair: QuantumKeyPair,
122 pub mnemonic: Option<String>,
123 pub metadata: std::collections::HashMap<String, String>,
124}
125
126pub struct Keystore {
128 storage_path: std::path::PathBuf,
129}
130
131impl Keystore {
132 pub fn new<P: AsRef<Path>>(storage_path: P) -> Self {
134 Self { storage_path: storage_path.as_ref().to_path_buf() }
135 }
136
137 pub fn save_wallet(&self, wallet: &EncryptedWallet) -> Result<()> {
139 let wallet_file = self.storage_path.join(format!("{}.json", wallet.name));
140 let wallet_json = serde_json::to_string_pretty(wallet)?;
141 std::fs::write(wallet_file, wallet_json)?;
142 Ok(())
143 }
144
145 pub fn load_wallet(&self, name: &str) -> Result<Option<EncryptedWallet>> {
147 let wallet_file = self.storage_path.join(format!("{name}.json"));
148
149 if !wallet_file.exists() {
150 return Ok(None);
151 }
152
153 let wallet_json = std::fs::read_to_string(wallet_file)?;
154 let wallet: EncryptedWallet = serde_json::from_str(&wallet_json)?;
155 Ok(Some(wallet))
156 }
157
158 pub fn list_wallets(&self) -> Result<Vec<String>> {
160 let mut wallets = Vec::new();
161
162 if !self.storage_path.exists() {
163 return Ok(wallets);
164 }
165
166 for entry in std::fs::read_dir(&self.storage_path)? {
167 let entry = entry?;
168 let path = entry.path();
169
170 if path.extension().and_then(|s| s.to_str()) == Some("json") {
171 if let Some(name) = path.file_stem().and_then(|s| s.to_str()) {
172 wallets.push(name.to_string());
173 }
174 }
175 }
176
177 Ok(wallets)
178 }
179
180 pub fn delete_wallet(&self, name: &str) -> Result<bool> {
182 let wallet_file = self.storage_path.join(format!("{name}.json"));
183
184 if wallet_file.exists() {
185 std::fs::remove_file(wallet_file)?;
186 Ok(true)
187 } else {
188 Ok(false)
189 }
190 }
191
192 pub fn encrypt_wallet_data(
195 &self,
196 data: &WalletData,
197 password: &str,
198 ) -> Result<EncryptedWallet> {
199 let mut argon2_salt = [0u8; 16];
201 rng().fill_bytes(&mut argon2_salt);
202
203 let argon2 = Argon2::default();
205 let salt_string = argon2::password_hash::SaltString::encode_b64(&argon2_salt)
206 .map_err(|e| WalletError::Encryption(e.to_string()))?;
207 let password_hash = argon2
208 .hash_password(password.as_bytes(), &salt_string)
209 .map_err(|e| WalletError::Encryption(e.to_string()))?;
210
211 let hash_bytes = password_hash.hash.as_ref().unwrap().as_bytes();
213 let aes_key = Key::<Aes256Gcm>::from_slice(&hash_bytes[..32]);
214 let cipher = Aes256Gcm::new(aes_key);
215
216 let nonce = Aes256Gcm::generate_nonce(&mut AesOsRng);
218 let serialized_data = serde_json::to_vec(data)?;
219 let encrypted_data = cipher
220 .encrypt(&nonce, serialized_data.as_ref())
221 .map_err(|e| WalletError::Encryption(e.to_string()))?;
222
223 Ok(EncryptedWallet {
224 name: data.name.clone(),
225 address: data.keypair.to_account_id_ss58check(), encrypted_data,
227 kyber_ciphertext: vec![], kyber_public_key: vec![], argon2_salt: argon2_salt.to_vec(),
230 argon2_params: password_hash.to_string(),
231 aes_nonce: nonce.to_vec(),
232 encryption_version: 1, created_at: chrono::Utc::now(),
234 })
235 }
236
237 pub fn decrypt_wallet_data(
239 &self,
240 encrypted: &EncryptedWallet,
241 password: &str,
242 ) -> Result<WalletData> {
243 let argon2 = Argon2::default();
245 let password_hash = PasswordHash::new(&encrypted.argon2_params)
246 .map_err(|_| WalletError::InvalidPassword)?;
247
248 argon2
249 .verify_password(password.as_bytes(), &password_hash)
250 .map_err(|_| WalletError::InvalidPassword)?;
251
252 let hash_bytes = password_hash.hash.as_ref().unwrap().as_bytes();
254 let aes_key = Key::<Aes256Gcm>::from_slice(&hash_bytes[..32]);
255 let cipher = Aes256Gcm::new(aes_key);
256
257 let nonce = Nonce::from_slice(&encrypted.aes_nonce);
259 let decrypted_data = cipher
260 .decrypt(nonce, encrypted.encrypted_data.as_ref())
261 .map_err(|_| WalletError::Decryption)?;
262
263 let wallet_data: WalletData = serde_json::from_slice(&decrypted_data)?;
265
266 Ok(wallet_data)
267 }
268}
269
270#[cfg(test)]
271mod tests {
272 use super::*;
273 use qp_dilithium_crypto::{crystal_alice, crystal_charlie, dilithium_bob};
274 use qp_rusty_crystals_dilithium::ml_dsa_87::Keypair;
275 use tempfile::TempDir;
276
277 #[test]
278 fn test_quantum_keypair_from_dilithium_keypair() {
279 let entropy = [1u8; 32];
281 let dilithium_keypair = Keypair::generate(Some(&entropy));
282
283 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
285
286 assert_eq!(quantum_keypair.public_key, dilithium_keypair.public.to_bytes().to_vec());
288 assert_eq!(quantum_keypair.private_key, dilithium_keypair.secret.to_bytes().to_vec());
289 }
290
291 #[test]
292 fn test_quantum_keypair_to_dilithium_keypair_roundtrip() {
293 let entropy = [2u8; 32];
295 let original_keypair = Keypair::generate(Some(&entropy));
296
297 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&original_keypair);
299 let converted_keypair =
300 quantum_keypair.to_dilithium_keypair().expect("Conversion should succeed");
301
302 assert_eq!(original_keypair.public.to_bytes(), converted_keypair.public.to_bytes());
304 assert_eq!(original_keypair.secret.to_bytes(), converted_keypair.secret.to_bytes());
305 }
306
307 #[test]
308 fn test_quantum_keypair_from_resonance_pair() {
309 let resonance_pair = crystal_alice();
311 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
312
313 assert_eq!(quantum_keypair.public_key, resonance_pair.public.as_ref().to_vec());
315 assert_eq!(quantum_keypair.private_key, resonance_pair.secret.as_ref().to_vec());
316 }
317
318 #[test]
319 fn test_quantum_keypair_to_resonance_pair_roundtrip() {
320 let original_pair = dilithium_bob();
322 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&original_pair);
323 let converted_pair =
324 quantum_keypair.to_resonance_pair().expect("Conversion should succeed");
325
326 assert_eq!(original_pair.public.as_ref(), converted_pair.public.as_ref());
328 assert_eq!(original_pair.secret.as_ref(), converted_pair.secret.as_ref());
329 }
330
331 #[test]
332 fn test_quantum_keypair_address_generation() {
333 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
334 let test_pairs = vec![
336 ("crystal_alice", crystal_alice()),
337 ("crystal_bob", dilithium_bob()),
338 ("crystal_charlie", crystal_charlie()),
339 ];
340
341 for (name, resonance_pair) in test_pairs {
342 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
343
344 let account_id = quantum_keypair.to_account_id_32();
346 let ss58_address = quantum_keypair.to_account_id_ss58check();
347
348 assert!(ss58_address.starts_with("qz"), "SS58 address for {name} should start with 5");
350 assert!(
351 ss58_address.len() >= 47,
352 "SS58 address for {name} should be at least 47 characters"
353 );
354
355 assert_eq!(
357 account_id.to_ss58check(),
358 ss58_address,
359 "Address methods should be consistent for {name}"
360 );
361
362 let expected_address = resonance_pair.public().into_account().to_ss58check();
364 assert_eq!(
365 ss58_address, expected_address,
366 "Address should match DilithiumPair method for {name}"
367 );
368 }
369 }
370
371 #[test]
372 fn test_ss58_to_account_id_conversion() {
373 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
374 let test_cases = vec![
376 crystal_alice().public().into_account().to_ss58check(),
377 dilithium_bob().public().into_account().to_ss58check(),
378 crystal_charlie().public().into_account().to_ss58check(),
379 ];
380
381 for ss58_address in test_cases {
382 let account_bytes = QuantumKeyPair::ss58_to_account_id(&ss58_address);
384
385 assert_eq!(account_bytes.len(), 32, "Account ID should be 32 bytes");
387
388 let account_id =
390 AccountId32::from_slice(&account_bytes).expect("Should create valid AccountId32");
391 let round_trip_address = account_id.to_ss58check();
392 assert_eq!(
393 ss58_address, round_trip_address,
394 "Round-trip conversion should preserve address"
395 );
396 }
397 }
398
399 #[test]
400 fn test_address_consistency_across_conversions() {
401 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
403
404 let entropy = [3u8; 32];
405 let dilithium_keypair = Keypair::generate(Some(&entropy));
406
407 let quantum_from_dilithium = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
409 let resonance_from_quantum =
410 quantum_from_dilithium.to_resonance_pair().expect("Should convert");
411 let quantum_from_resonance = QuantumKeyPair::from_resonance_pair(&resonance_from_quantum);
412
413 let addr1 = quantum_from_dilithium.to_account_id_ss58check();
415 let addr2 = quantum_from_resonance.to_account_id_ss58check();
416 let addr3 = resonance_from_quantum.public().into_account().to_ss58check();
417
418 assert_eq!(addr1, addr2, "Addresses should be consistent across conversion paths");
419 assert_eq!(addr2, addr3, "Address should match direct DilithiumPair calculation");
420 }
421
422 #[test]
423 fn test_known_test_wallet_addresses() {
424 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
426 let alice_pair = crystal_alice();
427 let bob_pair = dilithium_bob();
428 let charlie_pair = crystal_charlie();
429
430 let alice_quantum = QuantumKeyPair::from_resonance_pair(&alice_pair);
431 let bob_quantum = QuantumKeyPair::from_resonance_pair(&bob_pair);
432 let charlie_quantum = QuantumKeyPair::from_resonance_pair(&charlie_pair);
433
434 let alice_addr = alice_quantum.to_account_id_ss58check();
435 let bob_addr = bob_quantum.to_account_id_ss58check();
436 let charlie_addr = charlie_quantum.to_account_id_ss58check();
437
438 assert_ne!(alice_addr, bob_addr, "Alice and Bob should have different addresses");
440 assert_ne!(bob_addr, charlie_addr, "Bob and Charlie should have different addresses");
441 assert_ne!(alice_addr, charlie_addr, "Alice and Charlie should have different addresses");
442
443 assert!(alice_addr.starts_with("qz"), "Alice address should be valid SS58");
445 assert!(bob_addr.starts_with("qz"), "Bob address should be valid SS58");
446 assert!(charlie_addr.starts_with("qz"), "Charlie address should be valid SS58");
447
448 println!("Test wallet addresses:");
449 println!(" Alice: {alice_addr}");
450 println!(" Bob: {bob_addr}");
451 println!(" Charlie: {charlie_addr}");
452 }
453
454 #[test]
455 fn test_invalid_ss58_address_handling() {
456 let invalid_addresses = vec![
458 "invalid",
459 "5", "1234567890", "", ];
463
464 for invalid_addr in invalid_addresses {
465 let result =
466 std::panic::catch_unwind(|| QuantumKeyPair::ss58_to_account_id(invalid_addr));
467 assert!(result.is_err(), "Should panic on invalid address: {invalid_addr}");
468 }
469 }
470
471 #[test]
472 fn test_stored_wallet_address_generation() {
473 sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
474
475 let alice_pair = crystal_alice();
480 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
481
482 let mut metadata = std::collections::HashMap::new();
484 metadata.insert("version".to_string(), "1.0.0".to_string());
485 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
486 metadata.insert("test_wallet".to_string(), "true".to_string());
487
488 let wallet_data = WalletData {
489 name: "test_crystal_alice".to_string(),
490 keypair: quantum_keypair.clone(),
491 mnemonic: None,
492 metadata,
493 };
494
495 let result = std::panic::catch_unwind(|| wallet_data.keypair.to_account_id_ss58check());
497
498 match result {
499 Ok(address) => {
500 println!("✅ Address generation successful: {address}");
501 let expected = alice_pair.public().into_account().to_ss58check();
503 assert_eq!(address, expected, "Stored wallet should generate correct address");
504 },
505 Err(_) => {
506 panic!("❌ Address generation failed - this is the bug we need to fix!");
507 },
508 }
509 }
510
511 #[test]
512 fn test_encrypted_wallet_address_generation() {
513 let temp_dir = tempfile::TempDir::new().expect("Failed to create temp directory");
517 let keystore = Keystore::new(temp_dir.path());
518
519 let alice_pair = crystal_alice();
521 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
522
523 let mut metadata = std::collections::HashMap::new();
524 metadata.insert("version".to_string(), "1.0.0".to_string());
525 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
526 metadata.insert("test_wallet".to_string(), "true".to_string());
527
528 let wallet_data = WalletData {
529 name: "test_crystal_alice".to_string(),
530 keypair: quantum_keypair,
531 mnemonic: None,
532 metadata,
533 };
534
535 let encrypted_wallet = keystore
537 .encrypt_wallet_data(&wallet_data, "")
538 .expect("Encryption should succeed");
539
540 keystore.save_wallet(&encrypted_wallet).expect("Save should succeed");
542 let loaded_wallet = keystore
543 .load_wallet("test_crystal_alice")
544 .expect("Load should succeed")
545 .expect("Wallet should exist");
546
547 let decrypted_data = keystore
549 .decrypt_wallet_data(&loaded_wallet, "")
550 .expect("Decryption should succeed");
551
552 let result = std::panic::catch_unwind(|| decrypted_data.keypair.to_account_id_ss58check());
554
555 match result {
556 Ok(address) => {
557 println!("✅ Encrypted wallet address generation successful: {address}");
558 let expected = alice_pair.public().into_account().to_ss58check();
560 assert_eq!(address, expected, "Decrypted wallet should generate correct address");
561 },
562 Err(_) => {
563 panic!("❌ Encrypted wallet address generation failed - this reproduces the send command bug!");
564 },
565 }
566 }
567
568 #[test]
569 fn test_send_command_wallet_loading_flow() {
570 let temp_dir = TempDir::new().expect("Failed to create temp directory");
575 let keystore = Keystore::new(temp_dir.path());
576
577 let alice_pair = crystal_alice();
579 let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
580
581 let mut metadata = std::collections::HashMap::new();
582 metadata.insert("version".to_string(), "1.0.0".to_string());
583 metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
584 metadata.insert("test_wallet".to_string(), "true".to_string());
585
586 let wallet_data = WalletData {
587 name: "crystal_alice".to_string(),
588 keypair: quantum_keypair,
589 mnemonic: None,
590 metadata,
591 };
592
593 let encrypted_wallet = keystore
595 .encrypt_wallet_data(&wallet_data, "")
596 .expect("Encryption should succeed");
597 keystore.save_wallet(&encrypted_wallet).expect("Save should succeed");
598
599 use crate::wallet::WalletManager;
602 let wallet_manager = WalletManager { wallets_dir: temp_dir.path().to_path_buf() };
603 let loaded_wallet_data =
604 wallet_manager.load_wallet("crystal_alice", "").expect("Should load wallet");
605
606 let result = std::panic::catch_unwind(|| {
608 loaded_wallet_data.keypair.to_account_id_ss58check()
610 });
611
612 match result {
613 Ok(address) => {
614 println!("✅ Send command flow works: {address}");
615 let expected = alice_pair.public().into_account().to_ss58check();
617 assert_eq!(address, expected, "Loaded wallet should generate correct address");
618 },
619 Err(_) => {
620 println!("❌ Send command flow failed - this reproduces the bug!");
621 panic!(
623 "This test reproduces the send command bug - load_wallet returns dummy data!"
624 );
625 },
626 }
627 }
628
629 #[test]
630 fn test_keypair_data_integrity() {
631 for i in 0..5 {
633 let entropy = [i as u8; 32];
634 let dilithium_keypair = Keypair::generate(Some(&entropy));
635 let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
636
637 if i == 0 {
639 println!("Actual public key size: {}", quantum_keypair.public_key.len());
640 println!("Actual private key size: {}", quantum_keypair.private_key.len());
641 }
642
643 assert!(
645 quantum_keypair.public_key.len() > 1000,
646 "Public key should be reasonably large (actual: {})",
647 quantum_keypair.public_key.len()
648 );
649 assert!(
650 quantum_keypair.private_key.len() > 2000,
651 "Private key should be reasonably large (actual: {})",
652 quantum_keypair.private_key.len()
653 );
654
655 assert!(
657 quantum_keypair.public_key.iter().any(|&b| b != 0),
658 "Public key should not be all zeros"
659 );
660 assert!(
661 quantum_keypair.private_key.iter().any(|&b| b != 0),
662 "Private key should not be all zeros"
663 );
664 }
665 }
666}