diff --git a/Cargo.lock b/Cargo.lock index 26a3032..66adb24 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 a01bf05..a89a8e0 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 a5cf34d..13f16ff 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; @@ -23,12 +24,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 Ok(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))) } @@ -65,9 +101,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 2e4a66d..ed0ebfb 100644 --- a/src/commands/setup.rs +++ b/src/commands/setup.rs @@ -5,6 +5,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; use crate::config; @@ -82,12 +83,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!(); @@ -116,16 +136,16 @@ fn setup_wallet() -> Result
{ let signer = LocalSigner::from_str(&key) .context("Invalid private key")? .with_chain_id(Some(POLYGON)); - let hex = format!("{:#x}", signer.to_bytes()); - (signer.address(), hex) + (signer.address(), config::key_bytes_to_hex(&signer.credential().to_bytes())) } else { let signer = LocalSigner::random().with_chain_id(Some(POLYGON)); let address = signer.address(); - let hex = format!("{:#x}", signer.to_bytes()); - (address, hex) + (address, config::key_bytes_to_hex(&signer.credential().to_bytes())) }; - 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"); @@ -137,7 +157,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 d54a63b..d78acc4 100644 --- a/src/commands/wallet.rs +++ b/src/commands/wallet.rs @@ -5,6 +5,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; use crate::config; use crate::output::OutputFormat; @@ -47,6 +48,8 @@ pub enum WalletCommand { #[arg(long)] force: bool, }, + /// Export the private key (decrypts keystore, prints to stdout) + Export, } pub fn execute( @@ -64,17 +67,17 @@ pub fn execute( force, signature_type, } => cmd_import(&key, output, force, &signature_type), - WalletCommand::Address => cmd_address(output, private_key_flag), - WalletCommand::Show => cmd_show(output, private_key_flag), - WalletCommand::Reset { force } => cmd_reset(output, force), + 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(()) @@ -85,9 +88,11 @@ fn cmd_create(output: OutputFormat, force: bool, signature_type: &str) -> Result let signer = LocalSigner::random().with_chain_id(Some(POLYGON)); let address = signer.address(); - let key_hex = format!("{:#x}", signer.to_bytes()); + let key_hex = config::key_bytes_to_hex(&signer.credential().to_bytes()); - 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); @@ -112,8 +117,8 @@ fn cmd_create(output: OutputFormat, force: bool, signature_type: &str) -> Result 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(()) @@ -126,9 +131,11 @@ fn cmd_import(key: &str, output: OutputFormat, force: bool, signature_type: &str .context("Invalid private key")? .with_chain_id(Some(POLYGON)); let address = signer.address(); - let key_hex = format!("{:#x}", signer.to_bytes()); + let key_hex = config::key_bytes_to_hex(&signer.credential().to_bytes()); - 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); @@ -157,11 +164,10 @@ fn cmd_import(key: &str, output: OutputFormat, force: bool, signature_type: &str Ok(()) } -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))?; +fn cmd_address(output: &OutputFormat, private_key_flag: Option<&str>) -> Result<()> { + 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 { @@ -175,9 +181,22 @@ fn cmd_address(output: OutputFormat, private_key_flag: Option<&str>) -> Result<( Ok(()) } -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()); +fn cmd_show(output: &OutputFormat, private_key_flag: Option<&str>) -> Result<()> { + 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() @@ -217,8 +236,28 @@ fn cmd_show(output: OutputFormat, private_key_flag: Option<&str>) -> Result<()> Ok(()) } -fn cmd_reset(output: OutputFormat, force: bool) -> Result<()> { - if !config::config_exists() { +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() && !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 60c0179..8544500 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(crate) const DEFAULT_SIGNATURE_TYPE: &str = "proxy"; @@ -12,7 +13,8 @@ pub(crate) const NO_WALLET_MSG: &str = "No wallet configured. Run `polymarket wallet create` or `polymarket wallet import `"; #[derive(Serialize, Deserialize)] -pub(crate) struct Config { +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(crate) 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", } } @@ -94,7 +98,24 @@ pub fn resolve_signature_type(cli_flag: Option<&str>) -> Result { Ok(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().ok().flatten().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")?; @@ -104,8 +125,84 @@ 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(()) +} + +/// Convert raw key bytes to a 0x-prefixed hex SecretString. +pub(crate) fn key_bytes_to_hex(bytes: &[u8]) -> SecretString { + use std::fmt::Write as _; + + let mut hex = String::with_capacity(2 + bytes.len() * 2); + hex.push_str("0x"); + for b in bytes { + write!(hex, "{b:02x}").unwrap(); + } + SecretString::from(hex) +} + +/// Decrypt keystore.json and return the private key as 0x-prefixed hex. +pub fn load_key_encrypted(password: &str) -> Result { + 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}") + } + })?; + + Ok(key_bytes_to_hex(&signer.credential().to_bytes())) +} + +/// 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(), }; @@ -136,24 +233,25 @@ 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>) -> Result<(Option, KeySource)> { +pub fn resolve_key(cli_flag: Option<&str>) -> (Option, KeySource) { if let Some(key) = cli_flag { - return Ok((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 Ok((Some(key), KeySource::EnvVar)); + return (Some(SecretString::from(key)), KeySource::EnvVar); } - if let Some(config) = load_config()? { - return Ok((Some(config.private_key), KeySource::ConfigFile)); + if let Ok(Some(config)) = load_config() { + return (Some(SecretString::from(config.private_key)), KeySource::ConfigFile); } - Ok((None, KeySource::None)) + (None, KeySource::None) } #[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) @@ -171,8 +269,8 @@ mod tests { fn resolve_key_flag_overrides_env() { let _lock = ENV_LOCK.lock().unwrap(); unsafe { set(ENV_VAR, "env_key") }; - let (key, source) = resolve_key(Some("flag_key")).unwrap(); - assert_eq!(key.unwrap(), "flag_key"); + let (key, source) = resolve_key(Some("flag_key")); + assert_eq!(key.unwrap().expose_secret(), "flag_key"); assert!(matches!(source, KeySource::Flag)); unsafe { unset(ENV_VAR) }; } @@ -181,8 +279,8 @@ mod tests { fn resolve_key_env_var_returns_env_value() { let _lock = ENV_LOCK.lock().unwrap(); unsafe { set(ENV_VAR, "env_key_value") }; - let (key, source) = resolve_key(None).unwrap(); - assert_eq!(key.unwrap(), "env_key_value"); + let (key, source) = resolve_key(None); + assert_eq!(key.unwrap().expose_secret(), "env_key_value"); assert!(matches!(source, KeySource::EnvVar)); unsafe { unset(ENV_VAR) }; } @@ -191,7 +289,7 @@ mod tests { fn resolve_key_skips_empty_env_var() { let _lock = ENV_LOCK.lock().unwrap(); unsafe { set(ENV_VAR, "") }; - let (_, source) = resolve_key(None).unwrap(); + let (_, source) = resolve_key(None); assert!(!matches!(source, KeySource::EnvVar)); unsafe { unset(ENV_VAR) }; } @@ -222,4 +320,56 @@ mod tests { let result = resolve_signature_type(None).unwrap(); 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 2abb55f..30fd0ef 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")), ); }