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