diff --git a/Cargo.lock b/Cargo.lock index d49183a..d9cfa97 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,17 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + [[package]] name = "ahash" version = "0.7.8" @@ -472,6 +483,7 @@ dependencies = [ "alloy-primitives", "alloy-signer", "async-trait", + "eth-keystore", "k256", "rand 0.8.5", "thiserror 2.0.18", @@ -1235,6 +1247,16 @@ dependencies = [ "windows-link", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", +] + [[package]] name = "clap" version = "4.5.60" @@ -1440,6 +1462,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher", +] + [[package]] name = "darling" version = "0.21.3" @@ -1762,6 +1793,28 @@ version = "3.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59" +[[package]] +name = "eth-keystore" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fda3bf123be441da5260717e0661c25a2fd9cb2b2c1d20bf2e05580047158ab" +dependencies = [ + "aes", + "ctr", + "digest 0.10.7", + "hex", + "hmac", + "pbkdf2", + "rand 0.8.5", + "scrypt", + "serde", + "serde_json", + "sha2", + "sha3", + "thiserror 1.0.69", + "uuid 0.8.2", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -1798,7 +1851,7 @@ checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78" dependencies = [ "cfg-if", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -2430,6 +2483,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -2863,6 +2925,15 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pbkdf2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917" +dependencies = [ + "digest 0.10.7", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -2976,9 +3047,12 @@ dependencies = [ "dirs", "polymarket-client-sdk", "predicates", + "rand 0.8.5", + "rpassword", "rust_decimal", "rust_decimal_macros", "rustyline", + "secrecy", "serde", "serde_json", "tabled", @@ -3014,7 +3088,7 @@ dependencies = [ "sha2", "strum_macros", "url", - "uuid", + "uuid 1.21.0", ] [[package]] @@ -3540,7 +3614,7 @@ dependencies = [ "rkyv_derive", "seahash", "tinyvec", - "uuid", + "uuid 1.21.0", ] [[package]] @@ -3564,6 +3638,27 @@ dependencies = [ "rustc-hex", ] +[[package]] +name = "rpassword" +version = "7.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66d4c8b64f049c6721ec8ccec37ddfc3d641c4a7fca57e8f2a89de509c73df39" +dependencies = [ + "libc", + "rtoolbox", + "windows-sys 0.59.0", +] + +[[package]] +name = "rtoolbox" +version = "0.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7cc970b249fbe527d6e02e0a227762c9108b2f49d81094fe357ffc6d14d7f6f" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "ruint" version = "1.17.2" @@ -3789,6 +3884,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "salsa20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" +dependencies = [ + "cipher", +] + [[package]] name = "same-file" version = "1.0.6" @@ -3837,6 +3941,18 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scrypt" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f9e24d2b632954ded8ab2ef9fea0a0c769ea56ea98bddbafbad22caeeadf45d" +dependencies = [ + "hmac", + "pbkdf2", + "salsa20", + "sha2", +] + [[package]] name = "seahash" version = "4.1.0" @@ -4681,6 +4797,16 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc5cf98d8186244414c848017f0e2676b3fcb46807f6668a97dfe67359a3c4b7" +dependencies = [ + "getrandom 0.2.17", + "serde", +] + [[package]] name = "uuid" version = "1.21.0" diff --git a/Cargo.toml b/Cargo.toml index b7d516f..caca1e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ path = "src/main.rs" [dependencies] polymarket-client-sdk = { version = "0.4", features = ["gamma", "data", "bridge", "clob", "ctf"] } -alloy = { version = "1.6.3", default-features = false, features = ["providers", "sol-types", "contract", "reqwest", "reqwest-rustls-tls", "signer-local", "signers"] } +alloy = { version = "1.6.3", default-features = false, features = ["providers", "sol-types", "contract", "reqwest", "reqwest-rustls-tls", "signer-local", "signers", "signer-keystore"] } clap = { version = "4", features = ["derive"] } tokio = { version = "1", features = ["rt-multi-thread", "macros"] } serde_json = "1" @@ -26,6 +26,9 @@ anyhow = "1" chrono = "0.4" dirs = "6" rustyline = "15" +rpassword = "7" +rand = "0.8" +secrecy = "0.10" [dev-dependencies] assert_cmd = "2" diff --git a/src/auth.rs b/src/auth.rs index 15ad61e..e475e7c 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -6,6 +6,7 @@ use polymarket_client_sdk::auth::state::Authenticated; use polymarket_client_sdk::auth::{LocalSigner, Normal, Signer as _}; use polymarket_client_sdk::clob::types::SignatureType; use polymarket_client_sdk::{POLYGON, clob}; +use secrecy::{ExposeSecret, SecretString}; use crate::config; @@ -19,12 +20,47 @@ fn parse_signature_type(s: &str) -> SignatureType { } } +/// Resolve the private key hex string, prompting for password if needed. +pub(crate) fn resolve_key_string(private_key: Option<&str>) -> Result { + // 1. CLI flag (highest priority — never overridden by migration) + if let Some(key) = private_key { + return Ok(SecretString::from(key.to_string())); + } + // 2. Env var + if let Ok(key) = std::env::var(config::ENV_VAR) + && !key.is_empty() + { + return Ok(SecretString::from(key)); + } + + // Auto-migrate plaintext config to encrypted keystore + if config::needs_migration() { + eprintln!("Your wallet key is stored in plaintext. Encrypting it now..."); + let password = crate::password::prompt_new_password()?; + config::migrate_to_encrypted(&password)?; + eprintln!("Wallet key encrypted successfully."); + return config::load_key_encrypted(password.expose_secret()); + } + // 3. Old config (plaintext — for backward compat) + if let Some(cfg) = config::load_config() + && !cfg.private_key.is_empty() + { + return Ok(SecretString::from(cfg.private_key)); + } + // 4. Encrypted keystore with retry + if config::keystore_exists() { + return crate::password::prompt_password_with_retries(|pw| { + config::load_key_encrypted(pw) + }); + } + anyhow::bail!("{}", config::NO_WALLET_MSG) +} + pub fn resolve_signer( private_key: Option<&str>, ) -> Result { - let (key, _) = config::resolve_key(private_key); - let key = key.ok_or_else(|| anyhow::anyhow!("{}", config::NO_WALLET_MSG))?; - LocalSigner::from_str(&key) + let key = resolve_key_string(private_key)?; + LocalSigner::from_str(key.expose_secret()) .context("Invalid private key") .map(|s| s.with_chain_id(Some(POLYGON))) } @@ -61,9 +97,8 @@ pub async fn create_readonly_provider() -> Result, ) -> Result { - let (key, _) = config::resolve_key(private_key); - let key = key.ok_or_else(|| anyhow::anyhow!("{}", config::NO_WALLET_MSG))?; - let signer = LocalSigner::from_str(&key) + let key = resolve_key_string(private_key)?; + let signer = LocalSigner::from_str(key.expose_secret()) .context("Invalid private key")? .with_chain_id(Some(POLYGON)); ProviderBuilder::new() diff --git a/src/commands/setup.rs b/src/commands/setup.rs index dd04671..c69c104 100644 --- a/src/commands/setup.rs +++ b/src/commands/setup.rs @@ -6,6 +6,7 @@ use anyhow::{Context, Result}; use polymarket_client_sdk::auth::{LocalSigner, Signer as _}; use polymarket_client_sdk::types::Address; use polymarket_client_sdk::{POLYGON, derive_proxy_wallet}; +use secrecy::{ExposeSecret, SecretString}; use super::wallet::normalize_key; use crate::config; @@ -84,12 +85,31 @@ pub fn execute() -> Result<()> { step_header(1, total, "Wallet"); - let address = if config::config_exists() { - let (key, source) = config::resolve_key(None); - if let Some(k) = &key - && let Ok(signer) = LocalSigner::from_str(k) - { - let addr = signer.address(); + let address = if config::config_exists() || config::keystore_exists() { + // Try plaintext config / env var / flag first, then encrypted keystore + let (existing_addr, source) = { + let (key, src) = config::resolve_key(None); + let addr = key + .as_ref() + .and_then(|k| LocalSigner::from_str(k.expose_secret()).ok()) + .map(|s| s.address()); + (addr, src) + }; + let (existing_addr, source) = if existing_addr.is_some() { + (existing_addr, source) + } else if config::keystore_exists() { + let addr = crate::password::prompt_password_with_retries(|pw| { + config::load_key_encrypted(pw) + }) + .ok() + .and_then(|k| LocalSigner::from_str(k.expose_secret()).ok()) + .map(|s| s.address()); + (addr, config::KeySource::Keystore) + } else { + (None, source) + }; + + if let Some(addr) = existing_addr { println!(" ✓ Wallet already configured ({})", source.label()); println!(" Address: {addr}"); println!(); @@ -119,7 +139,7 @@ fn setup_wallet() -> Result
{ let signer = LocalSigner::from_str(&normalized) .context("Invalid private key")? .with_chain_id(Some(POLYGON)); - (signer.address(), normalized) + (signer.address(), SecretString::from(normalized)) } else { let signer = LocalSigner::random().with_chain_id(Some(POLYGON)); let address = signer.address(); @@ -129,10 +149,12 @@ fn setup_wallet() -> Result
{ for b in &bytes { write!(hex, "{b:02x}").unwrap(); } - (address, hex) + (address, SecretString::from(hex)) }; - config::save_wallet(&key_hex, POLYGON, config::DEFAULT_SIGNATURE_TYPE)?; + let password = crate::password::prompt_new_password()?; + config::save_key_encrypted(&key_hex, &password)?; + config::save_wallet_settings(POLYGON, config::DEFAULT_SIGNATURE_TYPE)?; if has_key { println!(" ✓ Wallet imported"); @@ -144,7 +166,7 @@ fn setup_wallet() -> Result
{ if !has_key { println!(); - println!(" ⚠ Back up your private key from the config file."); + println!(" ⚠ Remember your password. Use `polymarket wallet export` to back up your key."); println!(" If lost, your funds cannot be recovered."); } diff --git a/src/commands/wallet.rs b/src/commands/wallet.rs index ce7597f..bdb46e5 100644 --- a/src/commands/wallet.rs +++ b/src/commands/wallet.rs @@ -6,6 +6,7 @@ use clap::{Args, Subcommand}; use polymarket_client_sdk::auth::LocalSigner; use polymarket_client_sdk::auth::Signer as _; use polymarket_client_sdk::{POLYGON, derive_proxy_wallet}; +use secrecy::{ExposeSecret, SecretString}; use crate::config; use crate::output::OutputFormat; @@ -48,6 +49,8 @@ pub enum WalletCommand { #[arg(long)] force: bool, }, + /// Export the private key (decrypts keystore, prints to stdout) + Export, } pub fn execute( @@ -68,14 +71,14 @@ pub fn execute( WalletCommand::Address => cmd_address(output, private_key_flag), WalletCommand::Show => cmd_show(output, private_key_flag), WalletCommand::Reset { force } => cmd_reset(output, force), + WalletCommand::Export => cmd_export(output), } } fn guard_overwrite(force: bool) -> Result<()> { - if !force && config::config_exists() { + if !force && (config::config_exists() || config::keystore_exists()) { bail!( - "A wallet already exists at {}. Use --force to overwrite.", - config::config_path()?.display() + "A wallet already exists. Use --force to overwrite.", ); } Ok(()) @@ -95,13 +98,16 @@ fn cmd_create(output: &OutputFormat, force: bool, signature_type: &str) -> Resul let signer = LocalSigner::random().with_chain_id(Some(POLYGON)); let address = signer.address(); let bytes = signer.credential().to_bytes(); - let mut key_hex = String::with_capacity(2 + bytes.len() * 2); - key_hex.push_str("0x"); + let mut hex = String::with_capacity(2 + bytes.len() * 2); + hex.push_str("0x"); for b in &bytes { - write!(key_hex, "{b:02x}").unwrap(); + write!(hex, "{b:02x}").unwrap(); } + let key_hex = SecretString::from(hex); - config::save_wallet(&key_hex, POLYGON, signature_type)?; + let password = crate::password::prompt_new_password()?; + config::save_key_encrypted(&key_hex, &password)?; + config::save_wallet_settings(POLYGON, signature_type)?; let config_path = config::config_path()?; let proxy_addr = derive_proxy_wallet(address, POLYGON); @@ -126,8 +132,8 @@ fn cmd_create(output: &OutputFormat, force: bool, signature_type: &str) -> Resul println!("Signature type: {signature_type}"); println!("Config: {}", config_path.display()); println!(); - println!("IMPORTANT: Back up your private key from the config file."); - println!(" If lost, your funds cannot be recovered."); + println!("IMPORTANT: Remember your password. Use `polymarket wallet export`"); + println!(" to back up your private key. If lost, funds cannot be recovered."); } } Ok(()) @@ -141,8 +147,11 @@ fn cmd_import(key: &str, output: &OutputFormat, force: bool, signature_type: &st .context("Invalid private key")? .with_chain_id(Some(POLYGON)); let address = signer.address(); + let normalized = SecretString::from(normalized); - config::save_wallet(&normalized, POLYGON, signature_type)?; + let password = crate::password::prompt_new_password()?; + config::save_key_encrypted(&normalized, &password)?; + config::save_wallet_settings(POLYGON, signature_type)?; let config_path = config::config_path()?; let proxy_addr = derive_proxy_wallet(address, POLYGON); @@ -172,10 +181,9 @@ fn cmd_import(key: &str, output: &OutputFormat, force: bool, signature_type: &st } fn cmd_address(output: &OutputFormat, private_key_flag: Option<&str>) -> Result<()> { - let (key, _) = config::resolve_key(private_key_flag); - let key = key.ok_or_else(|| anyhow::anyhow!("{}", config::NO_WALLET_MSG))?; + let key = crate::auth::resolve_key_string(private_key_flag)?; - let signer = LocalSigner::from_str(&key).context("Invalid private key")?; + let signer = LocalSigner::from_str(key.expose_secret()).context("Invalid private key")?; let address = signer.address(); match output { @@ -190,8 +198,21 @@ fn cmd_address(output: &OutputFormat, private_key_flag: Option<&str>) -> Result< } fn cmd_show(output: &OutputFormat, private_key_flag: Option<&str>) -> Result<()> { - let (key, source) = config::resolve_key(private_key_flag); - let signer = key.as_deref().and_then(|k| LocalSigner::from_str(k).ok()); + let (key_result, source) = { + let (old_key, old_source) = config::resolve_key(private_key_flag); + if old_key.as_ref().is_some_and(|k| !k.expose_secret().is_empty()) { + (Ok(old_key.unwrap()), old_source) + } else if config::keystore_exists() { + let result = crate::password::prompt_password_with_retries(|pw| { + config::load_key_encrypted(pw) + }); + (result, config::KeySource::Keystore) + } else { + (Err(anyhow::anyhow!("not configured")), config::KeySource::None) + } + }; + + let signer = key_result.ok().and_then(|k| LocalSigner::from_str(k.expose_secret()).ok()); let address = signer.as_ref().map(|s| s.address().to_string()); let proxy_addr = signer .as_ref() @@ -231,8 +252,28 @@ fn cmd_show(output: &OutputFormat, private_key_flag: Option<&str>) -> Result<()> Ok(()) } +fn cmd_export(output: &OutputFormat) -> Result<()> { + if !config::keystore_exists() { + bail!("{}", config::NO_WALLET_MSG); + } + + let key = crate::password::prompt_password_with_retries(|pw| { + config::load_key_encrypted(pw) + })?; + + match output { + OutputFormat::Json => { + println!("{}", serde_json::json!({"private_key": key.expose_secret()})); + } + OutputFormat::Table => { + println!("{}", key.expose_secret()); + } + } + Ok(()) +} + fn cmd_reset(output: &OutputFormat, force: bool) -> Result<()> { - if !config::config_exists() { + if !config::config_exists() && !config::keystore_exists() { match output { OutputFormat::Table => println!("Nothing to reset. No config found."), OutputFormat::Json => { diff --git a/src/config.rs b/src/config.rs index d2f5395..72185e4 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2,9 +2,10 @@ use std::fs; use std::path::PathBuf; use anyhow::{Context, Result}; +use secrecy::{ExposeSecret, SecretString}; use serde::{Deserialize, Serialize}; -const ENV_VAR: &str = "POLYMARKET_PRIVATE_KEY"; +pub const ENV_VAR: &str = "POLYMARKET_PRIVATE_KEY"; const SIG_TYPE_ENV_VAR: &str = "POLYMARKET_SIGNATURE_TYPE"; pub const DEFAULT_SIGNATURE_TYPE: &str = "proxy"; @@ -13,6 +14,7 @@ pub const NO_WALLET_MSG: &str = #[derive(Serialize, Deserialize)] pub struct Config { + #[serde(default, skip_serializing_if = "String::is_empty")] pub private_key: String, pub chain_id: u64, #[serde(default = "default_signature_type")] @@ -27,6 +29,7 @@ pub enum KeySource { Flag, EnvVar, ConfigFile, + Keystore, None, } @@ -36,6 +39,7 @@ impl KeySource { Self::Flag => "--private-key flag", Self::EnvVar => "POLYMARKET_PRIVATE_KEY env var", Self::ConfigFile => "config file", + Self::Keystore => "encrypted keystore", Self::None => "not configured", } } @@ -84,7 +88,24 @@ pub fn resolve_signature_type(cli_flag: Option<&str>) -> String { DEFAULT_SIGNATURE_TYPE.to_string() } -pub fn save_wallet(key: &str, chain_id: u64, signature_type: &str) -> Result<()> { +pub fn keystore_path() -> Result { + Ok(config_dir()?.join("keystore.json")) +} + +pub fn keystore_exists() -> bool { + keystore_path().is_ok_and(|p| p.exists()) +} + +/// Returns true if old-format config has a plaintext private_key field but no keystore. +pub fn needs_migration() -> bool { + load_config().is_some_and(|c| !c.private_key.is_empty()) + && !keystore_exists() +} + +/// Encrypt a private key and save as keystore.json. +pub fn save_key_encrypted(key_hex: &SecretString, password: &SecretString) -> Result<()> { + use std::str::FromStr; + let dir = config_dir()?; fs::create_dir_all(&dir).context("Failed to create config directory")?; @@ -94,8 +115,80 @@ pub fn save_wallet(key: &str, chain_id: u64, signature_type: &str) -> Result<()> fs::set_permissions(&dir, fs::Permissions::from_mode(0o700))?; } + let signer = alloy::signers::local::LocalSigner::from_str(key_hex.expose_secret()) + .map_err(|e| anyhow::anyhow!("Invalid private key: {e}"))?; + let key_bytes = signer.credential().to_bytes(); + + let mut rng = rand::thread_rng(); + alloy::signers::local::LocalSigner::encrypt_keystore( + &dir, &mut rng, key_bytes, password.expose_secret(), Some("keystore"), + ) + .map_err(|e| anyhow::anyhow!("Failed to encrypt keystore: {e}"))?; + + // eth-keystore writes to dir/keystore — rename to keystore.json + let written = dir.join("keystore"); + let target = dir.join("keystore.json"); + if written.exists() && written != target { + fs::rename(&written, &target) + .context("Failed to rename keystore file")?; + } + + // Set restrictive permissions on keystore file + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(&target, fs::Permissions::from_mode(0o600))?; + } + + Ok(()) +} + +/// Decrypt keystore.json and return the private key as 0x-prefixed hex. +pub fn load_key_encrypted(password: &str) -> Result { + use std::fmt::Write as _; + + let path = keystore_path()?; + let signer = alloy::signers::local::LocalSigner::decrypt_keystore(&path, password) + .map_err(|e| { + let msg = e.to_string(); + if msg.contains("Mac Mismatch") { + anyhow::anyhow!("Wrong password") + } else { + anyhow::anyhow!("Failed to decrypt keystore: {e}") + } + })?; + + let bytes = signer.credential().to_bytes(); + let mut hex = String::with_capacity(2 + bytes.len() * 2); + hex.push_str("0x"); + for b in &bytes { + write!(hex, "{b:02x}").unwrap(); + } + Ok(SecretString::from(hex)) +} + +/// Migrate old plaintext config to encrypted keystore. +pub fn migrate_to_encrypted(password: &SecretString) -> Result<()> { + let config = load_config() + .ok_or_else(|| anyhow::anyhow!("No config file found to migrate"))?; + + if config.private_key.is_empty() { + anyhow::bail!("No private key found in config to migrate"); + } + + // Encrypt the key + save_key_encrypted(&SecretString::from(config.private_key), password)?; + + // Rewrite config.json without private_key + save_wallet_settings(config.chain_id, &config.signature_type)?; + + Ok(()) +} + +/// Save only non-sensitive settings to config.json (no private key). +pub fn save_wallet_settings(chain_id: u64, signature_type: &str) -> Result<()> { let config = Config { - private_key: key.to_string(), + private_key: String::new(), chain_id, signature_type: signature_type.to_string(), }; @@ -126,17 +219,17 @@ pub fn save_wallet(key: &str, chain_id: u64, signature_type: &str) -> Result<()> } /// Priority: CLI flag > env var > config file. -pub fn resolve_key(cli_flag: Option<&str>) -> (Option, KeySource) { +pub fn resolve_key(cli_flag: Option<&str>) -> (Option, KeySource) { if let Some(key) = cli_flag { - return (Some(key.to_string()), KeySource::Flag); + return (Some(SecretString::from(key.to_string())), KeySource::Flag); } if let Ok(key) = std::env::var(ENV_VAR) && !key.is_empty() { - return (Some(key), KeySource::EnvVar); + return (Some(SecretString::from(key)), KeySource::EnvVar); } if let Some(config) = load_config() { - return (Some(config.private_key), KeySource::ConfigFile); + return (Some(SecretString::from(config.private_key)), KeySource::ConfigFile); } (None, KeySource::None) } @@ -144,6 +237,7 @@ pub fn resolve_key(cli_flag: Option<&str>) -> (Option, KeySource) { #[cfg(test)] mod tests { use super::*; + use secrecy::ExposeSecret; use std::sync::Mutex; // Mutex to serialize env var tests (set_var is not thread-safe) @@ -162,7 +256,7 @@ mod tests { let _lock = ENV_LOCK.lock().unwrap(); unsafe { set(ENV_VAR, "env_key") }; let (key, source) = resolve_key(Some("flag_key")); - assert_eq!(key.unwrap(), "flag_key"); + assert_eq!(key.unwrap().expose_secret(), "flag_key"); assert!(matches!(source, KeySource::Flag)); unsafe { unset(ENV_VAR) }; } @@ -172,7 +266,7 @@ mod tests { let _lock = ENV_LOCK.lock().unwrap(); unsafe { set(ENV_VAR, "env_key_value") }; let (key, source) = resolve_key(None); - assert_eq!(key.unwrap(), "env_key_value"); + assert_eq!(key.unwrap().expose_secret(), "env_key_value"); assert!(matches!(source, KeySource::EnvVar)); unsafe { unset(ENV_VAR) }; } @@ -209,4 +303,56 @@ mod tests { let result = resolve_signature_type(None); assert!(!result.is_empty()); } + + #[test] + fn keystore_encrypt_decrypt_round_trip() { + use std::str::FromStr; + + let temp = std::env::temp_dir().join("polymarket_test_keystore"); + let _ = fs::remove_dir_all(&temp); + fs::create_dir_all(&temp).unwrap(); + + let key_hex = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"; + let password = "test_password_123"; + + let original = alloy::signers::local::LocalSigner::from_str(key_hex).unwrap(); + + let mut rng = rand::thread_rng(); + alloy::signers::local::LocalSigner::encrypt_keystore( + &temp, &mut rng, original.credential().to_bytes(), password, Some("test_ks"), + ) + .unwrap(); + + let recovered = + alloy::signers::local::LocalSigner::decrypt_keystore(temp.join("test_ks"), password) + .unwrap(); + assert_eq!(original.address(), recovered.address()); + + let _ = fs::remove_dir_all(&temp); + } + + #[test] + fn keystore_wrong_password_fails() { + use std::str::FromStr; + + let temp = std::env::temp_dir().join("polymarket_test_keystore_fail"); + let _ = fs::remove_dir_all(&temp); + fs::create_dir_all(&temp).unwrap(); + + let signer = alloy::signers::local::LocalSigner::from_str( + "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80", + ) + .unwrap(); + let mut rng = rand::thread_rng(); + alloy::signers::local::LocalSigner::encrypt_keystore( + &temp, &mut rng, signer.credential().to_bytes(), "correct", Some("test_ks2"), + ) + .unwrap(); + + let result = + alloy::signers::local::LocalSigner::decrypt_keystore(temp.join("test_ks2"), "wrong"); + assert!(result.is_err()); + + let _ = fs::remove_dir_all(&temp); + } } diff --git a/src/main.rs b/src/main.rs index 61af087..7ec098f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ mod auth; mod commands; mod config; mod output; +mod password; mod shell; use std::process::ExitCode; diff --git a/src/password.rs b/src/password.rs new file mode 100644 index 0000000..2bf4640 --- /dev/null +++ b/src/password.rs @@ -0,0 +1,51 @@ +use anyhow::{Result, bail}; +use secrecy::{ExposeSecret, SecretString}; + +const PASSWORD_ENV_VAR: &str = "POLYMARKET_PASSWORD"; + +/// Prompt for password, or read from POLYMARKET_PASSWORD env var. +pub fn prompt_password(prompt_msg: &str) -> Result { + if let Ok(pw) = std::env::var(PASSWORD_ENV_VAR) + && !pw.is_empty() + { + return Ok(SecretString::from(pw)); + } + rpassword::prompt_password(prompt_msg) + .map(SecretString::from) + .map_err(Into::into) +} + +/// Prompt for password with confirmation (for create/import). +pub fn prompt_new_password() -> Result { + let pw = prompt_password("Enter password to encrypt wallet: ")?; + if pw.expose_secret().is_empty() { + bail!("Password cannot be empty"); + } + let confirm = prompt_password("Confirm password: ")?; + if pw.expose_secret() != confirm.expose_secret() { + bail!("Passwords do not match"); + } + Ok(pw) +} + +/// Prompt for password with up to 3 retries, calling `try_fn` each time. +/// Returns the result of the first successful call to `try_fn`. +pub fn prompt_password_with_retries(try_fn: F) -> Result +where + F: Fn(&str) -> Result, +{ + for attempt in 1..=3 { + let pw = prompt_password("Enter wallet password: ")?; + match try_fn(pw.expose_secret()) { + Ok(val) => return Ok(val), + Err(e) => { + if attempt < 3 { + eprintln!("Wrong password. Try again. ({attempt}/3)"); + } else { + return Err(e); + } + } + } + } + unreachable!() +} diff --git a/tests/cli_integration.rs b/tests/cli_integration.rs index 41d3d11..317d3c0 100644 --- a/tests/cli_integration.rs +++ b/tests/cli_integration.rs @@ -7,6 +7,7 @@ fn polymarket() -> Command { let mut cmd = Command::cargo_bin("polymarket").unwrap(); cmd.env_remove("POLYMARKET_PRIVATE_KEY"); cmd.env_remove("POLYMARKET_SIGNATURE_TYPE"); + cmd.env_remove("POLYMARKET_PASSWORD"); cmd } @@ -79,7 +80,8 @@ fn wallet_help_lists_subcommands() { .and(predicate::str::contains("import")) .and(predicate::str::contains("address")) .and(predicate::str::contains("show")) - .and(predicate::str::contains("reset")), + .and(predicate::str::contains("reset")) + .and(predicate::str::contains("export")), ); }