From 70dee9daf937197cd9f5fafeb2108269c0c22c3e Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Tue, 10 Mar 2026 16:57:47 -0700 Subject: [PATCH 1/2] feat: encrypt private keys at rest using keystore Wallet create and import now encrypt private keys using the standard Ethereum keystore format (AES-128-CTR + scrypt), the same format used by MetaMask, Geth, and Foundry. Users are prompted for a password on key creation/import and when any command needs the key. The POLYMARKET_PASSWORD env var supports non-interactive use (CI/scripts). Use --no-password to opt into plaintext storage (not recommended). Adds wallet export subcommand to decrypt and display the key. Fixes #18 Co-Authored-By: Claude Opus 4.6 --- Cargo.lock | 131 ++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 4 +- src/commands/wallet.rs | 114 +++++++++++++++++++++++++++++++---- src/config.rs | 79 ++++++++++++++++++++++--- 4 files changed, 306 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 26a3032..bf9b37e 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,6 +3047,8 @@ dependencies = [ "dirs", "polymarket-client-sdk", "predicates", + "rand 0.8.5", + "rpassword", "rust_decimal", "rust_decimal_macros", "rustyline", @@ -3014,7 +3087,7 @@ dependencies = [ "sha2", "strum_macros", "url", - "uuid", + "uuid 1.21.0", ] [[package]] @@ -3540,7 +3613,7 @@ dependencies = [ "rkyv_derive", "seahash", "tinyvec", - "uuid", + "uuid 1.21.0", ] [[package]] @@ -3564,6 +3637,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 +3883,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 +3940,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 +4796,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..32e67b6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,9 @@ 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", "signer-keystore", "signers"] } +rpassword = "7" +rand = "0.8" clap = { version = "4", features = ["derive"] } tokio = { version = "1", features = ["rt-multi-thread", "macros"] } serde_json = "1" diff --git a/src/commands/wallet.rs b/src/commands/wallet.rs index d54a63b..49120bc 100644 --- a/src/commands/wallet.rs +++ b/src/commands/wallet.rs @@ -1,5 +1,6 @@ use std::str::FromStr; +use alloy::signers::local::PrivateKeySigner; use anyhow::{Context, Result, bail}; use clap::{Args, Subcommand}; use polymarket_client_sdk::auth::LocalSigner; @@ -25,6 +26,9 @@ pub enum WalletCommand { /// Signature type: eoa, proxy (default), or gnosis-safe #[arg(long, default_value = "proxy")] signature_type: String, + /// Store the key as plaintext (not recommended) + #[arg(long)] + no_password: bool, }, /// Import an existing private key Import { @@ -36,11 +40,16 @@ pub enum WalletCommand { /// Signature type: eoa, proxy (default), or gnosis-safe #[arg(long, default_value = "proxy")] signature_type: String, + /// Store the key as plaintext (not recommended) + #[arg(long)] + no_password: bool, }, /// Show the address of the configured wallet Address, /// Show wallet info (address, config path, key source) Show, + /// Export the private key (decrypts if encrypted) + Export, /// Delete all config and keys (fresh install) Reset { /// Skip confirmation prompt @@ -58,18 +67,45 @@ pub fn execute( WalletCommand::Create { force, signature_type, - } => cmd_create(output, force, &signature_type), + no_password, + } => cmd_create(output, force, &signature_type, no_password), WalletCommand::Import { key, force, signature_type, - } => cmd_import(&key, output, force, &signature_type), + no_password, + } => cmd_import(&key, output, force, &signature_type, no_password), WalletCommand::Address => cmd_address(output, private_key_flag), WalletCommand::Show => cmd_show(output, private_key_flag), + WalletCommand::Export => cmd_export(output, private_key_flag), WalletCommand::Reset { force } => cmd_reset(output, force), } } +fn save_encrypted(signer: &PrivateKeySigner, password: &str, signature_type: &str) -> Result<()> { + let ks_dir = config::keystore_dir()?; + std::fs::create_dir_all(&ks_dir).context("Failed to create keystore directory")?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&ks_dir, std::fs::Permissions::from_mode(0o700))?; + } + + // Remove any existing keystore files + if let Ok(entries) = std::fs::read_dir(&ks_dir) { + for entry in entries.flatten() { + let _ = std::fs::remove_file(entry.path()); + } + } + + let mut rng = rand::thread_rng(); + PrivateKeySigner::encrypt_keystore(&ks_dir, &mut rng, signer.to_bytes(), password, None) + .context("Failed to encrypt keystore")?; + + config::save_wallet_encrypted(POLYGON, signature_type) +} + fn guard_overwrite(force: bool) -> Result<()> { if !force && config::config_exists() { bail!( @@ -80,14 +116,26 @@ fn guard_overwrite(force: bool) -> Result<()> { Ok(()) } -fn cmd_create(output: OutputFormat, force: bool, signature_type: &str) -> Result<()> { +fn cmd_create( + output: OutputFormat, + force: bool, + signature_type: &str, + no_password: bool, +) -> Result<()> { guard_overwrite(force)?; let signer = LocalSigner::random().with_chain_id(Some(POLYGON)); let address = signer.address(); - let key_hex = format!("{:#x}", signer.to_bytes()); + let encrypted = !no_password; + + if encrypted { + let password = config::read_new_password()?; + save_encrypted(&signer, &password, signature_type)?; + } else { + let key_hex = format!("{:#x}", signer.to_bytes()); + config::save_wallet(&key_hex, POLYGON, signature_type)?; + } - config::save_wallet(&key_hex, POLYGON, signature_type)?; let config_path = config::config_path()?; let proxy_addr = derive_proxy_wallet(address, POLYGON); @@ -100,6 +148,7 @@ fn cmd_create(output: OutputFormat, force: bool, signature_type: &str) -> Result "proxy_address": proxy_addr.map(|a| a.to_string()), "signature_type": signature_type, "config_path": config_path.display().to_string(), + "encrypted": encrypted, }) ); } @@ -110,25 +159,42 @@ fn cmd_create(output: OutputFormat, force: bool, signature_type: &str) -> Result println!("Proxy wallet: {proxy}"); } println!("Signature type: {signature_type}"); + println!("Encrypted: {encrypted}"); 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."); + if !encrypted { + println!(); + println!( + "WARNING: Key stored as plaintext. Use without --no-password for encryption." + ); + } } } Ok(()) } -fn cmd_import(key: &str, output: OutputFormat, force: bool, signature_type: &str) -> Result<()> { +fn cmd_import( + key: &str, + output: OutputFormat, + force: bool, + signature_type: &str, + no_password: bool, +) -> Result<()> { guard_overwrite(force)?; let signer = LocalSigner::from_str(key) .context("Invalid private key")? .with_chain_id(Some(POLYGON)); let address = signer.address(); - let key_hex = format!("{:#x}", signer.to_bytes()); + let encrypted = !no_password; + + if encrypted { + let password = config::read_new_password()?; + save_encrypted(&signer, &password, signature_type)?; + } else { + let key_hex = format!("{:#x}", signer.to_bytes()); + config::save_wallet(&key_hex, POLYGON, signature_type)?; + } - config::save_wallet(&key_hex, POLYGON, signature_type)?; let config_path = config::config_path()?; let proxy_addr = derive_proxy_wallet(address, POLYGON); @@ -141,6 +207,7 @@ fn cmd_import(key: &str, output: OutputFormat, force: bool, signature_type: &str "proxy_address": proxy_addr.map(|a| a.to_string()), "signature_type": signature_type, "config_path": config_path.display().to_string(), + "encrypted": encrypted, }) ); } @@ -151,6 +218,7 @@ fn cmd_import(key: &str, output: OutputFormat, force: bool, signature_type: &str println!("Proxy wallet: {proxy}"); } println!("Signature type: {signature_type}"); + println!("Encrypted: {encrypted}"); println!("Config: {}", config_path.display()); } } @@ -217,6 +285,30 @@ fn cmd_show(output: OutputFormat, private_key_flag: Option<&str>) -> Result<()> Ok(()) } +fn cmd_export(output: OutputFormat, private_key_flag: Option<&str>) -> Result<()> { + let (key, source) = config::resolve_key(private_key_flag)?; + let key = key.ok_or_else(|| anyhow::anyhow!("{}", config::NO_WALLET_MSG))?; + + match output { + OutputFormat::Json => { + println!( + "{}", + serde_json::json!({ + "private_key": key, + "source": source.label(), + }) + ); + } + OutputFormat::Table => { + println!("Private key: {key}"); + println!("Source: {}", source.label()); + println!(); + println!("WARNING: Do not share this key. Anyone with it can access your funds."); + } + } + Ok(()) +} + fn cmd_reset(output: OutputFormat, force: bool) -> Result<()> { if !config::config_exists() { match output { diff --git a/src/config.rs b/src/config.rs index 60c0179..a01fe05 100644 --- a/src/config.rs +++ b/src/config.rs @@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize}; const ENV_VAR: &str = "POLYMARKET_PRIVATE_KEY"; const SIG_TYPE_ENV_VAR: &str = "POLYMARKET_SIGNATURE_TYPE"; +const PASSWORD_ENV_VAR: &str = "POLYMARKET_PASSWORD"; pub(crate) const DEFAULT_SIGNATURE_TYPE: &str = "proxy"; pub(crate) const NO_WALLET_MSG: &str = @@ -13,10 +14,13 @@ pub(crate) const NO_WALLET_MSG: &str = #[derive(Serialize, Deserialize)] pub(crate) struct Config { + #[serde(default)] pub private_key: String, pub chain_id: u64, #[serde(default = "default_signature_type")] pub signature_type: String, + #[serde(default)] + pub encrypted: bool, } fn default_signature_type() -> String { @@ -50,6 +54,29 @@ pub fn config_path() -> Result { Ok(config_dir()?.join("config.json")) } +pub fn keystore_dir() -> Result { + Ok(config_dir()?.join("keystore")) +} + +/// Prompt for a password, or read from `POLYMARKET_PASSWORD` env var. +pub fn read_password(prompt: &str) -> Result { + if let Ok(pw) = std::env::var(PASSWORD_ENV_VAR) + && !pw.is_empty() + { + return Ok(pw); + } + rpassword::prompt_password(prompt).context("Failed to read password") +} + +/// Prompt for a new password with confirmation. +pub fn read_new_password() -> Result { + let pw = read_password("Enter password to encrypt your key: ")?; + anyhow::ensure!(!pw.is_empty(), "Password cannot be empty"); + let confirm = read_password("Confirm password: ")?; + anyhow::ensure!(pw == confirm, "Passwords do not match"); + Ok(pw) +} + pub fn config_exists() -> bool { config_path().is_ok_and(|p| p.exists()) } @@ -94,7 +121,7 @@ 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<()> { +fn write_config_file(config: &Config) -> Result<()> { let dir = config_dir()?; fs::create_dir_all(&dir).context("Failed to create config directory")?; @@ -104,12 +131,7 @@ pub fn save_wallet(key: &str, chain_id: u64, signature_type: &str) -> Result<()> fs::set_permissions(&dir, fs::Permissions::from_mode(0o700))?; } - let config = Config { - private_key: key.to_string(), - chain_id, - signature_type: signature_type.to_string(), - }; - let json = serde_json::to_string_pretty(&config)?; + let json = serde_json::to_string_pretty(config)?; let path = config_path()?; #[cfg(unix)] @@ -135,7 +157,29 @@ pub fn save_wallet(key: &str, chain_id: u64, signature_type: &str) -> Result<()> Ok(()) } +/// Save wallet as plaintext (legacy, used only when `--no-password` is specified). +pub fn save_wallet(key: &str, chain_id: u64, signature_type: &str) -> Result<()> { + write_config_file(&Config { + private_key: key.to_string(), + chain_id, + signature_type: signature_type.to_string(), + encrypted: false, + }) +} + +/// Save wallet as encrypted keystore. The keystore file must already be +/// written to `keystore_dir()` by the caller. +pub fn save_wallet_encrypted(chain_id: u64, signature_type: &str) -> Result<()> { + write_config_file(&Config { + private_key: String::new(), + chain_id, + signature_type: signature_type.to_string(), + encrypted: true, + }) +} + /// Priority: CLI flag > env var > config file. +/// If the config file uses encrypted keystore, prompts for password. pub fn resolve_key(cli_flag: Option<&str>) -> Result<(Option, KeySource)> { if let Some(key) = cli_flag { return Ok((Some(key.to_string()), KeySource::Flag)); @@ -146,11 +190,32 @@ pub fn resolve_key(cli_flag: Option<&str>) -> Result<(Option, KeySource) return Ok((Some(key), KeySource::EnvVar)); } if let Some(config) = load_config()? { + if config.encrypted { + let key = decrypt_keystore()?; + return Ok((Some(key), KeySource::ConfigFile)); + } return Ok((Some(config.private_key), KeySource::ConfigFile)); } Ok((None, KeySource::None)) } +/// Decrypt the keystore file and return the private key hex. +fn decrypt_keystore() -> Result { + use polymarket_client_sdk::auth::LocalSigner; + + let ks_dir = keystore_dir()?; + let entry = fs::read_dir(&ks_dir) + .context("Failed to read keystore directory")? + .flatten() + .find(|e| e.path().is_file()) + .context("No keystore file found. Config says encrypted but keystore is missing.")?; + + let password = read_password("Enter password to unlock wallet: ")?; + let signer = LocalSigner::decrypt_keystore(entry.path(), password) + .context("Failed to decrypt keystore. Wrong password?")?; + Ok(format!("{:#x}", signer.to_bytes())) +} + #[cfg(test)] mod tests { use super::*; From d491d80ffbffc74e9ce781e57e3cdc3612cea393 Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Tue, 10 Mar 2026 17:28:04 -0700 Subject: [PATCH 2/2] fix(wallet): atomic keystore write and cleanup stale files in plaintext path Co-Authored-By: Claude Opus 4.6 --- src/commands/wallet.rs | 49 +++++++++++++++++++++++++++++++++++++----- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/src/commands/wallet.rs b/src/commands/wallet.rs index 49120bc..8411e17 100644 --- a/src/commands/wallet.rs +++ b/src/commands/wallet.rs @@ -92,20 +92,57 @@ fn save_encrypted(signer: &PrivateKeySigner, password: &str, signature_type: &st std::fs::set_permissions(&ks_dir, std::fs::Permissions::from_mode(0o700))?; } - // Remove any existing keystore files + // Write new keystore to a temp directory first so a failure doesn't destroy the old one + let tmp_dir = ks_dir.join(".tmp"); + if tmp_dir.exists() { + std::fs::remove_dir_all(&tmp_dir).context("Failed to clean temp keystore dir")?; + } + std::fs::create_dir_all(&tmp_dir).context("Failed to create temp keystore dir")?; + + let mut rng = rand::thread_rng(); + PrivateKeySigner::encrypt_keystore(&tmp_dir, &mut rng, signer.to_bytes(), password, None) + .context("Failed to encrypt keystore")?; + + // Find the newly written keystore file + let new_file = std::fs::read_dir(&tmp_dir)? + .flatten() + .next() + .ok_or_else(|| anyhow::anyhow!("encrypt_keystore produced no file"))? + .path(); + let file_name = new_file + .file_name() + .ok_or_else(|| anyhow::anyhow!("keystore file has no name"))? + .to_owned(); + + // Now safe to remove old keystore files if let Ok(entries) = std::fs::read_dir(&ks_dir) { for entry in entries.flatten() { - let _ = std::fs::remove_file(entry.path()); + let path = entry.path(); + if path != tmp_dir { + let _ = std::fs::remove_file(&path); + } } } - let mut rng = rand::thread_rng(); - PrivateKeySigner::encrypt_keystore(&ks_dir, &mut rng, signer.to_bytes(), password, None) - .context("Failed to encrypt keystore")?; + // Move new keystore into place and clean up temp dir + std::fs::rename(&new_file, ks_dir.join(&file_name)) + .context("Failed to move keystore into place")?; + let _ = std::fs::remove_dir_all(&tmp_dir); config::save_wallet_encrypted(POLYGON, signature_type) } +/// Remove any leftover keystore files (e.g. when switching from encrypted to plaintext). +fn cleanup_keystore_files() { + if let Ok(ks_dir) = config::keystore_dir() { + if let Ok(entries) = std::fs::read_dir(&ks_dir) { + for entry in entries.flatten() { + let _ = std::fs::remove_file(entry.path()); + } + } + } +} + fn guard_overwrite(force: bool) -> Result<()> { if !force && config::config_exists() { bail!( @@ -132,6 +169,7 @@ fn cmd_create( let password = config::read_new_password()?; save_encrypted(&signer, &password, signature_type)?; } else { + cleanup_keystore_files(); let key_hex = format!("{:#x}", signer.to_bytes()); config::save_wallet(&key_hex, POLYGON, signature_type)?; } @@ -191,6 +229,7 @@ fn cmd_import( let password = config::read_new_password()?; save_encrypted(&signer, &password, signature_type)?; } else { + cleanup_keystore_files(); let key_hex = format!("{:#x}", signer.to_bytes()); config::save_wallet(&key_hex, POLYGON, signature_type)?; }