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