use crate::error::{Result, WalletError};
use qp_rusty_crystals_dilithium::ml_dsa_87::{Keypair, PublicKey, SecretKey};
#[cfg(test)]
use qp_rusty_crystals_hdwallet::SensitiveBytes32;
use serde::{Deserialize, Serialize};
#[cfg(test)]
use sp_core::crypto::Ss58AddressFormat;
use sp_core::{
crypto::{AccountId32, Ss58Codec},
ByteArray,
};
use aes_gcm::{
aead::{Aead, AeadCore, KeyInit, OsRng as AesOsRng},
Aes256Gcm, Key, Nonce,
};
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
use rand::{rng, RngCore};
use std::path::Path;
use qp_dilithium_crypto::types::{DilithiumPair, DilithiumPublic};
use sp_runtime::traits::IdentifyAccount;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantumKeyPair {
pub public_key: Vec<u8>,
pub private_key: Vec<u8>,
}
impl QuantumKeyPair {
pub fn from_dilithium_keypair(keypair: &Keypair) -> Self {
Self {
public_key: keypair.public.to_bytes().to_vec(),
private_key: keypair.secret.to_bytes().to_vec(),
}
}
#[allow(dead_code)]
pub fn to_dilithium_keypair(&self) -> Result<Keypair> {
Ok(Keypair {
public: PublicKey::from_bytes(&self.public_key).expect("Failed to parse public key"),
secret: SecretKey::from_bytes(&self.private_key).expect("Failed to parse private key"),
})
}
pub fn to_resonance_pair(&self) -> Result<DilithiumPair> {
Ok(DilithiumPair::from_raw(&self.public_key, &self.private_key)
.map_err(|_| crate::error::WalletError::KeyGeneration)?)
}
pub fn from_resonance_pair(keypair: &DilithiumPair) -> Self {
use sp_core::Pair;
Self {
public_key: keypair.public().as_ref().to_vec(),
private_key: keypair.secret_bytes().to_vec(),
}
}
pub fn to_account_id_32(&self) -> AccountId32 {
let resonance_public =
DilithiumPublic::from_slice(&self.public_key).expect("Invalid public key");
resonance_public.into_account()
}
pub fn to_account_id_ss58check(&self) -> String {
use crate::cli::address_format::quantus_ss58_format;
let account = self.to_account_id_32();
account.to_ss58check_with_version(quantus_ss58_format())
}
pub fn to_subxt_signer(&self) -> Result<qp_dilithium_crypto::types::DilithiumPair> {
let resonance_pair = self.to_resonance_pair()?;
Ok(resonance_pair)
}
#[allow(dead_code)]
pub fn ss58_to_account_id(s: &str) -> Vec<u8> {
AsRef::<[u8]>::as_ref(&AccountId32::from_ss58check_with_version(s).unwrap().0).to_vec()
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct EncryptedWallet {
pub name: String,
pub address: String, pub encrypted_data: Vec<u8>,
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>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct WalletData {
pub name: String,
pub keypair: QuantumKeyPair,
pub mnemonic: Option<String>,
pub derivation_path: String,
pub metadata: std::collections::HashMap<String, String>,
}
pub struct Keystore {
storage_path: std::path::PathBuf,
}
impl Keystore {
pub fn new<P: AsRef<Path>>(storage_path: P) -> Self {
Self { storage_path: storage_path.as_ref().to_path_buf() }
}
pub fn save_wallet(&self, wallet: &EncryptedWallet) -> Result<()> {
let wallet_file = self.storage_path.join(format!("{}.json", wallet.name));
let wallet_json = serde_json::to_string_pretty(wallet)?;
std::fs::write(wallet_file, wallet_json)?;
Ok(())
}
pub fn load_wallet(&self, name: &str) -> Result<Option<EncryptedWallet>> {
let wallet_file = self.storage_path.join(format!("{name}.json"));
if !wallet_file.exists() {
return Ok(None);
}
let wallet_json = std::fs::read_to_string(wallet_file)?;
let wallet: EncryptedWallet = serde_json::from_str(&wallet_json)?;
Ok(Some(wallet))
}
pub fn list_wallets(&self) -> Result<Vec<String>> {
let mut wallets = Vec::new();
if !self.storage_path.exists() {
return Ok(wallets);
}
for entry in std::fs::read_dir(&self.storage_path)? {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|s| s.to_str()) == Some("json") {
if let Some(name) = path.file_stem().and_then(|s| s.to_str()) {
wallets.push(name.to_string());
}
}
}
Ok(wallets)
}
pub fn delete_wallet(&self, name: &str) -> Result<bool> {
let wallet_file = self.storage_path.join(format!("{name}.json"));
if wallet_file.exists() {
std::fs::remove_file(wallet_file)?;
Ok(true)
} else {
Ok(false)
}
}
pub fn encrypt_wallet_data(
&self,
data: &WalletData,
password: &str,
) -> Result<EncryptedWallet> {
let mut argon2_salt = [0u8; 16];
rng().fill_bytes(&mut argon2_salt);
let argon2 = Argon2::default();
let salt_string = argon2::password_hash::SaltString::encode_b64(&argon2_salt)
.map_err(|e| WalletError::Encryption(e.to_string()))?;
let password_hash = argon2
.hash_password(password.as_bytes(), &salt_string)
.map_err(|e| WalletError::Encryption(e.to_string()))?;
let hash_bytes = password_hash.hash.as_ref().unwrap().as_bytes();
let aes_key = Key::<Aes256Gcm>::from(<[u8; 32]>::try_from(&hash_bytes[..32]).unwrap());
let cipher = Aes256Gcm::new(&aes_key);
let nonce = Aes256Gcm::generate_nonce(&mut AesOsRng);
let serialized_data = serde_json::to_vec(data)?;
let encrypted_data = cipher
.encrypt(&nonce, serialized_data.as_ref())
.map_err(|e| WalletError::Encryption(e.to_string()))?;
Ok(EncryptedWallet {
name: data.name.clone(),
address: data.keypair.to_account_id_ss58check(), encrypted_data,
kyber_ciphertext: vec![], kyber_public_key: vec![], argon2_salt: argon2_salt.to_vec(),
argon2_params: password_hash.to_string(),
aes_nonce: nonce.to_vec(),
encryption_version: 1, created_at: chrono::Utc::now(),
})
}
pub fn decrypt_wallet_data(
&self,
encrypted: &EncryptedWallet,
password: &str,
) -> Result<WalletData> {
let argon2 = Argon2::default();
let password_hash = PasswordHash::new(&encrypted.argon2_params)
.map_err(|_| WalletError::InvalidPassword)?;
argon2
.verify_password(password.as_bytes(), &password_hash)
.map_err(|_| WalletError::InvalidPassword)?;
let hash_bytes = password_hash.hash.as_ref().unwrap().as_bytes();
let aes_key = Key::<Aes256Gcm>::from(<[u8; 32]>::try_from(&hash_bytes[..32]).unwrap());
let cipher = Aes256Gcm::new(&aes_key);
let nonce = Nonce::from(<[u8; 12]>::try_from(&encrypted.aes_nonce[..]).unwrap());
let decrypted_data = cipher
.decrypt(&nonce, encrypted.encrypted_data.as_ref())
.map_err(|_| WalletError::Decryption)?;
let wallet_data: WalletData = serde_json::from_slice(&decrypted_data)?;
Ok(wallet_data)
}
}
#[cfg(test)]
mod tests {
use super::*;
use qp_dilithium_crypto::{crystal_alice, crystal_charlie, dilithium_bob};
use qp_rusty_crystals_dilithium::ml_dsa_87::Keypair;
use sp_core::Pair;
use tempfile::TempDir;
#[test]
fn test_quantum_keypair_from_dilithium_keypair() {
let mut entropy = [1u8; 32];
let dilithium_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
assert_eq!(quantum_keypair.public_key, dilithium_keypair.public.to_bytes().to_vec());
assert_eq!(quantum_keypair.private_key, dilithium_keypair.secret.to_bytes().to_vec());
}
#[test]
fn test_quantum_keypair_to_dilithium_keypair_roundtrip() {
let mut entropy = [2u8; 32];
let original_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&original_keypair);
let converted_keypair =
quantum_keypair.to_dilithium_keypair().expect("Conversion should succeed");
assert_eq!(original_keypair.public.to_bytes(), converted_keypair.public.to_bytes());
assert_eq!(original_keypair.secret.to_bytes(), converted_keypair.secret.to_bytes());
}
#[test]
fn test_quantum_keypair_from_resonance_pair() {
let resonance_pair = crystal_alice();
let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
assert_eq!(quantum_keypair.public_key, resonance_pair.public().as_ref().to_vec());
assert_eq!(quantum_keypair.private_key.as_slice(), resonance_pair.secret_bytes());
}
#[test]
fn test_quantum_keypair_to_resonance_pair_roundtrip() {
let original_pair = dilithium_bob();
let quantum_keypair = QuantumKeyPair::from_resonance_pair(&original_pair);
let converted_pair =
quantum_keypair.to_resonance_pair().expect("Conversion should succeed");
assert_eq!(original_pair.public().as_ref(), converted_pair.public().as_ref());
assert_eq!(original_pair.secret_bytes(), converted_pair.secret_bytes());
}
#[test]
fn test_quantum_keypair_address_generation() {
sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
let test_pairs = vec![
("crystal_alice", crystal_alice()),
("crystal_bob", dilithium_bob()),
("crystal_charlie", crystal_charlie()),
];
for (name, resonance_pair) in test_pairs {
let quantum_keypair = QuantumKeyPair::from_resonance_pair(&resonance_pair);
let account_id = quantum_keypair.to_account_id_32();
let ss58_address = quantum_keypair.to_account_id_ss58check();
assert!(
ss58_address.starts_with("qz"),
"SS58 address for {name} should start with qz (Quantus prefix 189)"
);
assert!(
ss58_address.len() >= 47,
"SS58 address for {name} should be at least 47 characters"
);
use crate::cli::address_format::quantus_ss58_format;
assert_eq!(
account_id.to_ss58check_with_version(quantus_ss58_format()),
ss58_address,
"Address methods should be consistent for {name}"
);
let chain_expected_address = resonance_pair
.public()
.into_account()
.to_ss58check_with_version(quantus_ss58_format());
assert_eq!(
ss58_address, chain_expected_address,
"Wallet address for {name} must match chain dev account (same derivation and SS58 189)"
);
}
}
#[test]
fn test_ss58_to_account_id_conversion() {
sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
use crate::cli::address_format::quantus_ss58_format;
let test_cases = vec![
crystal_alice()
.public()
.into_account()
.to_ss58check_with_version(quantus_ss58_format()),
dilithium_bob()
.public()
.into_account()
.to_ss58check_with_version(quantus_ss58_format()),
crystal_charlie()
.public()
.into_account()
.to_ss58check_with_version(quantus_ss58_format()),
];
for ss58_address in test_cases {
let account_bytes = QuantumKeyPair::ss58_to_account_id(&ss58_address);
assert_eq!(account_bytes.len(), 32, "Account ID should be 32 bytes");
let account_id =
AccountId32::from_slice(&account_bytes).expect("Should create valid AccountId32");
let round_trip_address =
account_id.to_ss58check_with_version(Ss58AddressFormat::custom(189));
assert_eq!(
ss58_address, round_trip_address,
"Round-trip conversion should preserve address"
);
}
}
#[test]
fn test_address_consistency_across_conversions() {
sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
let mut entropy = [3u8; 32];
let dilithium_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
let quantum_from_dilithium = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
let resonance_from_quantum =
quantum_from_dilithium.to_resonance_pair().expect("Should convert");
let quantum_from_resonance = QuantumKeyPair::from_resonance_pair(&resonance_from_quantum);
let addr1 = quantum_from_dilithium.to_account_id_ss58check();
let addr2 = quantum_from_resonance.to_account_id_ss58check();
let addr3 = resonance_from_quantum
.public()
.into_account()
.to_ss58check_with_version(Ss58AddressFormat::custom(189));
assert_eq!(addr1, addr2, "Addresses should be consistent across conversion paths");
assert_eq!(addr2, addr3, "Address should match direct DilithiumPair calculation");
}
#[test]
fn test_known_test_wallet_addresses() {
sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
let alice_pair = crystal_alice();
let bob_pair = dilithium_bob();
let charlie_pair = crystal_charlie();
let alice_quantum = QuantumKeyPair::from_resonance_pair(&alice_pair);
let bob_quantum = QuantumKeyPair::from_resonance_pair(&bob_pair);
let charlie_quantum = QuantumKeyPair::from_resonance_pair(&charlie_pair);
let alice_addr = alice_quantum.to_account_id_ss58check();
let bob_addr = bob_quantum.to_account_id_ss58check();
let charlie_addr = charlie_quantum.to_account_id_ss58check();
assert_ne!(alice_addr, bob_addr, "Alice and Bob should have different addresses");
assert_ne!(bob_addr, charlie_addr, "Bob and Charlie should have different addresses");
assert_ne!(alice_addr, charlie_addr, "Alice and Charlie should have different addresses");
assert!(alice_addr.starts_with("qz"), "Alice address should be valid SS58");
assert!(bob_addr.starts_with("qz"), "Bob address should be valid SS58");
assert!(charlie_addr.starts_with("qz"), "Charlie address should be valid SS58");
println!("Test wallet addresses:");
println!(" Alice: {alice_addr}");
println!(" Bob: {bob_addr}");
println!(" Charlie: {charlie_addr}");
}
#[test]
fn test_invalid_ss58_address_handling() {
let invalid_addresses = vec![
"invalid",
"5", "1234567890", "", ];
for invalid_addr in invalid_addresses {
let result =
std::panic::catch_unwind(|| QuantumKeyPair::ss58_to_account_id(invalid_addr));
assert!(result.is_err(), "Should panic on invalid address: {invalid_addr}");
}
}
#[test]
fn test_stored_wallet_address_generation() {
sp_core::crypto::set_default_ss58_version(sp_core::crypto::Ss58AddressFormat::custom(189));
let alice_pair = crystal_alice();
let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
let mut metadata = std::collections::HashMap::new();
metadata.insert("version".to_string(), "1.0.0".to_string());
metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
metadata.insert("test_wallet".to_string(), "true".to_string());
let wallet_data = WalletData {
name: "test_crystal_alice".to_string(),
keypair: quantum_keypair.clone(),
mnemonic: None,
derivation_path: "m/".to_string(),
metadata,
};
let result = std::panic::catch_unwind(|| wallet_data.keypair.to_account_id_ss58check());
match result {
Ok(address) => {
println!("✅ Address generation successful: {address}");
let expected = alice_pair
.public()
.into_account()
.to_ss58check_with_version(Ss58AddressFormat::custom(189));
assert_eq!(address, expected, "Stored wallet should generate correct address");
},
Err(_) => {
panic!("❌ Address generation failed - this is the bug we need to fix!");
},
}
}
#[test]
fn test_encrypted_wallet_address_generation() {
let temp_dir = tempfile::TempDir::new().expect("Failed to create temp directory");
let keystore = Keystore::new(temp_dir.path());
let alice_pair = crystal_alice();
let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
let mut metadata = std::collections::HashMap::new();
metadata.insert("version".to_string(), "1.0.0".to_string());
metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
metadata.insert("test_wallet".to_string(), "true".to_string());
let wallet_data = WalletData {
name: "test_crystal_alice".to_string(),
keypair: quantum_keypair,
mnemonic: None,
derivation_path: "m/".to_string(),
metadata,
};
let encrypted_wallet = keystore
.encrypt_wallet_data(&wallet_data, "")
.expect("Encryption should succeed");
keystore.save_wallet(&encrypted_wallet).expect("Save should succeed");
let loaded_wallet = keystore
.load_wallet("test_crystal_alice")
.expect("Load should succeed")
.expect("Wallet should exist");
let decrypted_data = keystore
.decrypt_wallet_data(&loaded_wallet, "")
.expect("Decryption should succeed");
let result = std::panic::catch_unwind(|| decrypted_data.keypair.to_account_id_ss58check());
match result {
Ok(address) => {
println!("✅ Encrypted wallet address generation successful: {address}");
let expected = alice_pair
.public()
.into_account()
.to_ss58check_with_version(Ss58AddressFormat::custom(189));
assert_eq!(address, expected, "Decrypted wallet should generate correct address");
},
Err(_) => {
panic!("❌ Encrypted wallet address generation failed - this reproduces the send command bug!");
},
}
}
#[test]
fn test_send_command_wallet_loading_flow() {
let temp_dir = TempDir::new().expect("Failed to create temp directory");
let keystore = Keystore::new(temp_dir.path());
let alice_pair = crystal_alice();
let quantum_keypair = QuantumKeyPair::from_resonance_pair(&alice_pair);
let mut metadata = std::collections::HashMap::new();
metadata.insert("version".to_string(), "1.0.0".to_string());
metadata.insert("algorithm".to_string(), "ML-DSA-87".to_string());
metadata.insert("test_wallet".to_string(), "true".to_string());
let wallet_data = WalletData {
name: "crystal_alice".to_string(),
keypair: quantum_keypair,
mnemonic: None,
derivation_path: "m/".to_string(),
metadata,
};
let encrypted_wallet = keystore
.encrypt_wallet_data(&wallet_data, "")
.expect("Encryption should succeed");
keystore.save_wallet(&encrypted_wallet).expect("Save should succeed");
use crate::wallet::WalletManager;
let wallet_manager = WalletManager { wallets_dir: temp_dir.path().to_path_buf() };
let loaded_wallet_data =
wallet_manager.load_wallet("crystal_alice", "").expect("Should load wallet");
let result = std::panic::catch_unwind(|| {
loaded_wallet_data.keypair.to_account_id_ss58check()
});
match result {
Ok(address) => {
println!("✅ Send command flow works: {address}");
let expected = alice_pair
.public()
.into_account()
.to_ss58check_with_version(Ss58AddressFormat::custom(189));
assert_eq!(address, expected, "Loaded wallet should generate correct address");
},
Err(_) => {
println!("❌ Send command flow failed - this reproduces the bug!");
panic!(
"This test reproduces the send command bug - load_wallet returns dummy data!"
);
},
}
}
#[test]
fn test_keypair_data_integrity() {
for i in 0..5 {
let mut entropy = [i as u8; 32];
let dilithium_keypair = Keypair::generate(SensitiveBytes32::from(&mut entropy));
let quantum_keypair = QuantumKeyPair::from_dilithium_keypair(&dilithium_keypair);
if i == 0 {
println!("Actual public key size: {}", quantum_keypair.public_key.len());
println!("Actual private key size: {}", quantum_keypair.private_key.len());
}
assert!(
quantum_keypair.public_key.len() > 1000,
"Public key should be reasonably large (actual: {})",
quantum_keypair.public_key.len()
);
assert!(
quantum_keypair.private_key.len() > 2000,
"Private key should be reasonably large (actual: {})",
quantum_keypair.private_key.len()
);
assert!(
quantum_keypair.public_key.iter().any(|&b| b != 0),
"Public key should not be all zeros"
);
assert!(
quantum_keypair.private_key.iter().any(|&b| b != 0),
"Private key should not be all zeros"
);
}
}
}