use crate::error::Result;
use crate::{
ansible::{inventory::AnsibleInventoryType, provisioning::AnsibleProvisioner},
error::Error,
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};
use std::collections::HashMap;
use std::str::FromStr;
const DEFAULT_TOKEN_AMOUNT: &str = "100_000_000_000_000_000_000";
const DEFAULT_GAS_AMOUNT: &str = "100_000_000_000_000_000";
pub struct FundingOptions {
pub evm_network: EvmNetwork,
pub evm_data_payments_address: Option<String>,
pub evm_payment_token_address: Option<String>,
pub evm_rpc_url: Option<String>,
pub funding_wallet_secret_key: Option<String>,
pub uploaders_count: Option<u16>,
pub token_amount: Option<U256>,
pub gas_amount: Option<U256>,
}
impl AnsibleProvisioner {
pub fn get_uploader_secret_keys(
&self,
) -> Result<HashMap<VirtualMachine, Vec<PrivateKeySigner>>> {
let uploaders_count = self.get_current_uploader_count()?;
debug!("Fetching uploader secret keys");
let mut uploader_secret_keys = HashMap::new();
if uploaders_count.is_empty() {
debug!("No uploaders VMs found");
return Err(Error::EmptyInventory(AnsibleInventoryType::Uploaders));
}
for (vm, count) in uploaders_count {
if count == 0 {
warn!("No uploader instances found for {:?}, ", vm.name);
uploader_secret_keys.insert(vm.clone(), Vec::new());
} else {
let sks = self.get_uploader_secret_key_per_vm(&vm, count)?;
uploader_secret_keys.insert(vm.clone(), sks);
}
}
Ok(uploader_secret_keys)
}
pub async fn deposit_funds_to_uploaders(
&self,
options: &FundingOptions,
) -> Result<HashMap<VirtualMachine, Vec<PrivateKeySigner>>> {
debug!("Funding all the uploader wallets");
let mut uploader_secret_keys = self.get_uploader_secret_keys()?;
for (vm, keys) in uploader_secret_keys.iter_mut() {
if let Some(provided_count) = options.uploaders_count {
if provided_count < keys.len() as u16 {
error!("Provided {provided_count} is less than the existing {} uploaders count for {}", keys.len(), vm.name);
return Err(Error::InvalidUpscaleDesiredUploaderCount);
}
let missing_keys_count = provided_count - keys.len() as u16;
debug!(
"Found {} secret keys for {}, missing {missing_keys_count} keys",
keys.len(),
vm.name
);
if missing_keys_count > 0 {
debug!(
"Generating {missing_keys_count} secret keys for {}",
vm.name
);
for _ in 0..missing_keys_count {
let sk = PrivateKeySigner::random();
debug!("Generated key with address: {}", sk.address());
keys.push(sk);
}
}
}
}
let funding_wallet_sk = if let Some(sk) = &options.funding_wallet_secret_key {
Some(sk.parse().map_err(|_| Error::FailedToParseKey)?)
} else {
None
};
self.deposit_funds(funding_wallet_sk, &uploader_secret_keys, options)
.await?;
Ok(uploader_secret_keys)
}
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()
)
})?;
println!(
"Draining {token_balance} tokens from {} to {to_address:?}",
from_wallet.address()
);
debug!(
"Draining {token_balance} tokens from {} to {to_address:?}",
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:?}",
)
})?;
println!(
"Drained {token_balance} tokens from {} to {to_address:?}",
from_wallet.address()
);
debug!(
"Drained {token_balance} tokens from {} to {to_address:?}",
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()
)
})?;
println!(
"Draining {gas_balance} gas from {} to {to_address:?}",
from_wallet.address()
);
debug!(
"Draining {gas_balance} gas from {} to {to_address:?}",
from_wallet.address()
);
if gas_balance.is_zero() {
debug!("No gas tokens to drain from wallet: {to_address}");
} else {
from_wallet
.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 from {to_address} with err: {err:?}",
)
})?;
println!(
"Drained {gas_balance} gas from {} to {to_address:?}",
from_wallet.address()
);
debug!(
"Drained {gas_balance} gas from {} to {to_address:?}",
from_wallet.address()
);
}
}
}
println!("All funds drained to {to_address:?} successfully");
debug!("All funds drained to {to_address:?} successfully");
Ok(())
}
fn get_current_uploader_count(&self) -> Result<HashMap<VirtualMachine, usize>> {
let uploader_inventories = self
.ansible_runner
.get_inventory(AnsibleInventoryType::Uploaders, true)?;
if uploader_inventories.is_empty() {
debug!("No uploaders VMs found");
return Err(Error::EmptyInventory(AnsibleInventoryType::Uploaders));
}
let mut uploader_count = HashMap::new();
for vm in uploader_inventories {
debug!(
"Fetching uploader count for {} @ {}",
vm.name, vm.public_ip_addr
);
let cmd = "systemctl list-units --type=service --all | grep ant_uploader_ | wc -l";
let result = self
.ssh_client
.run_command(&vm.public_ip_addr, "root", cmd, true);
match result {
Ok(count) => {
debug!("Count found to be {count:?}, parsing");
let count = count
.first()
.ok_or_else(|| {
error!("No count found for {}", vm.name);
Error::SecretKeyNotFound
})?
.trim()
.parse()
.map_err(|_| Error::FailedToParseKey)?;
uploader_count.insert(vm.clone(), count);
}
Err(Error::ExternalCommandRunFailed {
binary,
exit_status,
}) => {
if let Some(1) = exit_status.code() {
debug!("No uploaders found for {:?}", vm.public_ip_addr);
uploader_count.insert(vm.clone(), 0);
} else {
debug!("Error while fetching uploader count with different exit code {exit_status:?}",);
return Err(Error::ExternalCommandRunFailed {
binary,
exit_status,
});
}
}
Err(err) => {
debug!("Error while fetching uploader count: {err:?}",);
return Err(err);
}
}
}
Ok(uploader_count)
}
fn get_uploader_secret_key_per_vm(
&self,
vm: &VirtualMachine,
instance_count: usize,
) -> Result<Vec<PrivateKeySigner>> {
let mut sks_per_vm = Vec::new();
debug!(
"Fetching uploader secret key for {} @ {}",
vm.name, vm.public_ip_addr
);
for count in 1..=instance_count {
let cmd = format!(
"systemctl show ant_uploader_{count}.service --property=Environment | grep SECRET_KEY | cut -d= -f3 | awk '{{print $1}}'"
);
debug!("Fetching secret key for {} instance {count}", vm.name);
let result = self
.ssh_client
.run_command(&vm.public_ip_addr, "root", &cmd, true);
match result {
Ok(secret_keys) => {
let sk_str = secret_keys
.iter()
.map(|sk| sk.trim().to_string())
.collect::<Vec<String>>();
let sk_str = sk_str.first().ok_or({
debug!("No secret key found for {}", vm.name);
Error::SecretKeyNotFound
})?;
let sk = sk_str.parse().map_err(|_| Error::FailedToParseKey)?;
debug!("Secret keys found for {} instance {count}: {sk:?}", vm.name,);
sks_per_vm.push(sk);
}
Err(err) => {
debug!("Error while fetching secret key: {err}");
return Err(err);
}
}
}
Ok(sks_per_vm)
}
async fn deposit_funds(
&self,
funding_wallet_sk: Option<PrivateKeySigner>,
all_secret_keys: &HashMap<VirtualMachine, Vec<PrivateKeySigner>>,
options: &FundingOptions,
) -> Result<()> {
if all_secret_keys.is_empty() {
error!("No uploader secret keys found");
return Err(Error::SecretKeyNotFound);
}
let _sk_count = all_secret_keys.values().map(|v| v.len()).sum::<usize>();
let from_wallet = match &options.evm_network {
EvmNetwork::Anvil | EvmNetwork::Custom => {
let network = if let (
Some(evm_data_payments_address),
Some(evm_payment_token_address),
Some(evm_rpc_url),
) = (
options.evm_data_payments_address.as_ref(),
options.evm_payment_token_address.as_ref(),
options.evm_rpc_url.as_ref(),
) {
Network::new_custom(
evm_rpc_url,
evm_payment_token_address,
evm_data_payments_address,
)
} else {
error!("Custom evm network data not provided");
return Err(Error::EvmTestnetDataNotFound);
};
let Some(deployer_wallet_sk) = &options.funding_wallet_secret_key else {
error!("Deployer wallet secret key not provided");
return Err(Error::SecretKeyNotFound);
};
let deployer_wallet_sk: PrivateKeySigner = deployer_wallet_sk
.parse()
.map_err(|_| Error::FailedToParseKey)?;
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;
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;
Wallet::new(network.clone(), EthereumWallet::new(funding_wallet_sk))
}
};
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?;
println!("Funding wallet token balance: {token_balance}");
println!("Funding wallet gas balance: {gas_balance}");
debug!("Funding wallet token balance: {token_balance:?} and gas balance {gas_balance}");
let default_token_amount = U256::from_str(DEFAULT_TOKEN_AMOUNT).unwrap();
let default_gas_amount = U256::from_str(DEFAULT_GAS_AMOUNT).unwrap();
let token_amount = options.token_amount.unwrap_or(default_token_amount);
let gas_amount = options.gas_amount.unwrap_or(default_gas_amount);
println!("Transferring {token_amount} tokens and {gas_amount} gas tokens to each uploader");
debug!("Transferring {token_amount} tokens and {gas_amount} gas tokens to each uploader");
for (vm, vm_secret_keys) in all_secret_keys.iter() {
println!("Transferring funds for uploader vm: {}", vm.name);
for sk in vm_secret_keys.iter() {
sk.address();
if !token_amount.is_zero() {
print!("Transferring {token_amount} tokens to {}...", sk.address());
from_wallet
.transfer_tokens(sk.address(), token_amount)
.await
.inspect_err(|err| {
debug!(
"Failed to transfer {token_amount} tokens to {}: {err:?}",
sk.address()
)
})?;
println!("Transfer complete");
}
if !gas_amount.is_zero() {
print!("Transferring {gas_amount} gas to {}...", sk.address());
from_wallet
.transfer_gas_tokens(sk.address(), gas_amount)
.await
.inspect_err(|err| {
debug!(
"Failed to transfer {gas_amount} gas to {}: {err:?}",
sk.address()
)
})?;
println!("Transfer complete");
}
}
}
println!("All funds transferred successfully");
debug!("All funds transferred successfully");
Ok(())
}
}
pub fn get_address_from_sk(secret_key: &str) -> Result<Address> {
let sk: PrivateKeySigner = secret_key.parse().map_err(|_| Error::FailedToParseKey)?;
Ok(sk.address())
}