feat(aggregation-mode): support keystores in cli (#2243)

Author: MarcosNicolau

aggregation_mode/Cargo.lock

@@ -13,3 +13,4 @@ alloy = { workspace = true }
 agg_mode_sdk = { path = "../sdk"}
 sp1-sdk = "5.0.0"
 clap = { version = "4.5.4", features = ["derive"] }
+rpassword = "7.3.1"

aggregation_mode/cli/Cargo.toml

@@ -3,20 +3,18 @@ use alloy::{
     network::{EthereumWallet, TransactionBuilder},
     primitives::{Address, U256},
     providers::{Provider, ProviderBuilder},
-    signers::local::LocalSigner,
 };
 use clap::{self, Args};
-use std::str::FromStr;
 
-use crate::commands::helpers::parse_network;
+use crate::commands::helpers::{parse_network, PrivateKeyType};
 
 const PAYMENT_AMOUNT: &str = "1"; // ether
 
 /// Send 1 ether to the aggregation mode payment service to fund proof submissions
 #[derive(Debug, Clone, Args)]
 pub struct SendPaymentArgs {
-    #[arg(long = "private-key")]
-    private_key: String,
+    #[command(flatten)]
+    private_key_type: PrivateKeyType,
     #[arg(short = 'n', long = "network", default_value = "devnet", value_parser = parse_network)]
     network: Network,
     #[arg(long = "rpc-url")]
@@ -29,10 +27,10 @@ pub async fn run(args: SendPaymentArgs) {
         args.network
     );
 
-    let signer = match LocalSigner::from_str(args.private_key.trim()) {
+    let signer = match args.private_key_type.into_signer() {
         Ok(s) => s,
         Err(e) => {
-            tracing::error!("Failed to parse private key: {e}");
+            tracing::error!("{e}");
             return;
         }
     };

aggregation_mode/cli/src/commands/deposit.rs

@@ -1,4 +1,6 @@
-use clap::{self, ValueEnum};
+use alloy::signers::local::{LocalSigner, PrivateKeySigner};
+use clap::{self, Args, ValueEnum};
+use std::path::PathBuf;
 use std::str::FromStr;
 
 use agg_mode_sdk::types::Network;
@@ -14,3 +16,30 @@ pub enum ProvingSystemArg {
     #[clap(name = "Risc0")]
     Risc0,
 }
+
+#[derive(Args, Debug, Clone)]
+#[group(required = true, multiple = false)]
+pub struct PrivateKeyType {
+    #[arg(name = "keystore_path", long = "keystore-path")]
+    pub keystore_path: Option<PathBuf>,
+    #[arg(name = "private_key", long = "private-key")]
+    pub private_key: Option<String>,
+}
+
+impl PrivateKeyType {
+    /// Creates a LocalSigner from either a keystore file or a raw private key.
+    /// If a keystore path is provided, prompts for the password interactively.
+    pub fn into_signer(self) -> Result<PrivateKeySigner, String> {
+        if let Some(keystore_path) = self.keystore_path {
+            let password = rpassword::prompt_password("Please enter your keystore password: ")
+                .map_err(|e| format!("Failed to read password: {e}"))?;
+            LocalSigner::decrypt_keystore(&keystore_path, password)
+                .map_err(|e| format!("Failed to decrypt keystore: {e}"))
+        } else if let Some(private_key) = self.private_key {
+            LocalSigner::from_str(private_key.trim())
+                .map_err(|e| format!("Failed to parse private key: {e}"))
+        } else {
+            Err("Either --keystore-path or --private-key must be provided".to_string())
+        }
+    }
+}

aggregation_mode/cli/src/commands/helpers.rs

@@ -1,10 +1,9 @@
 use agg_mode_sdk::{gateway::provider::AggregationModeGatewayProvider, types::Network};
-use alloy::signers::local::LocalSigner;
 use clap::{command, Args, Subcommand};
 use sp1_sdk::{SP1ProofWithPublicValues, SP1VerifyingKey};
-use std::{path::PathBuf, str::FromStr};
+use std::path::PathBuf;
 
-use crate::commands::helpers::parse_network;
+use crate::commands::helpers::{parse_network, PrivateKeyType};
 
 #[derive(Debug, Subcommand)]
 pub enum SubmitCommand {
@@ -18,8 +17,8 @@ pub struct SubmitSP1Args {
     proof_path: PathBuf,
     #[arg(long = "vk")]
     verifying_key_path: PathBuf,
-    #[arg(long = "private-key")]
-    private_key: String,
+    #[command(flatten)]
+    private_key_type: PrivateKeyType,
     #[arg(short = 'n', long = "network", default_value = "devnet", value_parser = parse_network)]
     network: Network,
 }
@@ -30,8 +29,13 @@ pub async fn run(args: SubmitSP1Args) {
     let proof = load_proof(&args.proof_path).expect("Valid proof");
     let vk = load_vk(&args.verifying_key_path).expect("Valid vk");
 
-    let signer =
-        LocalSigner::from_str(args.private_key.trim()).expect("failed to parse private key: {e}");
+    let signer = match args.private_key_type.into_signer() {
+        Ok(s) => s,
+        Err(e) => {
+            tracing::error!("{e}");
+            return;
+        }
+    };
 
     let provider = AggregationModeGatewayProvider::new_with_signer(args.network.clone(), signer)
         .expect("failed to initialize gateway client: {e:?}");