Add --save-path Option to VanityArgs for Custom Wallet File Saving

crates/cast/bin/cmd/wallet/vanity.rs

@@ -6,7 +6,12 @@ use eyre::Result;
 use foundry_common::types::ToAlloy;
 use rayon::iter::{self, ParallelIterator};
 use regex::Regex;
-use std::time::Instant;
+use serde::{Deserialize, Serialize};
+use std::{
+    fs,
+    path::{Path, PathBuf},
+    time::Instant,
+};
 
 /// Type alias for the result of [generate_wallet].
 pub type GeneratedWallet = (SigningKey, Address);
@@ -32,11 +37,44 @@ pub struct VanityArgs {
     /// nonce.
     #[clap(long)]
     pub nonce: Option<u64>,
+
+    /// Path to save the generated vanity contract address to.
+    ///
+    /// If provided, the generated vanity addresses will appended to a JSON array in the specified
+    /// file.
+    #[clap(
+        long,
+        value_hint = clap::ValueHint::FilePath,
+        value_name = "PATH",
+    )]
+    pub save_path: Option<PathBuf>,
+}
+
+/// WalletData contains address and private_key information for a wallet.
+#[derive(Serialize, Deserialize)]
+struct WalletData {
+    address: String,
+    private_key: String,
+}
+
+/// Wallets is a collection of WalletData.
+#[derive(Default, Serialize, Deserialize)]
+struct Wallets {
+    wallets: Vec<WalletData>,
+}
+
+impl WalletData {
+    pub fn new(wallet: &LocalWallet) -> Self {
+        WalletData {
+            address: wallet.address().to_alloy().to_checksum(None),
+            private_key: format!("0x{}", hex::encode(wallet.signer().to_bytes())),
+        }
+    }
 }
 
 impl VanityArgs {
     pub fn run(self) -> Result<LocalWallet> {
-        let Self { starts_with, ends_with, nonce } = self;
+        let Self { starts_with, ends_with, nonce, save_path } = self;
         let mut left_exact_hex = None;
         let mut left_regex = None;
         let mut right_exact_hex = None;
@@ -108,6 +146,11 @@ impl VanityArgs {
         }
         .expect("failed to generate vanity wallet");
 
+        // If a save path is provided, save the generated vanity wallet to the specified path.
+        if let Some(save_path) = save_path {
+            save_wallet_to_file(&wallet, &save_path)?;
+        }
+
         println!(
             "Successfully found vanity address in {} seconds.{}{}\nAddress: {}\nPrivate Key: 0x{}",
             timer.elapsed().as_secs(),
@@ -125,6 +168,23 @@ impl VanityArgs {
     }
 }
 
+/// Saves the specified `wallet` to a 'vanity_addresses.json' file at the given `save_path`.
+/// If the file exists, the wallet data is appended to the existing content;
+/// otherwise, a new file is created.
+fn save_wallet_to_file(wallet: &LocalWallet, path: &Path) -> Result<()> {
+    let mut wallets = if path.exists() {
+        let data = fs::read_to_string(&path)?;
+        serde_json::from_str::<Wallets>(&data).unwrap_or_default()
+    } else {
+        Wallets::default()
+    };
+
+    wallets.wallets.push(WalletData::new(wallet));
+
+    fs::write(&path, serde_json::to_string_pretty(&wallets)?)?;
+    Ok(())
+}
+
 /// Generates random wallets until `matcher` matches the wallet address, returning the wallet.
 pub fn find_vanity_address<T: VanityMatcher>(matcher: T) -> Option<LocalWallet> {
     wallet_generator().find_any(create_matcher(matcher)).map(|(key, _)| key.into())
@@ -328,4 +388,21 @@ mod tests {
         let addr = format!("{addr:x}");
         assert!(addr.ends_with("00"));
     }
+
+    #[test]
+    fn save_path() {
+        let tmp = tempfile::NamedTempFile::new().unwrap();
+        let args: VanityArgs = VanityArgs::parse_from([
+            "foundry-cli",
+            "--starts-with",
+            "00",
+            "--save-path",
+            tmp.path().to_str().unwrap(),
+        ]);
+        args.run().unwrap();
+        assert!(tmp.path().exists());
+        let s = fs::read_to_string(tmp.path()).unwrap();
+        let wallets: Wallets = serde_json::from_str(&s).unwrap();
+        assert!(!wallets.wallets.is_empty());
+    }
 }