From 215f6f495d41ddfeb63315ec909ecc282f8bf3fe Mon Sep 17 00:00:00 2001
From: rplusq <rafaplusq@gmail.com>
Date: Sun, 6 Aug 2023 21:17:15 +0100
Subject: [PATCH] MultiWallet/Wallet now leverage imported keystores

---
 cli/src/opts/wallet/mod.rs          | 28 +++++++++++++++++++--
 cli/src/opts/wallet/multi_wallet.rs | 39 +++++++++++++++++++++++++++--
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/cli/src/opts/wallet/mod.rs b/cli/src/opts/wallet/mod.rs
index 023a8bc291c5..d3a5fb97d18f 100644
--- a/cli/src/opts/wallet/mod.rs
+++ b/cli/src/opts/wallet/mod.rs
@@ -110,6 +110,16 @@ pub struct Wallet {
     )]
     pub keystore_path: Option<String>,
 
+    // Use a keystore from the default folder by its name
+    #[clap(
+        long = "account",
+        help_heading = "Wallet options - keystore",
+        value_name = "ACCOUNT_NAME",
+        env = "ETH_KEYSTORE_ACCOUNT",
+        conflicts_with = "keystore_path"
+    )]
+    pub keystore_account_name: Option<String>,
+
     /// The keystore password.
     ///
     /// Used with --keystore.
@@ -169,8 +179,17 @@ impl Wallet {
     }
 
     pub fn keystore(&self) -> Result<Option<LocalWallet>> {
+        let default_keystore_dir = dirs::home_dir()
+        .ok_or_else(|| eyre::eyre!("Failed to get home directory"))?
+        .join(".foundry")
+        .join("keystores");
+    
+        // If keystore path is provided, use it, otherwise use default path + keystore account name
+        let keystore_path: Option<String> = self.keystore_path.clone()
+            .or_else(|| self.keystore_account_name.as_ref().map(|keystore_name| default_keystore_dir.join(keystore_name).to_string_lossy().into_owned()));
+
         self.get_from_keystore(
-            self.keystore_path.as_ref(),
+            keystore_path.as_ref(),
             self.keystore_password.as_ref(),
             self.keystore_password_file.as_ref(),
         )
@@ -366,6 +385,7 @@ pub trait WalletTrait {
         keystore_password_file: Option<&String>,
     ) -> Result<Option<LocalWallet>> {
         Ok(match (keystore_path, keystore_password, keystore_password_file) {
+            // Path and password provided
             (Some(path), Some(password), _) => {
                 let path = self.find_keystore_file(path)?;
                 Some(
@@ -373,18 +393,21 @@ pub trait WalletTrait {
                         .wrap_err_with(|| format!("Failed to decrypt keystore {path:?}"))?,
                 )
             }
-            (Some(path), _, Some(password_file)) => {
+            // Path and password file provided
+            (Some(path),  _, Some(password_file)) => {
                 let path = self.find_keystore_file(path)?;
                 Some(
                     LocalWallet::decrypt_keystore(&path, self.password_from_file(password_file)?)
                         .wrap_err_with(|| format!("Failed to decrypt keystore {path:?} with password file {password_file:?}"))?,
                 )
             }
+            // Only Path provided -> interactive
             (Some(path), None, None) => {
                 let path = self.find_keystore_file(path)?;
                 let password = rpassword::prompt_password("Enter keystore password:")?;
                 Some(LocalWallet::decrypt_keystore(path, password)?)
             }
+            // Nothing provided
             (None, _, _) => None,
         })
     }
@@ -574,6 +597,7 @@ mod tests {
             },
             from: None,
             keystore_path: None,
+            keystore_account_name: None,
             keystore_password: None,
             keystore_password_file: None,
             ledger: false,
diff --git a/cli/src/opts/wallet/multi_wallet.rs b/cli/src/opts/wallet/multi_wallet.rs
index 0cb3e520f415..186ba849eb03 100644
--- a/cli/src/opts/wallet/multi_wallet.rs
+++ b/cli/src/opts/wallet/multi_wallet.rs
@@ -153,6 +153,16 @@ pub struct MultiWallet {
     )]
     pub keystore_paths: Option<Vec<String>>,
 
+    // Use a keystore from the default folder by its name
+    #[clap(
+        long = "account",
+        help_heading = "Wallet options - keystore",
+        value_name = "ACCOUNT_NAMES",
+        env = "ETH_KEYSTORE_ACCOUNT",
+        conflicts_with = "keystore_paths"
+    )]
+    pub keystore_account_names: Option<Vec<String>>,
+
     /// The keystore password.
     ///
     /// Used with --keystore.
@@ -286,7 +296,32 @@ impl MultiWallet {
     ///
     /// Returns `Ok(None)` if no keystore provided.
     pub fn keystores(&self) -> Result<Option<Vec<LocalWallet>>> {
-        if let Some(keystore_paths) = &self.keystore_paths {
+        let default_keystore_dir = dirs::home_dir()
+            .ok_or_else(|| eyre::eyre!("Failed to get home directory"))?
+            .join(".foundry")
+            .join("keystores");
+
+        // If keystore paths are provided, use them, otherwise use default path + keystore account names
+        let keystore_paths = self
+            .keystore_paths
+            .clone()
+            .or_else(|| {
+                self.keystore_account_names
+                    .as_ref()
+                    .map(|keystore_names| {
+                        keystore_names
+                            .iter()
+                            .map(|keystore_name| {
+                                default_keystore_dir
+                                    .join(keystore_name)
+                                    .to_string_lossy()
+                                    .into_owned()
+                            })
+                            .collect()
+                    })
+            });
+
+        if let Some(keystore_paths) = keystore_paths {
             let mut wallets = Vec::with_capacity(keystore_paths.len());
 
             let mut passwords_iter =
@@ -296,7 +331,7 @@ impl MultiWallet {
                 self.keystore_password_files.clone().unwrap_or_default().into_iter();
 
             for path in keystore_paths {
-                wallets.push(self.get_from_keystore(Some(path), passwords_iter.next().as_ref(), password_files_iter.next().as_ref())?.wrap_err("Keystore paths do not have the same length as provided passwords or password files.")?);
+                wallets.push(self.get_from_keystore(Some(&path), passwords_iter.next().as_ref(), password_files_iter.next().as_ref())?.wrap_err("Keystore paths do not have the same length as provided passwords or password files.")?);
             }
             return Ok(Some(wallets))
         }