1use base64::engine::general_purpose;
34use base64::Engine;
35use ring::digest;
36
37#[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
50pub mod interactive;
52
53pub mod bot_helper;
55
56pub mod operations;
58
59#[cfg(any(feature = "solana-ops", feature = "sol-trade-sdk"))]
61pub mod solana_utils;
62
63pub use solana_sdk::signature::{Keypair, Signer};
65pub use solana_sdk::pubkey::Pubkey;
66
67fn xor_encrypt_decrypt(data: &[u8], key: &[u8; 32]) -> Vec<u8> {
73 let mut result = Vec::with_capacity(data.len());
74
75 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 for (i, &byte) in data.iter().enumerate() {
89 result.push(byte ^ keystream[i % keystream.len()]);
90 }
91
92 result
93}
94
95pub 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
104pub 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
112pub 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
120fn 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
134pub const MIN_PASSWORD_LENGTH: usize = 10;
136
137pub const MAX_PASSWORD_LENGTH: usize = 20;
139
140const PASSWORD_SALT: &[u8] = b"sol-safekey-v1-salt-2025";
143
144pub fn generate_encryption_key_simple(password: &str) -> [u8; 32] {
155 let mut salted_password = password.as_bytes().to_vec();
157 salted_password.extend_from_slice(PASSWORD_SALT);
158
159 let hash = digest::digest(&digest::SHA256, &salted_password);
161
162 let mut key = [0u8; 32];
164 key[0..16].copy_from_slice(&hash.as_ref()[0..16]);
165
166 key[16..32].copy_from_slice(&hash.as_ref()[0..16]);
169
170 key
171}
172
173pub type EncryptionResult<T> = Result<T, String>;
179
180pub struct KeyManager;
185
186impl KeyManager {
187 pub fn generate_keypair() -> Keypair {
198 Keypair::new()
199 }
200
201 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 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 pub fn get_public_key(private_key: &str) -> EncryptionResult<String> {
267 use solana_sdk::signature::Keypair;
268
269 let keypair = Keypair::from_base58_string(private_key);
271
272 Ok(keypair.pubkey().to_string())
273 }
274
275 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 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 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 if decrypted.len() == 64 {
333 if let Ok(k) = Keypair::try_from(decrypted.as_slice()) {
334 return Ok(k);
335 }
336 }
337
338 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#[cfg(feature = "2fa")]
366#[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]; 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")]
392pub 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")]
423fn 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#[cfg(feature = "2fa")]
450pub 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")]
484pub 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")]
521pub 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_current_totp_code(&twofa_secret, twofa_code)?;
559
560 Ok((private_key, twofa_secret, question_index))
561}
562
563#[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}