feat: drain funds on cleanup

src/ansible/provisioning.rs

@@ -458,7 +458,8 @@ impl AnsibleProvisioner {
         let start = Instant::now();
 
         let sk_map = self
-            .fund_uploader_wallets(&FundingOptions {
+            .deposit_funds_to_uploaders(&FundingOptions {
+                custom_evm_testnet_data: evm_testnet_data.clone(),
                 uploaders_count: options.uploaders_count,
                 evm_network: options.evm_network.clone(),
                 funding_wallet_secret_key: options.funding_wallet_secret_key.clone(),

src/bootstrap.rs

@@ -56,6 +56,7 @@ impl TestnetDeployer {
                 environment_type: options.environment_type.clone(),
                 evm_network: options.evm_network.clone(),
                 evm_testnet_data: None,
+                funding_wallet_address: None,
                 rewards_address: options.rewards_address.clone(),
             },
         )

src/deploy.rs

@@ -7,10 +7,12 @@
 use crate::{
     ansible::{inventory::AnsibleInventoryType, provisioning::ProvisionOptions},
     error::Result,
+    funding::get_address_from_sk,
     get_evm_testnet_data, get_genesis_multiaddr, write_environment_details, BinaryOption,
     DeploymentInventory, DeploymentType, EnvironmentDetails, EnvironmentType, EvmNetwork,
     InfraRunOptions, LogFormat, NodeType, TestnetDeployer,
 };
+use alloy::hex::ToHexExt;
 use colored::Colorize;
 use std::{net::SocketAddr, path::PathBuf};
 
@@ -104,8 +106,9 @@ impl TestnetDeployer {
                 deployment_type: DeploymentType::New,
                 environment_type: options.environment_type.clone(),
                 evm_network: options.evm_network.clone(),
-                rewards_address: options.rewards_address.clone(),
                 evm_testnet_data: None,
+                funding_wallet_address: None,
+                rewards_address: options.rewards_address.clone(),
             },
         )
         .await?;
@@ -133,6 +136,17 @@ impl TestnetDeployer {
             None
         };
 
+        let funding_wallet_address = if let Some(secret_key) = &options.funding_wallet_secret_key {
+            let address = get_address_from_sk(secret_key)?;
+            Some(address.encode_hex())
+        } else if let Some(emv_data) = &evm_testnet_data {
+            let address = get_address_from_sk(&emv_data.deployer_wallet_private_key)?;
+            Some(address.encode_hex())
+        } else {
+            log::error!("Funding wallet address not provided");
+            None
+        };
+
         write_environment_details(
             &self.s3_repository,
             &options.name,
@@ -141,6 +155,7 @@ impl TestnetDeployer {
                 environment_type: options.environment_type.clone(),
                 evm_network: options.evm_network.clone(),
                 evm_testnet_data: evm_testnet_data.clone(),
+                funding_wallet_address,
                 rewards_address: options.rewards_address.clone(),
             },
         )

src/error.rs

@@ -55,6 +55,8 @@ pub enum Error {
         binary: String,
         exit_status: std::process::ExitStatus,
     },
+    #[error("Failed to parse key")]
+    FailedToParseKey,
     #[error("Failed to retrieve filename")]
     FilenameNotRetrieved,
     #[error(transparent)]
@@ -133,8 +135,6 @@ pub enum Error {
     SafeBinaryDownloadError,
     #[error("Error in byte stream when attempting to retrieve S3 object")]
     S3ByteStreamError,
-    #[error("Failed to parse secret key")]
-    SecretKeyParseError,
     #[error("The secret key was not found in the environment")]
     SecretKeyNotFound,
     #[error(transparent)]

src/funding.rs

@@ -5,13 +5,14 @@
 // Please see the LICENSE file for more details.
 
 use crate::error::Result;
+use crate::EvmCustomTestnetData;
 use crate::{
     ansible::{inventory::AnsibleInventoryType, provisioning::AnsibleProvisioner},
     error::Error,
-    get_evm_testnet_data,
     inventory::VirtualMachine,
     EvmNetwork,
 };
+use alloy::primitives::Address;
 use alloy::{network::EthereumWallet, signers::local::PrivateKeySigner};
 use evmlib::{common::U256, wallet::Wallet, Network};
 use log::{debug, error, warn};
@@ -20,6 +21,7 @@ use std::str::FromStr;
 
 pub struct FundingOptions {
     pub evm_network: EvmNetwork,
+    pub custom_evm_testnet_data: Option<EvmCustomTestnetData>,
     pub funding_wallet_secret_key: Option<String>,
     /// Have to specify during upscale and deploy
     pub uploaders_count: Option<u16>,
@@ -59,10 +61,10 @@ impl AnsibleProvisioner {
         Ok(uploader_secret_keys)
     }
 
-    /// Send funds from the funding_wallet_secret_key to the uploader wallets
+    /// Deposit funds from the funding_wallet_secret_key to the uploader wallets
     /// If FundingOptions::uploaders_count is provided, it will generate the missing secret keys.
     /// If not provided, we'll just fund the existing uploader wallets
-    pub async fn fund_uploader_wallets(
+    pub async fn deposit_funds_to_uploaders(
         &self,
         options: &FundingOptions,
     ) -> Result<HashMap<VirtualMachine, Vec<PrivateKeySigner>>> {
@@ -96,17 +98,95 @@ impl AnsibleProvisioner {
         }
 
         let funding_wallet_sk = if let Some(sk) = &options.funding_wallet_secret_key {
-            Some(sk.parse().map_err(|_| Error::SecretKeyParseError)?)
+            Some(sk.parse().map_err(|_| Error::FailedToParseKey)?)
         } else {
             None
         };
 
-        self.transfer_funds(funding_wallet_sk, &uploader_secret_keys, options)
+        self.deposit_funds(funding_wallet_sk, &uploader_secret_keys, options)
             .await?;
 
         Ok(uploader_secret_keys)
     }
 
+    /// Drain all the funds from the uploader wallets to the provided wallet
+    pub async fn drain_funds_from_uploaders(
+        &self,
+        to_address: Address,
+        evm_network: Network,
+    ) -> Result<()> {
+        debug!("Draining all the uploader wallets to {to_address:?}");
+        println!("Draining all the uploader wallets to {to_address:?}");
+        let uploader_secret_keys = self.get_uploader_secret_keys()?;
+
+        for (vm, keys) in uploader_secret_keys.iter() {
+            debug!(
+                "Draining funds for uploader vm: {} to {to_address:?}",
+                vm.name
+            );
+            for uploader_sk in keys.iter() {
+                debug!(
+                    "Draining funds for uploader vm: {} with key: {uploader_sk:?}",
+                    vm.name,
+                );
+
+                let from_wallet = Wallet::new(
+                    evm_network.clone(),
+                    EthereumWallet::new(uploader_sk.clone()),
+                );
+
+                let token_balance = from_wallet.balance_of_tokens().await.inspect_err(|err| {
+                    debug!(
+                        "Failed to get token balance for {} with err: {err:?}",
+                        from_wallet.address()
+                    )
+                })?;
+                let gas_balance = from_wallet
+                    .balance_of_gas_tokens()
+                    .await
+                    .inspect_err(|err| {
+                        debug!(
+                            "Failed to get gas token balance for {} with err: {err:?}",
+                            from_wallet.address()
+                        )
+                    })?;
+
+                if token_balance.is_zero() {
+                    debug!(
+                        "No tokens to drain from wallet: {} with token balance",
+                        from_wallet.address()
+                    );
+                } else {
+                    from_wallet
+                        .transfer_tokens(to_address, token_balance)
+                        .await
+                        .inspect_err(|err| {
+                            debug!(
+                                "Failed to transfer {token_balance} tokens from {to_address} with err: {err:?}",
+                            )
+                        })?;
+                }
+
+                if gas_balance.is_zero() {
+                    debug!("No gas tokens to drain from wallet: {to_address}");
+                } else {
+                    from_wallet
+                    // 0.00001 gas
+                        .transfer_gas_tokens(to_address, gas_balance - U256::from_str("10_000_000_000_000").unwrap()).await
+                        .inspect_err(|err| {
+                            debug!(
+                                "Failed to transfer {gas_balance} gas tokens from {to_address} with err: {err:?}",
+                            )
+                        })?;
+                }
+            }
+        }
+        println!("All funds drained to {to_address:?} successfully");
+        debug!("All funds drained to {to_address:?} successfully");
+
+        Ok(())
+    }
+
     /// Return the (vm name, uploader count) for all uploader VMs
     fn get_current_uploader_count(&self) -> Result<HashMap<VirtualMachine, usize>> {
         let uploader_inventories = self
@@ -139,7 +219,7 @@ impl AnsibleProvisioner {
                         })?
                         .trim()
                         .parse()
-                        .map_err(|_| Error::SecretKeyParseError)?;
+                        .map_err(|_| Error::FailedToParseKey)?;
                     uploader_count.insert(vm.clone(), count);
                 }
                 Err(Error::ExternalCommandRunFailed {
@@ -198,7 +278,7 @@ impl AnsibleProvisioner {
                         debug!("No secret key found for {}", vm.name);
                         Error::SecretKeyNotFound
                     })?;
-                    let sk = sk_str.parse().map_err(|_| Error::SecretKeyParseError)?;
+                    let sk = sk_str.parse().map_err(|_| Error::FailedToParseKey)?;
 
                     debug!("Secret keys found for {} instance {count}: {sk:?}", vm.name,);
 
@@ -214,7 +294,7 @@ impl AnsibleProvisioner {
         Ok(sks_per_vm)
     }
 
-    async fn transfer_funds(
+    async fn deposit_funds(
         &self,
         funding_wallet_sk: Option<PrivateKeySigner>,
         all_secret_keys: &HashMap<VirtualMachine, Vec<PrivateKeySigner>>,
@@ -227,10 +307,13 @@ impl AnsibleProvisioner {
 
         let _sk_count = all_secret_keys.values().map(|v| v.len()).sum::<usize>();
 
-        let (from_wallet, network) = match &options.evm_network {
+        let from_wallet = match &options.evm_network {
             EvmNetwork::Custom => {
                 let evm_testnet_data =
-                    get_evm_testnet_data(&self.ansible_runner, &self.ssh_client)?;
+                    options.custom_evm_testnet_data.as_ref().ok_or_else(|| {
+                        error!("Custom Evm testnet data not provided");
+                        Error::EvmTestnetDataNotFound
+                    })?;
                 let network = Network::new_custom(
                     &evm_testnet_data.rpc_url,
                     &evm_testnet_data.payment_token_address,
@@ -239,31 +322,28 @@ impl AnsibleProvisioner {
                 let deployer_wallet_sk: PrivateKeySigner = evm_testnet_data
                     .deployer_wallet_private_key
                     .parse()
-                    .map_err(|_| Error::SecretKeyParseError)?;
+                    .map_err(|_| Error::FailedToParseKey)?;
 
-                let wallet = Wallet::new(network.clone(), EthereumWallet::new(deployer_wallet_sk));
-                (wallet, network)
+                Wallet::new(network.clone(), EthereumWallet::new(deployer_wallet_sk))
             }
             EvmNetwork::ArbitrumOne => {
                 let funding_wallet_sk = funding_wallet_sk.ok_or_else(|| {
                     error!("Funding wallet secret key not provided");
                     Error::SecretKeyNotFound
                 })?;
                 let network = Network::ArbitrumOne;
-                let wallet = Wallet::new(network.clone(), EthereumWallet::new(funding_wallet_sk));
-                (wallet, network)
+                Wallet::new(network.clone(), EthereumWallet::new(funding_wallet_sk))
             }
             EvmNetwork::ArbitrumSepolia => {
                 let funding_wallet_sk = funding_wallet_sk.ok_or_else(|| {
                     error!("Funding wallet secret key not provided");
                     Error::SecretKeyNotFound
                 })?;
                 let network = Network::ArbitrumSepolia;
-                let wallet = Wallet::new(network.clone(), EthereumWallet::new(funding_wallet_sk));
-                (wallet, network)
+                Wallet::new(network.clone(), EthereumWallet::new(funding_wallet_sk))
             }
         };
-        debug!("Using emv network: {network:?}",);
+        debug!("Using emv network: {:?}", options.evm_network);
 
         let token_balance = from_wallet.balance_of_tokens().await?;
         let gas_balance = from_wallet.balance_of_gas_tokens().await?;
@@ -289,26 +369,26 @@ impl AnsibleProvisioner {
         for (vm, sks_per_machine) in all_secret_keys.iter() {
             debug!("Transferring funds for uploader vm: {}", vm.name);
             for sk in sks_per_machine.iter() {
-                let to_wallet = Wallet::new(network.clone(), EthereumWallet::new(sk.clone()));
+                sk.address();
                 debug!(
                     "Transferring funds for uploader vm: {} with public key: {}",
                     vm.name,
-                    to_wallet.address()
+                    sk.address()
                 );
 
                 from_wallet
-                    .transfer_tokens(to_wallet.address(), tokens_for_each_uploader)
+                    .transfer_tokens(sk.address(), tokens_for_each_uploader)
                     .await.inspect_err(|err| {
                         debug!(
-                            "Failed to transfer {tokens_for_each_uploader} tokens to {} with err: {err:?}", to_wallet.address()
+                            "Failed to transfer {tokens_for_each_uploader} tokens to {} with err: {err:?}", sk.address()
                         )
                     })?;
                 from_wallet
-                    .transfer_gas_tokens(to_wallet.address(), gas_tokens_for_each_uploader)
+                    .transfer_gas_tokens(sk.address(), gas_tokens_for_each_uploader)
                     .await
                     .inspect_err(|err| {
                         debug!(
-                            "Failed to transfer {gas_tokens_for_each_uploader} gas tokens to {} with err: {err:?}", to_wallet.address()
+                            "Failed to transfer {gas_tokens_for_each_uploader} gas tokens to {} with err: {err:?}", sk.address()
                         )
                     })
                     ?;
@@ -320,3 +400,9 @@ impl AnsibleProvisioner {
         Ok(())
     }
 }
+
+/// Get the Address of the funding wallet from the secret key string
+pub fn get_address_from_sk(secret_key: &str) -> Result<Address> {
+    let sk: PrivateKeySigner = secret_key.parse().map_err(|_| Error::FailedToParseKey)?;
+    Ok(sk.address())
+}

src/lib.rs

@@ -36,6 +36,8 @@ use crate::{
     ssh::SshClient,
     terraform::TerraformRunner,
 };
+use alloy::primitives::Address;
+use evmlib::Network;
 use flate2::read::GzDecoder;
 use indicatif::{ProgressBar, ProgressStyle};
 use log::{debug, trace};
@@ -164,6 +166,7 @@ pub struct EnvironmentDetails {
     pub environment_type: EnvironmentType,
     pub evm_network: EvmNetwork,
     pub evm_testnet_data: Option<EvmCustomTestnetData>,
+    pub funding_wallet_address: Option<String>,
     pub rewards_address: String,
 }
 
@@ -718,11 +721,34 @@ impl TestnetDeployer {
     }
 
     pub async fn clean(&self) -> Result<()> {
-        let environment_type =
+        let environment_details =
             get_environment_details(&self.environment_name, &self.s3_repository).await?;
+
+        let evm_network = match environment_details.evm_network {
+            EvmNetwork::Custom => None,
+            EvmNetwork::ArbitrumOne => Some(Network::ArbitrumOne),
+            EvmNetwork::ArbitrumSepolia => Some(Network::ArbitrumSepolia),
+        };
+        if let (Some(network), Some(address)) =
+            (evm_network, environment_details.funding_wallet_address)
+        {
+            self.ansible_provisioner
+                .drain_funds_from_uploaders(
+                    Address::from_str(&address).map_err(|err| {
+                        log::error!("Invalid funding wallet public key: {err:?}");
+                        Error::FailedToParseKey
+                    })?,
+                    network,
+                )
+                .await?;
+        } else {
+            println!("Custom network provided. Not draining funds.");
+            log::info!("Custom network provided. Not draining funds.");
+        }
+
         do_clean(
             &self.environment_name,
-            Some(environment_type.environment_type),
+            Some(environment_details.environment_type),
             self.working_directory_path.clone(),
             &self.terraform_runner,
             None,