Skip to main content

sol_safekey/
lib.rs

1//! # Sol SafeKey
2//!
3//! A powerful Solana key management library with military-grade encryption.
4//!
5//! ## Features
6//!
7//! - **Simple Encryption**: Password-based encryption for Solana private keys
8//! - **Triple-Factor Authentication**: Hardware fingerprint + master password + security question
9//! - **2FA Support**: TOTP-based two-factor authentication
10//! - **Cross-Platform**: Works on macOS, Linux, and Windows
11//!
12//! ## Quick Start
13//!
14//! ```rust
15//! use sol_safekey::{KeyManager, EncryptionResult};
16//!
17//! // Generate a new Solana keypair
18//! let keypair = KeyManager::generate_keypair();
19//!
20//! // Encrypt with password
21//! let encrypted = KeyManager::encrypt_with_password(
22//!     &keypair.to_base58_string(),
23//!     "my_strong_password"
24//! ).unwrap();
25//!
26//! // Decrypt with password
27//! let decrypted = KeyManager::decrypt_with_password(
28//!     &encrypted,
29//!     "my_strong_password"
30//! ).unwrap();
31//! ```
32
33use base64::engine::general_purpose;
34use base64::Engine;
35use ring::digest;
36
37// Re-export modules for advanced usage (conditional compilation)
38#[cfg(feature = "2fa")]
39pub mod totp;
40
41#[cfg(feature = "2fa")]
42pub mod secure_totp;
43
44#[cfg(feature = "2fa")]
45pub mod hardware_fingerprint;
46
47#[cfg(feature = "2fa")]
48pub mod security_question;
49
50// Interactive menu module - needed for bot integration
51pub mod interactive;
52
53// Bot helper module for easy bot integration (no CLI dependency)
54pub mod bot_helper;
55
56// Solana operations interactive menu
57pub mod operations;
58
59// Solana utilities for token operations
60#[cfg(any(feature = "solana-ops", feature = "sol-trade-sdk"))]
61pub mod solana_utils;
62
63// Re-export commonly used types
64pub use solana_sdk::signature::{Keypair, Signer};
65pub use solana_sdk::pubkey::Pubkey;
66
67// ============================================================================
68// Core Encryption/Decryption Functions
69// ============================================================================
70
71/// Simple XOR encryption/decryption using a 32-byte key
72fn xor_encrypt_decrypt(data: &[u8], key: &[u8; 32]) -> Vec<u8> {
73    let mut result = Vec::with_capacity(data.len());
74
75    // Generate keystream from the key
76    let mut keystream = Vec::new();
77    let mut i: u32 = 0;
78    while keystream.len() < data.len() {
79        let mut ctx = digest::Context::new(&digest::SHA256);
80        ctx.update(key);
81        ctx.update(&i.to_le_bytes());
82        let hash = ctx.finish();
83        keystream.extend_from_slice(hash.as_ref());
84        i += 1;
85    }
86
87    // XOR operation
88    for (i, &byte) in data.iter().enumerate() {
89        result.push(byte ^ keystream[i % keystream.len()]);
90    }
91
92    result
93}
94
95/// Encrypt a string with a 32-byte encryption key
96///
97/// Returns base64-encoded encrypted data
98pub fn encrypt_key(secret_key: &str, encryption_key: &[u8; 32]) -> Result<String, String> {
99    let data = secret_key.as_bytes();
100    let encrypted = xor_encrypt_decrypt(data, encryption_key);
101    Ok(general_purpose::STANDARD.encode(encrypted))
102}
103
104/// Decrypt a base64-encoded encrypted string with a 32-byte encryption key
105///
106/// Returns the original plaintext string
107pub fn decrypt_key(encrypted_data: &str, encryption_key: &[u8; 32]) -> Result<String, String> {
108    let decrypted = decrypt_key_to_bytes(encrypted_data, encryption_key)?;
109    String::from_utf8(decrypted).map_err(|_| "Invalid UTF-8 data in decrypted content".to_string())
110}
111
112/// Decrypt to raw bytes (used when plaintext may be base58 string or 64-byte keypair).
113pub fn decrypt_key_to_bytes(encrypted_data: &str, encryption_key: &[u8; 32]) -> Result<Vec<u8>, String> {
114    let ciphertext = general_purpose::STANDARD
115        .decode(encrypted_data)
116        .map_err(|_| "Invalid encrypted data format".to_string())?;
117    Ok(xor_encrypt_decrypt(&ciphertext, encryption_key))
118}
119
120/// Strip trailing bytes that are not valid base58 (e.g. 0x00, \n, \r, or non-ASCII).
121fn trim_trailing_non_base58(bytes: &[u8]) -> &[u8] {
122    let mut end = bytes.len();
123    while end > 0 {
124        let b = bytes[end - 1];
125        if b == 0x00 || b == b'\n' || b == b'\r' || b > 127 {
126            end -= 1;
127        } else {
128            break;
129        }
130    }
131    &bytes[..end]
132}
133
134/// Minimum password length for encryption/decryption
135pub const MIN_PASSWORD_LENGTH: usize = 10;
136
137/// Maximum password length for encryption/decryption
138pub const MAX_PASSWORD_LENGTH: usize = 20;
139
140/// Fixed salt for password hashing
141/// This prevents rainbow table attacks and ensures consistent key derivation
142const PASSWORD_SALT: &[u8] = b"sol-safekey-v1-salt-2025";
143
144/// Generate a 16-byte encryption key from a password using SHA-256
145///
146/// This function:
147/// 1. Combines password with fixed salt
148/// 2. Hashes using SHA-256 (produces 32 bytes)
149/// 3. Takes the first 16 bytes of the hash as the encryption key
150///
151/// Password requirements:
152/// - Minimum length: 10 characters
153/// - Maximum length: 20 characters
154pub fn generate_encryption_key_simple(password: &str) -> [u8; 32] {
155    // Combine password with fixed salt
156    let mut salted_password = password.as_bytes().to_vec();
157    salted_password.extend_from_slice(PASSWORD_SALT);
158
159    // Hash the salted password using SHA-256
160    let hash = digest::digest(&digest::SHA256, &salted_password);
161
162    // Take the first 16 bytes of the hash
163    let mut key = [0u8; 32];
164    key[0..16].copy_from_slice(&hash.as_ref()[0..16]);
165
166    // Fill the remaining 16 bytes by repeating the first 16 bytes
167    // This ensures we have a 32-byte key for compatibility
168    key[16..32].copy_from_slice(&hash.as_ref()[0..16]);
169
170    key
171}
172
173// ============================================================================
174// High-Level Key Management API (简单集成用)
175// ============================================================================
176
177/// Result type for encryption operations
178pub type EncryptionResult<T> = Result<T, String>;
179
180/// Main interface for key management operations
181///
182/// This is the recommended API for library integration.
183/// It provides simple, safe methods for common key operations.
184pub struct KeyManager;
185
186impl KeyManager {
187    /// Generate a new Solana keypair
188    ///
189    /// # Example
190    ///
191    /// ```
192    /// use sol_safekey::KeyManager;
193    ///
194    /// let keypair = KeyManager::generate_keypair();
195    /// println!("Public key: {}", keypair.pubkey());
196    /// ```
197    pub fn generate_keypair() -> Keypair {
198        Keypair::new()
199    }
200
201    /// Encrypt a private key with a password
202    ///
203    /// # Arguments
204    ///
205    /// * `private_key` - The private key in base58 string format
206    /// * `password` - The password to use for encryption
207    ///
208    /// # Returns
209    ///
210    /// Base64-encoded encrypted string
211    ///
212    /// # Example
213    ///
214    /// ```
215    /// use sol_safekey::KeyManager;
216    ///
217    /// let keypair = KeyManager::generate_keypair();
218    /// let private_key = keypair.to_base58_string();
219    ///
220    /// let encrypted = KeyManager::encrypt_with_password(
221    ///     &private_key,
222    ///     "my_password"
223    /// ).unwrap();
224    /// ```
225    pub fn encrypt_with_password(private_key: &str, password: &str) -> EncryptionResult<String> {
226        let key = generate_encryption_key_simple(password);
227        encrypt_key(private_key, &key)
228    }
229
230    /// Decrypt a private key with a password
231    ///
232    /// # Arguments
233    ///
234    /// * `encrypted_data` - Base64-encoded encrypted data
235    /// * `password` - The password used for encryption
236    ///
237    /// # Returns
238    ///
239    /// The original private key in base58 string format
240    ///
241    /// # Example
242    ///
243    /// ```
244    /// use sol_safekey::KeyManager;
245    ///
246    /// let encrypted = "..."; // from encryption
247    /// let decrypted = KeyManager::decrypt_with_password(
248    ///     encrypted,
249    ///     "my_password"
250    /// ).unwrap();
251    /// ```
252    pub fn decrypt_with_password(encrypted_data: &str, password: &str) -> EncryptionResult<String> {
253        let key = generate_encryption_key_simple(password);
254        decrypt_key(encrypted_data, &key)
255    }
256
257    /// Get public key from a private key
258    ///
259    /// # Arguments
260    ///
261    /// * `private_key` - Private key in base58 string format
262    ///
263    /// # Returns
264    ///
265    /// Public key as a base58 string
266    pub fn get_public_key(private_key: &str) -> EncryptionResult<String> {
267        use solana_sdk::signature::Keypair;
268
269        // Solana 3.0 uses from_base58_string directly
270        let keypair = Keypair::from_base58_string(private_key);
271
272        Ok(keypair.pubkey().to_string())
273    }
274
275    /// Encrypt a keypair to a JSON keystore format
276    ///
277    /// This creates a standard encrypted keystore file compatible with Solana tools.
278    ///
279    /// # Arguments
280    ///
281    /// * `keypair` - The Solana keypair to encrypt
282    /// * `password` - The password for encryption
283    ///
284    /// # Returns
285    ///
286    /// JSON string containing the encrypted keystore
287    pub fn keypair_to_encrypted_json(keypair: &Keypair, password: &str) -> EncryptionResult<String> {
288        use serde_json::json;
289        use chrono::Utc;
290
291        let private_key = keypair.to_base58_string();
292        let public_key = keypair.pubkey().to_string();
293
294        let encrypted = Self::encrypt_with_password(&private_key, password)?;
295
296        let keystore = json!({
297            "encrypted_private_key": encrypted,
298            "public_key": public_key,
299            "encryption_type": "password_only",
300            "created_at": Utc::now().to_rfc3339(),
301        });
302
303        Ok(keystore.to_string())
304    }
305
306    /// Decrypt a keypair from encrypted JSON keystore.
307    /// 与 GitHub 最新版兼容:优先按「base58 私钥」解密;若解密结果非 UTF-8,再尝试「64 字节 keypair」。
308    pub fn keypair_from_encrypted_json(json_data: &str, password: &str) -> EncryptionResult<Keypair> {
309        use serde_json::Value;
310
311        let data: Value = serde_json::from_str(json_data)
312            .map_err(|_| "Invalid JSON format")?;
313
314        let encrypted = data["encrypted_private_key"]
315            .as_str()
316            .ok_or("Missing encrypted_private_key field")?;
317
318        let key = generate_encryption_key_simple(password);
319        let decrypted = decrypt_key_to_bytes(encrypted, &key)
320            .map_err(|e| format!("解密失败: {}", e))?;
321
322        // 1) 标准格式:UTF-8 base58 私钥(keypair_to_encrypted_json 生成)。去除首尾空白避免换行/空格导致解析失败
323        if let Ok(s) = String::from_utf8(decrypted.clone()) {
324            let s = s.trim_end_matches(|c| c == '\n' || c == '\r').trim();
325            if !s.is_empty() {
326                let keypair = Keypair::from_base58_string(s);
327                return Ok(keypair);
328            }
329        }
330
331        // 2) 兼容:部分工具加密的是 64 字节原始 keypair,解密后非 UTF-8
332        if decrypted.len() == 64 {
333            if let Ok(k) = Keypair::try_from(decrypted.as_slice()) {
334                return Ok(k);
335            }
336        }
337
338        // 3) 解密长度在典型 base58 范围内 (80–96):可能带有尾部垃圾,先剥掉再试
339        if (80..=96).contains(&decrypted.len()) {
340            let trimmed = trim_trailing_non_base58(&decrypted);
341            if let Ok(s) = String::from_utf8(trimmed.to_vec()) {
342                let s = s.trim_end_matches(|c| c == '\n' || c == '\r').trim();
343                if !s.is_empty() {
344                    let keypair = Keypair::from_base58_string(s);
345                    return Ok(keypair);
346                }
347            }
348        }
349
350        Err(format!(
351            "解密结果既非有效 base58 私钥,也非 64 字节 keypair(解密长度: {})。请确认:1) 密码正确 2) keystore 由 sol-safekey 生成 3) 密码长度 10–20 字符",
352            decrypted.len()
353        ))
354    }
355}
356
357// ============================================================================
358// Advanced 2FA Functions (CLI 工具使用,库集成可选)
359// ============================================================================
360
361// ============================================================================
362// 2FA Functions (only available with "2fa" feature)
363// ============================================================================
364
365#[cfg(feature = "2fa")]
366/// Derive a TOTP secret from password
367///
368/// This is used internally for deterministic 2FA key generation.
369#[allow(dead_code)]
370fn derive_totp_secret_from_password(password: &str, account: &str, issuer: &str) -> Result<String, String> {
371    use ring::pbkdf2;
372    use data_encoding::BASE32_NOPAD;
373    use std::num::NonZeroU32;
374
375    let salt = format!("sol-safekey-totp-{}-{}", issuer, account);
376    let iterations = NonZeroU32::new(100_000)
377        .ok_or("Invalid iteration count")?;
378
379    let mut secret = [0u8; 20]; // 160 bits for TOTP
380    pbkdf2::derive(
381        pbkdf2::PBKDF2_HMAC_SHA256,
382        iterations,
383        salt.as_bytes(),
384        password.as_bytes(),
385        &mut secret,
386    );
387
388    Ok(BASE32_NOPAD.encode(&secret))
389}
390
391#[cfg(feature = "2fa")]
392/// Derive TOTP secret from hardware fingerprint and password
393///
394/// This creates a deterministic 2FA key bound to specific hardware.
395pub fn derive_totp_secret_from_hardware_and_password(
396    hardware_fingerprint: &str,
397    master_password: &str,
398    account: &str,
399    issuer: &str,
400) -> Result<String, String> {
401    use ring::pbkdf2;
402    use data_encoding::BASE32_NOPAD;
403    use std::num::NonZeroU32;
404
405    let key_material = format!("{}::{}", hardware_fingerprint, master_password);
406    let salt = format!("sol-safekey-2fa-{}-{}", issuer, account);
407    let iterations = NonZeroU32::new(100_000)
408        .ok_or("Invalid iteration count")?;
409
410    let mut secret = [0u8; 20];
411    pbkdf2::derive(
412        pbkdf2::PBKDF2_HMAC_SHA256,
413        iterations,
414        salt.as_bytes(),
415        key_material.as_bytes(),
416        &mut secret,
417    );
418
419    Ok(BASE32_NOPAD.encode(&secret))
420}
421
422#[cfg(feature = "2fa")]
423/// Verify a TOTP code
424fn verify_current_totp_code(totp_secret: &str, current_code: &str) -> Result<(), String> {
425    use crate::totp::{TOTPConfig, TOTPManager};
426
427    let config = TOTPConfig {
428        secret: totp_secret.to_string(),
429        account: "wallet".to_string(),
430        issuer: "Sol-SafeKey".to_string(),
431        algorithm: "SHA1".to_string(),
432        digits: 6,
433        step: 30,
434    };
435
436    let totp_manager = TOTPManager::new(config);
437
438    match totp_manager.verify_code(current_code) {
439        Ok(true) => Ok(()),
440        Ok(false) => Err("验证失败,请检查主密码、安全问题答案或2FA验证码".to_string()),
441        Err(e) => Err(format!("验证失败: {}", e)),
442    }
443}
444
445// ============================================================================
446// Triple-Factor Encryption (only available with "2fa" feature)
447// ============================================================================
448
449#[cfg(feature = "2fa")]
450/// Generate a triple-factor encryption key
451///
452/// Combines hardware fingerprint + master password + security answer
453pub fn generate_triple_factor_key(
454    hardware_fingerprint: &str,
455    master_password: &str,
456    security_answer: &str,
457) -> [u8; 32] {
458    use ring::pbkdf2;
459    use std::num::NonZeroU32;
460
461    let key_material = format!(
462        "HW:{}|PASS:{}|QA:{}",
463        hardware_fingerprint,
464        master_password,
465        security_answer.trim().to_lowercase()
466    );
467
468    let salt = b"sol-safekey-triple-factor-v1";
469    let iterations = NonZeroU32::new(200_000).unwrap();
470
471    let mut key = [0u8; 32];
472    pbkdf2::derive(
473        pbkdf2::PBKDF2_HMAC_SHA256,
474        iterations,
475        salt,
476        key_material.as_bytes(),
477        &mut key,
478    );
479
480    key
481}
482
483#[cfg(feature = "2fa")]
484/// Encrypt with triple-factor authentication
485///
486/// Used by CLI for maximum security with device binding.
487pub fn encrypt_with_triple_factor(
488    private_key: &str,
489    twofa_secret: &str,
490    hardware_fingerprint: &str,
491    master_password: &str,
492    question_index: usize,
493    security_answer: &str,
494) -> Result<String, String> {
495    use serde_json::json;
496
497    let encryption_key = generate_triple_factor_key(
498        hardware_fingerprint,
499        master_password,
500        security_answer,
501    );
502
503    let data_package = json!({
504        "private_key": private_key,
505        "twofa_secret": twofa_secret,
506        "question_index": question_index,
507        "version": "triple_factor_v1",
508        "created_at": std::time::SystemTime::now()
509            .duration_since(std::time::UNIX_EPOCH)
510            .unwrap()
511            .as_secs()
512    });
513
514    let package_str = data_package.to_string();
515    let encrypted = encrypt_key(&package_str, &encryption_key)?;
516
517    Ok(encrypted)
518}
519
520#[cfg(feature = "2fa")]
521/// Decrypt with triple-factor authentication and verify 2FA code
522///
523/// Used by CLI for unlocking triple-factor encrypted wallets.
524pub fn decrypt_with_triple_factor_and_2fa(
525    encrypted_data: &str,
526    hardware_fingerprint: &str,
527    master_password: &str,
528    security_answer: &str,
529    twofa_code: &str,
530) -> Result<(String, String, usize), String> {
531    let decryption_key = generate_triple_factor_key(
532        hardware_fingerprint,
533        master_password,
534        security_answer,
535    );
536
537    let decrypted = decrypt_key(encrypted_data, &decryption_key)
538        .map_err(|_| "解密失败,请检查主密码、安全问题答案是否正确")?;
539
540    let data: serde_json::Value = serde_json::from_str(&decrypted)
541        .map_err(|_| "解密失败,请检查主密码、安全问题答案是否正确")?;
542
543    let private_key = data["private_key"]
544        .as_str()
545        .ok_or("缺少私钥数据")?
546        .to_string();
547
548    let twofa_secret = data["twofa_secret"]
549        .as_str()
550        .ok_or("缺少2FA密钥数据")?
551        .to_string();
552
553    let question_index = data["question_index"]
554        .as_u64()
555        .ok_or("缺少安全问题索引")? as usize;
556
557    // Verify 2FA code
558    verify_current_totp_code(&twofa_secret, twofa_code)?;
559
560    Ok((private_key, twofa_secret, question_index))
561}
562
563// ============================================================================
564// Tests
565// ============================================================================
566
567#[cfg(test)]
568mod tests {
569    use super::*;
570
571    #[test]
572    fn test_generate_keypair() {
573        let keypair = KeyManager::generate_keypair();
574        assert_eq!(keypair.to_bytes().len(), 64);
575    }
576
577    #[test]
578    fn test_encrypt_decrypt_with_password() {
579        let keypair = KeyManager::generate_keypair();
580        let private_key = keypair.to_base58_string();
581        let password = "test_password_123";
582
583        let encrypted = KeyManager::encrypt_with_password(&private_key, password).unwrap();
584        let decrypted = KeyManager::decrypt_with_password(&encrypted, password).unwrap();
585
586        assert_eq!(private_key, decrypted);
587    }
588
589    #[test]
590    fn test_get_public_key() {
591        let keypair = KeyManager::generate_keypair();
592        let private_key = keypair.to_base58_string();
593        let expected_pubkey = keypair.pubkey().to_string();
594
595        let pubkey = KeyManager::get_public_key(&private_key).unwrap();
596        assert_eq!(pubkey, expected_pubkey);
597    }
598
599    #[test]
600    fn test_keystore_json_round_trip() {
601        let keypair = KeyManager::generate_keypair();
602        let password = "secure_password";
603
604        let json = KeyManager::keypair_to_encrypted_json(&keypair, password).unwrap();
605        let restored_keypair = KeyManager::keypair_from_encrypted_json(&json, password).unwrap();
606
607        assert_eq!(keypair.to_bytes(), restored_keypair.to_bytes());
608    }
609
610    #[test]
611    fn test_wrong_password_fails() {
612        let keypair = KeyManager::generate_keypair();
613        let private_key = keypair.to_base58_string();
614
615        let encrypted = KeyManager::encrypt_with_password(&private_key, "correct").unwrap();
616        let result = KeyManager::decrypt_with_password(&encrypted, "wrong");
617
618        assert!(result.is_err());
619    }
620}