1use crate::{
2 utils,
3 wallet_signer::{PendingSigner, WalletSigner},
4};
5use alloy_primitives::map::AddressHashMap;
6use alloy_signer::Signer;
7use clap::Parser;
8use derive_builder::Builder;
9use eyre::Result;
10use foundry_config::Config;
11use serde::Serialize;
12use std::path::PathBuf;
13
14#[derive(Debug, Default)]
16pub struct MultiWallet {
17 pending_signers: Vec<PendingSigner>,
20 signers: AddressHashMap<WalletSigner>,
22}
23
24impl MultiWallet {
25 pub fn new(pending_signers: Vec<PendingSigner>, signers: Vec<WalletSigner>) -> Self {
26 let signers = signers.into_iter().map(|signer| (signer.address(), signer)).collect();
27 Self { pending_signers, signers }
28 }
29
30 fn maybe_unlock_pending(&mut self) -> Result<()> {
31 for pending in self.pending_signers.drain(..) {
32 let signer = pending.unlock()?;
33 self.signers.insert(signer.address(), signer);
34 }
35 Ok(())
36 }
37
38 pub fn signers(&mut self) -> Result<&AddressHashMap<WalletSigner>> {
39 self.maybe_unlock_pending()?;
40 Ok(&self.signers)
41 }
42
43 pub fn into_signers(mut self) -> Result<AddressHashMap<WalletSigner>> {
44 self.maybe_unlock_pending()?;
45 Ok(self.signers)
46 }
47
48 pub fn add_signer(&mut self, signer: WalletSigner) {
49 self.signers.insert(signer.address(), signer);
50 }
51}
52
53macro_rules! create_hw_wallets {
57 ($self:ident, $create_signer:expr, $signers:ident) => {
58 let mut $signers = vec![];
59
60 if let Some(hd_paths) = &$self.hd_paths {
61 for path in hd_paths {
62 let hw = $create_signer(Some(path), 0).await?;
63 $signers.push(hw);
64 }
65 }
66
67 if let Some(mnemonic_indexes) = &$self.mnemonic_indexes {
68 for index in mnemonic_indexes {
69 let hw = $create_signer(None, *index).await?;
70 $signers.push(hw);
71 }
72 }
73
74 if $signers.is_empty() {
75 let hw = $create_signer(None, 0).await?;
76 $signers.push(hw);
77 }
78 };
79}
80
81#[derive(Builder, Clone, Debug, Default, Serialize, Parser)]
90#[command(next_help_heading = "Wallet options", about = None, long_about = None)]
91pub struct MultiWalletOpts {
92 #[arg(
96 long,
97 short,
98 help_heading = "Wallet options - raw",
99 default_value = "0",
100 value_name = "NUM"
101 )]
102 pub interactives: u32,
103
104 #[arg(long, help_heading = "Wallet options - raw", value_name = "RAW_PRIVATE_KEYS")]
106 #[builder(default = "None")]
107 pub private_keys: Option<Vec<String>>,
108
109 #[arg(
111 long,
112 help_heading = "Wallet options - raw",
113 conflicts_with = "private_keys",
114 value_name = "RAW_PRIVATE_KEY"
115 )]
116 #[builder(default = "None")]
117 pub private_key: Option<String>,
118
119 #[arg(long, alias = "mnemonic-paths", help_heading = "Wallet options - raw")]
121 #[builder(default = "None")]
122 pub mnemonics: Option<Vec<String>>,
123
124 #[arg(long, help_heading = "Wallet options - raw", value_name = "PASSPHRASE")]
126 #[builder(default = "None")]
127 pub mnemonic_passphrases: Option<Vec<String>>,
128
129 #[arg(
133 long = "mnemonic-derivation-paths",
134 alias = "hd-paths",
135 help_heading = "Wallet options - raw",
136 value_name = "PATH"
137 )]
138 #[builder(default = "None")]
139 pub hd_paths: Option<Vec<String>>,
140
141 #[arg(
145 long,
146 conflicts_with = "hd_paths",
147 help_heading = "Wallet options - raw",
148 default_value = "0",
149 value_name = "INDEXES"
150 )]
151 pub mnemonic_indexes: Option<Vec<u32>>,
152
153 #[arg(
155 long = "keystore",
156 visible_alias = "keystores",
157 help_heading = "Wallet options - keystore",
158 value_name = "PATHS",
159 env = "ETH_KEYSTORE"
160 )]
161 #[builder(default = "None")]
162 pub keystore_paths: Option<Vec<String>>,
163
164 #[arg(
166 long = "account",
167 visible_alias = "accounts",
168 help_heading = "Wallet options - keystore",
169 value_name = "ACCOUNT_NAMES",
170 env = "ETH_KEYSTORE_ACCOUNT",
171 conflicts_with = "keystore_paths"
172 )]
173 #[builder(default = "None")]
174 pub keystore_account_names: Option<Vec<String>>,
175
176 #[arg(
180 long = "password",
181 help_heading = "Wallet options - keystore",
182 requires = "keystore_paths",
183 value_name = "PASSWORDS"
184 )]
185 #[builder(default = "None")]
186 pub keystore_passwords: Option<Vec<String>>,
187
188 #[arg(
192 long = "password-file",
193 help_heading = "Wallet options - keystore",
194 requires = "keystore_paths",
195 value_name = "PATHS",
196 env = "ETH_PASSWORD"
197 )]
198 #[builder(default = "None")]
199 pub keystore_password_files: Option<Vec<String>>,
200
201 #[arg(long, short, help_heading = "Wallet options - hardware wallet")]
203 pub ledger: bool,
204
205 #[arg(long, short, help_heading = "Wallet options - hardware wallet")]
207 pub trezor: bool,
208
209 #[arg(long, help_heading = "Wallet options - remote", hide = !cfg!(feature = "aws-kms"))]
214 pub aws: bool,
215
216 #[arg(long, help_heading = "Wallet options - remote", hide = !cfg!(feature = "gcp-kms"))]
223 pub gcp: bool,
224}
225
226impl MultiWalletOpts {
227 pub async fn get_multi_wallet(&self) -> Result<MultiWallet> {
229 let mut pending = Vec::new();
230 let mut signers: Vec<WalletSigner> = Vec::new();
231
232 if let Some(ledgers) = self.ledgers().await? {
233 signers.extend(ledgers);
234 }
235 if let Some(trezors) = self.trezors().await? {
236 signers.extend(trezors);
237 }
238 if let Some(aws_signers) = self.aws_signers().await? {
239 signers.extend(aws_signers);
240 }
241 if let Some(gcp_signer) = self.gcp_signers().await? {
242 signers.extend(gcp_signer);
243 }
244 if let Some((pending_keystores, unlocked)) = self.keystores()? {
245 pending.extend(pending_keystores);
246 signers.extend(unlocked);
247 }
248 if let Some(pks) = self.private_keys()? {
249 signers.extend(pks);
250 }
251 if let Some(mnemonics) = self.mnemonics()? {
252 signers.extend(mnemonics);
253 }
254 if self.interactives > 0 {
255 pending.extend(std::iter::repeat_n(
256 PendingSigner::Interactive,
257 self.interactives as usize,
258 ));
259 }
260
261 Ok(MultiWallet::new(pending, signers))
262 }
263
264 pub fn private_keys(&self) -> Result<Option<Vec<WalletSigner>>> {
265 let mut pks = vec![];
266 if let Some(private_key) = &self.private_key {
267 pks.push(private_key);
268 }
269 if let Some(private_keys) = &self.private_keys {
270 for pk in private_keys {
271 pks.push(pk);
272 }
273 }
274 if !pks.is_empty() {
275 let wallets = pks
276 .into_iter()
277 .map(|pk| utils::create_private_key_signer(pk))
278 .collect::<Result<Vec<_>>>()?;
279 Ok(Some(wallets))
280 } else {
281 Ok(None)
282 }
283 }
284
285 fn keystore_paths(&self) -> Result<Option<Vec<PathBuf>>> {
286 if let Some(keystore_paths) = &self.keystore_paths {
287 return Ok(Some(keystore_paths.iter().map(PathBuf::from).collect()));
288 }
289 if let Some(keystore_account_names) = &self.keystore_account_names {
290 let default_keystore_dir = Config::foundry_keystores_dir()
291 .ok_or_else(|| eyre::eyre!("Could not find the default keystore directory."))?;
292 return Ok(Some(
293 keystore_account_names
294 .iter()
295 .map(|keystore_name| default_keystore_dir.join(keystore_name))
296 .collect(),
297 ));
298 }
299 Ok(None)
300 }
301
302 pub fn keystores(&self) -> Result<Option<(Vec<PendingSigner>, Vec<WalletSigner>)>> {
306 if let Some(keystore_paths) = self.keystore_paths()? {
307 let mut pending = Vec::new();
308 let mut signers = Vec::new();
309
310 let mut passwords_iter =
311 self.keystore_passwords.iter().flat_map(|passwords| passwords.iter());
312
313 let mut password_files_iter = self
314 .keystore_password_files
315 .iter()
316 .flat_map(|password_files| password_files.iter());
317
318 for path in &keystore_paths {
319 let (maybe_signer, maybe_pending) = utils::create_keystore_signer(
320 path,
321 passwords_iter.next().map(|password| password.as_str()),
322 password_files_iter.next().map(|password_file| password_file.as_str()),
323 )?;
324 if let Some(pending_signer) = maybe_pending {
325 pending.push(pending_signer);
326 } else if let Some(signer) = maybe_signer {
327 signers.push(signer);
328 }
329 }
330 return Ok(Some((pending, signers)));
331 }
332 Ok(None)
333 }
334
335 pub fn mnemonics(&self) -> Result<Option<Vec<WalletSigner>>> {
336 if let Some(ref mnemonics) = self.mnemonics {
337 let mut wallets = vec![];
338
339 let mut hd_paths_iter =
340 self.hd_paths.iter().flat_map(|paths| paths.iter().map(String::as_str));
341
342 let mut passphrases_iter = self
343 .mnemonic_passphrases
344 .iter()
345 .flat_map(|passphrases| passphrases.iter().map(String::as_str));
346
347 let mut indexes_iter =
348 self.mnemonic_indexes.iter().flat_map(|indexes| indexes.iter().copied());
349
350 for mnemonic in mnemonics {
351 let wallet = utils::create_mnemonic_signer(
352 mnemonic,
353 passphrases_iter.next(),
354 hd_paths_iter.next(),
355 indexes_iter.next().unwrap_or(0),
356 )?;
357 wallets.push(wallet);
358 }
359 return Ok(Some(wallets));
360 }
361 Ok(None)
362 }
363
364 pub async fn ledgers(&self) -> Result<Option<Vec<WalletSigner>>> {
365 if self.ledger {
366 let mut args = self.clone();
367
368 if let Some(paths) = &args.hd_paths {
369 if paths.len() > 1 {
370 eyre::bail!("Ledger only supports one signer.");
371 }
372 args.mnemonic_indexes = None;
373 }
374
375 create_hw_wallets!(args, utils::create_ledger_signer, wallets);
376 return Ok(Some(wallets));
377 }
378 Ok(None)
379 }
380
381 pub async fn trezors(&self) -> Result<Option<Vec<WalletSigner>>> {
382 if self.trezor {
383 create_hw_wallets!(self, utils::create_trezor_signer, wallets);
384 return Ok(Some(wallets));
385 }
386 Ok(None)
387 }
388
389 pub async fn aws_signers(&self) -> Result<Option<Vec<WalletSigner>>> {
390 #[cfg(feature = "aws-kms")]
391 if self.aws {
392 let mut wallets = vec![];
393 let aws_keys = std::env::var("AWS_KMS_KEY_IDS")
394 .or(std::env::var("AWS_KMS_KEY_ID"))?
395 .split(',')
396 .map(|k| k.to_string())
397 .collect::<Vec<_>>();
398
399 for key in aws_keys {
400 let aws_signer = WalletSigner::from_aws(key).await?;
401 wallets.push(aws_signer)
402 }
403
404 return Ok(Some(wallets));
405 }
406
407 Ok(None)
408 }
409
410 pub async fn gcp_signers(&self) -> Result<Option<Vec<WalletSigner>>> {
421 #[cfg(feature = "gcp-kms")]
422 if self.gcp {
423 let mut wallets = vec![];
424
425 let project_id = std::env::var("GCP_PROJECT_ID")?;
426 let location = std::env::var("GCP_LOCATION")?;
427 let key_ring = std::env::var("GCP_KEY_RING")?;
428 let key_names = std::env::var("GCP_KEY_NAME")?;
429 let key_version = std::env::var("GCP_KEY_VERSION")?;
430
431 let gcp_signer = WalletSigner::from_gcp(
432 project_id,
433 location,
434 key_ring,
435 key_names,
436 key_version.parse()?,
437 )
438 .await?;
439 wallets.push(gcp_signer);
440
441 return Ok(Some(wallets));
442 }
443
444 Ok(None)
445 }
446}
447
448#[cfg(test)]
449mod tests {
450 use super::*;
451 use alloy_primitives::address;
452 use std::path::Path;
453
454 #[test]
455 fn parse_keystore_args() {
456 let args: MultiWalletOpts =
457 MultiWalletOpts::parse_from(["foundry-cli", "--keystores", "my/keystore/path"]);
458 assert_eq!(args.keystore_paths, Some(vec!["my/keystore/path".to_string()]));
459
460 unsafe {
461 std::env::set_var("ETH_KEYSTORE", "MY_KEYSTORE");
462 }
463 let args: MultiWalletOpts = MultiWalletOpts::parse_from(["foundry-cli"]);
464 assert_eq!(args.keystore_paths, Some(vec!["MY_KEYSTORE".to_string()]));
465
466 unsafe {
467 std::env::remove_var("ETH_KEYSTORE");
468 }
469 }
470
471 #[test]
472 fn parse_keystore_password_file() {
473 let keystore =
474 Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/../cast/tests/fixtures/keystore"));
475 let keystore_file = keystore
476 .join("UTC--2022-12-20T10-30-43.591916000Z--ec554aeafe75601aaab43bd4621a22284db566c2");
477
478 let keystore_password_file = keystore.join("password-ec554").into_os_string();
479
480 let args: MultiWalletOpts = MultiWalletOpts::parse_from([
481 "foundry-cli",
482 "--keystores",
483 keystore_file.to_str().unwrap(),
484 "--password-file",
485 keystore_password_file.to_str().unwrap(),
486 ]);
487 assert_eq!(
488 args.keystore_password_files,
489 Some(vec![keystore_password_file.to_str().unwrap().to_string()])
490 );
491
492 let (_, unlocked) = args.keystores().unwrap().unwrap();
493 assert_eq!(unlocked.len(), 1);
494 assert_eq!(unlocked[0].address(), address!("0xec554aeafe75601aaab43bd4621a22284db566c2"));
495 }
496
497 #[test]
499 fn should_not_require_the_mnemonics_flag_with_mnemonic_indexes() {
500 let wallet_options = vec![
501 ("ledger", "--mnemonic-indexes", 1),
502 ("trezor", "--mnemonic-indexes", 2),
503 ("aws", "--mnemonic-indexes", 10),
504 ];
505
506 for test_case in wallet_options {
507 let args: MultiWalletOpts = MultiWalletOpts::parse_from([
508 "foundry-cli",
509 &format!("--{}", test_case.0),
510 test_case.1,
511 &test_case.2.to_string(),
512 ]);
513
514 match test_case.0 {
515 "ledger" => assert!(args.ledger),
516 "trezor" => assert!(args.trezor),
517 "aws" => assert!(args.aws),
518 _ => panic!("Should have matched one of the previous wallet options"),
519 }
520
521 assert_eq!(
522 args.mnemonic_indexes.expect("--mnemonic-indexes should have been set")[0],
523 test_case.2
524 )
525 }
526 }
527}