Skip to main content

sol_safekey/
bot_helper.rs

1//! Bot Helper Module
2//!
3//! Provides easy-to-use functions for bot integration with sol-safekey.
4//! Bots can integrate the interactive key management tool with just 1 line of code.
5//! **No CLI dependency required** - uses the library directly.
6//!
7//! # Quick Start
8//!
9//! ```no_run
10//! use sol_safekey::bot_helper;
11//!
12//! let keypair = bot_helper::ensure_wallet_ready("wallet.json").unwrap();
13//! ```
14
15use std::path::Path;
16use std::fs;
17use std::io::Write;
18use serde_json::Value;
19use solana_sdk::signature::Keypair;
20use solana_sdk::signer::Signer;
21use crate::{decrypt_key, generate_encryption_key_simple, interactive};
22
23/// Result type for bot helper operations
24pub type Result<T> = std::result::Result<T, String>;
25
26/// Ensure wallet is ready to use - creates or unlocks interactively
27///
28/// Main function for bot integration:
29/// - If wallet exists: prompts for password and unlocks
30/// - If wallet doesn't exist: launches interactive creation, then unlocks
31///
32/// # Arguments
33///
34/// * `wallet_path` - Path to the encrypted wallet file
35///
36/// # Example
37///
38/// ```no_run
39/// use sol_safekey::bot_helper;
40/// use solana_sdk::signer::Signer;
41///
42/// let keypair = bot_helper::ensure_wallet_ready("wallet.json")?;
43/// println!("Wallet ready: {}", keypair.pubkey());
44/// # Ok::<(), String>(())
45/// ```
46pub fn ensure_wallet_ready(wallet_path: &str) -> Result<Keypair> {
47    if wallet_exists(wallet_path) {
48        println!("โœ… Wallet found at: {}", wallet_path);
49        println!("๐Ÿ”“ Starting interactive wallet unlock...\n");
50        unlock_wallet(wallet_path)
51    } else {
52        println!("โš ๏ธ  Wallet not found at: {}", wallet_path);
53        println!("๐Ÿ“ Starting interactive wallet creation...\n");
54        create_wallet(wallet_path)?;
55
56        println!("\nNow unlocking the newly created wallet...\n");
57        unlock_wallet(wallet_path)
58    }
59}
60
61/// Check if a wallet file exists at the given path
62pub fn wallet_exists(path: &str) -> bool {
63    Path::new(path).exists()
64}
65
66/// Get public key from wallet without unlocking
67///
68/// # Arguments
69///
70/// * `wallet_path` - Path to the encrypted wallet file
71pub fn get_wallet_pubkey(wallet_path: &str) -> Result<String> {
72    let content = fs::read_to_string(wallet_path)
73        .map_err(|e| format!("Failed to read wallet: {}", e))?;
74
75    let data: Value = serde_json::from_str(&content)
76        .map_err(|e| format!("Invalid wallet format: {}", e))?;
77
78    data["public_key"]
79        .as_str()
80        .map(String::from)
81        .ok_or_else(|| "Public key not found in wallet file".to_string())
82}
83
84/// Create wallet interactively
85fn create_wallet(output_path: &str) -> Result<()> {
86    println!("๐Ÿ” Interactive Wallet Creation");
87    println!("โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”");
88    println!("\nYou will be guided through the wallet creation process.");
89    println!("The wallet will be saved to: {}\n", output_path);
90
91    interactive::show_main_menu()?;
92
93    if !wallet_exists(output_path) {
94        println!("\nโš ๏ธ  Note: Wallet was not saved to the expected path.");
95        println!("   Expected: {}", output_path);
96        println!("   Please make sure to save your wallet to this location.");
97        return Err(format!("Wallet not created at expected path: {}", output_path));
98    }
99
100    println!("\nโœ… Wallet created successfully!");
101    println!("๐Ÿ“ Location: {}", output_path);
102    Ok(())
103}
104
105/// Unlock wallet interactively (prompts for password)
106fn unlock_wallet(wallet_path: &str) -> Result<Keypair> {
107    if !wallet_exists(wallet_path) {
108        return Err(format!("Wallet file not found: {}", wallet_path));
109    }
110
111    println!("๐Ÿ”“ Unlocking wallet: {}", wallet_path);
112
113    let content = fs::read_to_string(wallet_path)
114        .map_err(|e| format!("Failed to read wallet: {}", e))?;
115
116    let data: Value = serde_json::from_str(&content)
117        .map_err(|e| format!("Invalid wallet format: {}", e))?;
118
119    let encrypted_key = data["encrypted_private_key"]
120        .as_str()
121        .ok_or_else(|| "Encrypted private key not found in wallet file".to_string())?;
122
123    print!("๐Ÿ”‘ Enter wallet password: ");
124    std::io::stdout().flush().unwrap();
125
126    // ไธŽ wick-bot ๅฎŒๅ…จไธ€่‡ด๏ผšๅฏ†็ ๅŽŸๆ ทไฝฟ็”จ๏ผŒไธๅš trim๏ผˆwick-bot ็š„ bot_helper ไบฆไธ trim๏ผ‰
127    let password = rpassword::read_password()
128        .map_err(|e| format!("Failed to read password: {}", e))?;
129
130    let encryption_key = generate_encryption_key_simple(&password);
131    let private_key = decrypt_key(encrypted_key, &encryption_key)?;
132    let keypair = Keypair::from_base58_string(&private_key);
133
134    println!("โœ… Wallet unlocked successfully!");
135    println!("๐Ÿ“ Address: {}", keypair.pubkey());
136
137    Ok(keypair)
138}