apex_sdk_substrate/
wallet.rs

1//! Substrate wallet and account management
2//!
3//! This module provides comprehensive wallet functionality including:
4//! - Key pair generation (SR25519, ED25519)
5//! - Mnemonic phrase support (BIP-39)
6//! - SS58 address encoding
7//! - Message and transaction signing
8//! - Multi-wallet management
9
10use crate::{Error, Result};
11use parking_lot::RwLock;
12use sp_core::crypto::{Ss58AddressFormat, Ss58Codec};
13use sp_core::{ed25519, sr25519, Pair as PairTrait};
14use std::collections::HashMap;
15use std::sync::Arc;
16use tracing::{debug, info};
17
18/// Supported key pair types
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum KeyPairType {
21    /// SR25519 (Schnorrkel) - Default for Substrate
22    Sr25519,
23    /// ED25519 - Alternative signing algorithm
24    Ed25519,
25}
26
27impl Default for KeyPairType {
28    fn default() -> Self {
29        Self::Sr25519
30    }
31}
32
33/// A unified wallet that can hold either SR25519 or ED25519 keys
34///
35/// # Security
36/// This struct implements `Clone` to support wallet management operations.
37/// Be aware that cloning duplicates private key material in memory.
38/// For shared access without duplication, consider wrapping in `Arc<Wallet>`.
39#[derive(Clone)]
40pub struct Wallet {
41    /// The key pair type
42    key_type: KeyPairType,
43    /// SR25519 pair (if applicable)
44    sr25519_pair: Option<sr25519::Pair>,
45    /// ED25519 pair (if applicable)
46    ed25519_pair: Option<ed25519::Pair>,
47    /// SS58 address format (network prefix)
48    ss58_format: Ss58AddressFormat,
49}
50
51impl Wallet {
52    /// Create a new random wallet with SR25519 keys
53    pub fn new_random() -> Self {
54        Self::new_random_with_type(KeyPairType::Sr25519)
55    }
56
57    /// Create a new random wallet with specified key type
58    pub fn new_random_with_type(key_type: KeyPairType) -> Self {
59        info!("Creating new random {:?} wallet", key_type);
60
61        match key_type {
62            KeyPairType::Sr25519 => {
63                let (pair, _seed) = sr25519::Pair::generate();
64                Self {
65                    key_type,
66                    sr25519_pair: Some(pair),
67                    ed25519_pair: None,
68                    ss58_format: Ss58AddressFormat::custom(42), // Default to generic
69                }
70            }
71            KeyPairType::Ed25519 => {
72                let (pair, _seed) = ed25519::Pair::generate();
73                Self {
74                    key_type,
75                    sr25519_pair: None,
76                    ed25519_pair: Some(pair),
77                    ss58_format: Ss58AddressFormat::custom(42),
78                }
79            }
80        }
81    }
82
83    /// Create wallet from mnemonic phrase
84    #[allow(clippy::result_large_err)]
85    pub fn from_mnemonic(mnemonic: &str, key_type: KeyPairType) -> Result<Self> {
86        Self::from_mnemonic_with_path(mnemonic, None, key_type)
87    }
88
89    /// Create wallet from mnemonic phrase with derivation path
90    #[allow(clippy::result_large_err)]
91    pub fn from_mnemonic_with_path(
92        mnemonic: &str,
93        path: Option<&str>,
94        key_type: KeyPairType,
95    ) -> Result<Self> {
96        info!("Creating wallet from mnemonic with {:?} keys", key_type);
97
98        // Validate mnemonic
99        let _ = bip39::Mnemonic::parse(mnemonic)
100            .map_err(|e| Error::Wallet(format!("Invalid mnemonic: {}", e)))?;
101
102        // Create derivation path string
103        let full_path = if let Some(p) = path {
104            format!("{}//{}", mnemonic, p)
105        } else {
106            mnemonic.to_string()
107        };
108
109        match key_type {
110            KeyPairType::Sr25519 => {
111                let pair = sr25519::Pair::from_string(&full_path, None)
112                    .map_err(|e| Error::Wallet(format!("Failed to derive key: {:?}", e)))?;
113
114                Ok(Self {
115                    key_type,
116                    sr25519_pair: Some(pair),
117                    ed25519_pair: None,
118                    ss58_format: Ss58AddressFormat::custom(42),
119                })
120            }
121            KeyPairType::Ed25519 => {
122                let pair = ed25519::Pair::from_string(&full_path, None)
123                    .map_err(|e| Error::Wallet(format!("Failed to derive key: {:?}", e)))?;
124
125                Ok(Self {
126                    key_type,
127                    sr25519_pair: None,
128                    ed25519_pair: Some(pair),
129                    ss58_format: Ss58AddressFormat::custom(42),
130                })
131            }
132        }
133    }
134
135    /// Create wallet from private key (seed)
136    #[allow(clippy::result_large_err)]
137    pub fn from_seed(seed: &[u8], key_type: KeyPairType) -> Result<Self> {
138        info!("Creating wallet from seed with {:?} keys", key_type);
139
140        if seed.len() != 32 {
141            return Err(Error::Wallet("Seed must be 32 bytes".to_string()));
142        }
143
144        let mut seed_array = [0u8; 32];
145        seed_array.copy_from_slice(seed);
146
147        match key_type {
148            KeyPairType::Sr25519 => {
149                let pair = sr25519::Pair::from_seed(&seed_array);
150                Ok(Self {
151                    key_type,
152                    sr25519_pair: Some(pair),
153                    ed25519_pair: None,
154                    ss58_format: Ss58AddressFormat::custom(42),
155                })
156            }
157            KeyPairType::Ed25519 => {
158                let pair = ed25519::Pair::from_seed(&seed_array);
159                Ok(Self {
160                    key_type,
161                    sr25519_pair: None,
162                    ed25519_pair: Some(pair),
163                    ss58_format: Ss58AddressFormat::custom(42),
164                })
165            }
166        }
167    }
168
169    /// Generate a new mnemonic phrase
170    pub fn generate_mnemonic() -> String {
171        use bip39::{Language, Mnemonic};
172        use rand::RngCore;
173
174        let mut entropy = [0u8; 32];
175        rand::rng().fill_bytes(&mut entropy);
176
177        Mnemonic::from_entropy_in(Language::English, &entropy)
178            .expect("Failed to generate mnemonic")
179            .to_string()
180    }
181
182    /// Set the SS58 address format (network prefix)
183    pub fn with_ss58_format(mut self, format: u16) -> Self {
184        self.ss58_format = Ss58AddressFormat::custom(format);
185        self
186    }
187
188    /// Get the public key as bytes
189    pub fn public_key(&self) -> Vec<u8> {
190        match self.key_type {
191            KeyPairType::Sr25519 => self.sr25519_pair.as_ref().unwrap().public().0.to_vec(),
192            KeyPairType::Ed25519 => self.ed25519_pair.as_ref().unwrap().public().0.to_vec(),
193        }
194    }
195
196    /// Get the SS58-encoded address
197    pub fn address(&self) -> String {
198        match self.key_type {
199            KeyPairType::Sr25519 => {
200                let public = self.sr25519_pair.as_ref().unwrap().public();
201                public.to_ss58check_with_version(self.ss58_format)
202            }
203            KeyPairType::Ed25519 => {
204                let public = self.ed25519_pair.as_ref().unwrap().public();
205                public.to_ss58check_with_version(self.ss58_format)
206            }
207        }
208    }
209
210    /// Get the key pair type
211    pub fn key_type(&self) -> KeyPairType {
212        self.key_type
213    }
214
215    /// Sign a message
216    pub fn sign(&self, message: &[u8]) -> Vec<u8> {
217        match self.key_type {
218            KeyPairType::Sr25519 => {
219                let pair = self.sr25519_pair.as_ref().unwrap();
220                pair.sign(message).0.to_vec()
221            }
222            KeyPairType::Ed25519 => {
223                let pair = self.ed25519_pair.as_ref().unwrap();
224                pair.sign(message).0.to_vec()
225            }
226        }
227    }
228
229    /// Verify a signature
230    pub fn verify(&self, message: &[u8], signature: &[u8]) -> bool {
231        match self.key_type {
232            KeyPairType::Sr25519 => {
233                if signature.len() != 64 {
234                    return false;
235                }
236                let mut sig_array = [0u8; 64];
237                sig_array.copy_from_slice(signature);
238                let sig = sr25519::Signature::from_raw(sig_array);
239                let public = self.sr25519_pair.as_ref().unwrap().public();
240                sr25519::Pair::verify(&sig, message, &public)
241            }
242            KeyPairType::Ed25519 => {
243                if signature.len() != 64 {
244                    return false;
245                }
246                let mut sig_array = [0u8; 64];
247                sig_array.copy_from_slice(signature);
248                let sig = ed25519::Signature::from_raw(sig_array);
249                let public = self.ed25519_pair.as_ref().unwrap().public();
250                ed25519::Pair::verify(&sig, message, &public)
251            }
252        }
253    }
254
255    /// Get the seed/private key (if available)
256    /// Note: This should be kept secure and not exposed in production
257    pub fn seed(&self) -> Option<[u8; 32]> {
258        match self.key_type {
259            KeyPairType::Sr25519 => {
260                // SR25519 doesn't expose seed directly in a simple way
261                None
262            }
263            KeyPairType::Ed25519 => {
264                // ED25519 also doesn't expose seed directly
265                None
266            }
267        }
268    }
269
270    /// Get the SR25519 pair for signing (if this is an SR25519 wallet)
271    pub fn sr25519_pair(&self) -> Option<&sr25519::Pair> {
272        self.sr25519_pair.as_ref()
273    }
274
275    /// Get the ED25519 pair for signing (if this is an ED25519 wallet)
276    pub fn ed25519_pair(&self) -> Option<&ed25519::Pair> {
277        self.ed25519_pair.as_ref()
278    }
279}
280
281impl std::fmt::Debug for Wallet {
282    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
283        f.debug_struct("Wallet")
284            .field("key_type", &self.key_type)
285            .field("address", &self.address())
286            .field("ss58_format", &self.ss58_format)
287            .finish()
288    }
289}
290
291/// Manager for multiple wallets
292pub struct WalletManager {
293    wallets: Arc<RwLock<HashMap<String, Wallet>>>,
294    default_key_type: KeyPairType,
295}
296
297impl WalletManager {
298    /// Create a new wallet manager
299    pub fn new() -> Self {
300        Self {
301            wallets: Arc::new(RwLock::new(HashMap::new())),
302            default_key_type: KeyPairType::Sr25519,
303        }
304    }
305
306    /// Create a new wallet manager with default key type
307    pub fn with_key_type(key_type: KeyPairType) -> Self {
308        Self {
309            wallets: Arc::new(RwLock::new(HashMap::new())),
310            default_key_type: key_type,
311        }
312    }
313
314    /// Create and add a new random wallet
315    pub fn create_wallet(&self, name: impl Into<String>) -> Wallet {
316        let wallet = Wallet::new_random_with_type(self.default_key_type);
317        let name = name.into();
318
319        debug!("Creating wallet '{}' at address {}", name, wallet.address());
320
321        self.wallets.write().insert(name.clone(), wallet.clone());
322        wallet
323    }
324
325    /// Add an existing wallet
326    pub fn add_wallet(&self, name: impl Into<String>, wallet: Wallet) {
327        let name = name.into();
328        debug!("Adding wallet '{}' at address {}", name, wallet.address());
329        self.wallets.write().insert(name, wallet);
330    }
331
332    /// Get a wallet by name
333    pub fn get_wallet(&self, name: &str) -> Option<Wallet> {
334        self.wallets.read().get(name).cloned()
335    }
336
337    /// Remove a wallet
338    pub fn remove_wallet(&self, name: &str) -> Option<Wallet> {
339        debug!("Removing wallet '{}'", name);
340        self.wallets.write().remove(name)
341    }
342
343    /// List all wallet names
344    pub fn list_wallets(&self) -> Vec<String> {
345        self.wallets.read().keys().cloned().collect()
346    }
347
348    /// Get number of wallets
349    pub fn wallet_count(&self) -> usize {
350        self.wallets.read().len()
351    }
352
353    /// Clear all wallets
354    pub fn clear(&self) {
355        debug!("Clearing all wallets");
356        self.wallets.write().clear();
357    }
358}
359
360impl Default for WalletManager {
361    fn default() -> Self {
362        Self::new()
363    }
364}
365
366#[cfg(test)]
367mod tests {
368    use super::*;
369
370    #[test]
371    fn test_create_random_wallet() {
372        let wallet = Wallet::new_random();
373        assert_eq!(wallet.key_type(), KeyPairType::Sr25519);
374        assert!(!wallet.address().is_empty());
375        assert!(!wallet.public_key().is_empty());
376    }
377
378    #[test]
379    fn test_create_wallet_types() {
380        let sr25519_wallet = Wallet::new_random_with_type(KeyPairType::Sr25519);
381        assert_eq!(sr25519_wallet.key_type(), KeyPairType::Sr25519);
382
383        let ed25519_wallet = Wallet::new_random_with_type(KeyPairType::Ed25519);
384        assert_eq!(ed25519_wallet.key_type(), KeyPairType::Ed25519);
385    }
386
387    #[test]
388    fn test_sign_and_verify() {
389        let wallet = Wallet::new_random();
390        let message = b"Hello, Substrate!";
391
392        let signature = wallet.sign(message);
393        assert_eq!(signature.len(), 64);
394
395        assert!(wallet.verify(message, &signature));
396        assert!(!wallet.verify(b"Different message", &signature));
397    }
398
399    #[test]
400    fn test_generate_mnemonic() {
401        let mnemonic = Wallet::generate_mnemonic();
402        assert!(!mnemonic.is_empty());
403
404        // Should be able to create a wallet from it
405        let wallet = Wallet::from_mnemonic(&mnemonic, KeyPairType::Sr25519);
406        assert!(wallet.is_ok());
407    }
408
409    #[test]
410    fn test_wallet_from_mnemonic() {
411        let mnemonic = "bottom drive obey lake curtain smoke basket hold race lonely fit walk";
412
413        let wallet1 = Wallet::from_mnemonic(mnemonic, KeyPairType::Sr25519).unwrap();
414        let wallet2 = Wallet::from_mnemonic(mnemonic, KeyPairType::Sr25519).unwrap();
415
416        // Same mnemonic should produce same address
417        assert_eq!(wallet1.address(), wallet2.address());
418    }
419
420    #[test]
421    fn test_wallet_with_ss58_format() {
422        let wallet = Wallet::new_random().with_ss58_format(0); // Polkadot
423        let address = wallet.address();
424
425        // Polkadot addresses start with '1'
426        assert!(address.starts_with('1'));
427    }
428
429    #[test]
430    fn test_wallet_from_seed() {
431        let seed = [42u8; 32];
432        let wallet1 = Wallet::from_seed(&seed, KeyPairType::Sr25519).unwrap();
433        let wallet2 = Wallet::from_seed(&seed, KeyPairType::Sr25519).unwrap();
434
435        // Same seed should produce same address
436        assert_eq!(wallet1.address(), wallet2.address());
437    }
438
439    #[test]
440    fn test_wallet_manager() {
441        let manager = WalletManager::new();
442
443        let wallet1 = manager.create_wallet("wallet1");
444        assert_eq!(manager.wallet_count(), 1);
445
446        let wallet2 = Wallet::new_random();
447        manager.add_wallet("wallet2", wallet2.clone());
448        assert_eq!(manager.wallet_count(), 2);
449
450        let retrieved = manager.get_wallet("wallet1").unwrap();
451        assert_eq!(retrieved.address(), wallet1.address());
452
453        let names = manager.list_wallets();
454        assert_eq!(names.len(), 2);
455        assert!(names.contains(&"wallet1".to_string()));
456        assert!(names.contains(&"wallet2".to_string()));
457
458        manager.remove_wallet("wallet1");
459        assert_eq!(manager.wallet_count(), 1);
460
461        manager.clear();
462        assert_eq!(manager.wallet_count(), 0);
463    }
464
465    #[test]
466    fn test_different_key_types_produce_different_addresses() {
467        let seed = [42u8; 32];
468        let sr25519_wallet = Wallet::from_seed(&seed, KeyPairType::Sr25519).unwrap();
469        let ed25519_wallet = Wallet::from_seed(&seed, KeyPairType::Ed25519).unwrap();
470
471        // Different cryptographic algorithms (SR25519 vs ED25519) interpret the same seed differently,
472        // resulting in distinct key pairs and thus different addresses.
473        assert_ne!(sr25519_wallet.address(), ed25519_wallet.address());
474    }
475}
apex_sdk_substrate/wallet.rs

apex_sdk_substrate/
wallet.rs
