From b81aafbc2334ba6c45c3aa0347ddf4789957d074 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Wed, 3 Dec 2025 13:10:07 +0200 Subject: [PATCH] feat: UX improvements --- integration_tests/src/test_suite_map.rs | 85 ++++--------- .../src/key_management/key_tree/mod.rs | 114 ++++++++++++++---- key_protocol/src/key_protocol_core/mod.rs | 37 ++++-- wallet/src/cli/account.rs | 14 ++- wallet/src/cli/config.rs | 9 -- wallet/src/cli/mod.rs | 48 +++++--- wallet/src/lib.rs | 10 +- wallet/src/main.rs | 19 +-- 8 files changed, 192 insertions(+), 144 deletions(-) diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs index 8eb849f..d5bed4c 100644 --- a/integration_tests/src/test_suite_map.rs +++ b/integration_tests/src/test_suite_map.rs @@ -87,9 +87,7 @@ pub fn prepare_function_map() -> HashMap { #[nssa_integration_test] pub async fn test_success_move_to_another_account() { info!("########## test_success_move_to_another_account ##########"); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), - })); + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })); let wallet_config = fetch_config().await.unwrap(); @@ -293,9 +291,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -306,9 +302,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -319,9 +313,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -453,9 +445,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -466,9 +456,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -479,9 +467,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -614,9 +600,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -627,9 +611,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -640,9 +622,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -756,9 +736,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -769,9 +747,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -782,9 +758,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -898,9 +872,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -911,9 +883,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -924,9 +894,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -1127,9 +1095,8 @@ pub fn prepare_function_map() -> HashMap { ); let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap(); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: ChainIndex::root(), - })); + let command = + Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); let SubcommandReturnValue::RegisterAccount { @@ -1490,9 +1457,7 @@ pub fn prepare_function_map() -> HashMap { #[nssa_integration_test] pub async fn test_authenticated_transfer_initialize_function() { info!("########## test initialize account for authenticated transfer ##########"); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), - })); + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })); let SubcommandReturnValue::RegisterAccount { account_id } = wallet::cli::execute_subcommand(command).await.unwrap() else { @@ -1590,9 +1555,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: winner_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -1676,7 +1639,7 @@ pub fn prepare_function_map() -> HashMap { let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap(); let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); @@ -1688,7 +1651,7 @@ pub fn prepare_function_map() -> HashMap { }; let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: ChainIndex::from_str("/0").unwrap(), + cci: Some(ChainIndex::from_str("/0").unwrap()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); @@ -1726,7 +1689,7 @@ pub fn prepare_function_map() -> HashMap { let from: AccountId = ACC_SENDER.parse().unwrap(); let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); @@ -1738,7 +1701,7 @@ pub fn prepare_function_map() -> HashMap { }; let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::from_str("/0").unwrap(), + cci: Some(ChainIndex::from_str("/0").unwrap()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 3d55516..beb0f95 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeMap, HashMap}, + collections::{BTreeMap, HashMap, VecDeque}, sync::Arc, }; @@ -20,6 +20,8 @@ pub mod keys_private; pub mod keys_public; pub mod traits; +pub const DEPTH_SOFT_CAP: u32 = 20; + #[derive(Debug, Serialize, Deserialize, Clone)] pub struct KeyTree { pub key_map: BTreeMap, @@ -101,10 +103,13 @@ impl KeyTree { } } - pub fn generate_new_node(&mut self, parent_cci: ChainIndex) -> Option { - let father_keys = self.key_map.get(&parent_cci)?; + pub fn generate_new_node( + &mut self, + parent_cci: &ChainIndex, + ) -> Option<(nssa::AccountId, ChainIndex)> { + let father_keys = self.key_map.get(parent_cci)?; let next_child_id = self - .find_next_last_child_of_id(&parent_cci) + .find_next_last_child_of_id(parent_cci) .expect("Can be None only if parent is not present"); let next_cci = parent_cci.nth_child(next_child_id); @@ -113,9 +118,43 @@ impl KeyTree { let account_id = child_keys.account_id(); self.key_map.insert(next_cci.clone(), child_keys); - self.account_id_map.insert(account_id, next_cci); + self.account_id_map.insert(account_id, next_cci.clone()); - Some(account_id) + Some((account_id, next_cci)) + } + + fn have_child_slot_capped(&self, cci: &ChainIndex) -> bool { + let depth = cci.depth(); + + self.find_next_last_child_of_id(cci) + .map(|inn| inn + 1 + depth < DEPTH_SOFT_CAP) + .unwrap_or(false) + } + + pub fn search_new_parent_capped(&self) -> Option { + let mut parent_list = VecDeque::new(); + parent_list.push_front(ChainIndex::root()); + + let mut search_res = None; + + while let Some(next_parent) = parent_list.pop_back() { + if self.have_child_slot_capped(&next_parent) { + search_res = Some(next_parent); + break; + } else { + let last_child = self.find_next_last_child_of_id(&next_parent)?; + + for id in 0..last_child { + parent_list.push_front(next_parent.nth_child(id)); + } + } + } + + search_res + } + + pub fn generate_new_node_capped(&mut self) -> Option<(nssa::AccountId, ChainIndex)> { + self.generate_new_node(&self.search_new_parent_capped()?) } pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> { @@ -150,7 +189,7 @@ impl KeyTree { let mut id_stack = vec![ChainIndex::root()]; while let Some(curr_id) = id_stack.pop() { - self.generate_new_node(curr_id.clone()); + self.generate_new_node(&curr_id); let mut next_id = curr_id.nth_child(0); @@ -268,7 +307,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -281,12 +320,12 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -307,7 +346,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -320,7 +359,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - let key_opt = tree.generate_new_node(ChainIndex::from_str("/3").unwrap()); + let key_opt = tree.generate_new_node(&ChainIndex::from_str("/3").unwrap()); assert_eq!(key_opt, None); } @@ -337,7 +376,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -350,7 +389,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -363,7 +402,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 2); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -377,7 +416,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/0").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -391,7 +430,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/1").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -405,7 +444,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/2").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0/1").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0/1").unwrap()) .unwrap(); assert!( @@ -419,4 +458,35 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); } + + #[test] + fn test_tree_balancing_automatic() { + let seed_holder = seed_holder_for_tests(); + + let mut tree = KeyTreePublic::new(&seed_holder); + + for _ in 0..19 { + tree.generate_new_node_capped().unwrap(); + } + + let next_suitable_parent = tree.search_new_parent_capped().unwrap(); + + assert_eq!(next_suitable_parent, ChainIndex::from_str("/0").unwrap()); + + for _ in 0..18 { + tree.generate_new_node_capped().unwrap(); + } + + let next_suitable_parent = tree.search_new_parent_capped().unwrap(); + + assert_eq!(next_suitable_parent, ChainIndex::from_str("/1").unwrap()); + + for _ in 0..17 { + tree.generate_new_node_capped().unwrap(); + } + + let next_suitable_parent = tree.search_new_parent_capped().unwrap(); + + assert_eq!(next_suitable_parent, ChainIndex::from_str("/2").unwrap()); + } } diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index ac5ee48..41a686b 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -89,9 +89,18 @@ impl NSSAUserData { /// Returns the account_id of new account pub fn generate_new_public_transaction_private_key( &mut self, - parent_cci: ChainIndex, - ) -> nssa::AccountId { - self.public_key_tree.generate_new_node(parent_cci).unwrap() + parent_cci: Option, + ) -> (nssa::AccountId, ChainIndex) { + match parent_cci { + Some(parent_cci) => self + .public_key_tree + .generate_new_node(&parent_cci) + .expect("Parent must be present in a tree"), + None => self + .public_key_tree + .generate_new_node_capped() + .expect("No slots left"), + } } /// Returns the signing key for public transaction signatures @@ -113,9 +122,18 @@ impl NSSAUserData { /// Returns the account_id of new account pub fn generate_new_privacy_preserving_transaction_key_chain( &mut self, - parent_cci: ChainIndex, - ) -> nssa::AccountId { - self.private_key_tree.generate_new_node(parent_cci).unwrap() + parent_cci: Option, + ) -> (nssa::AccountId, ChainIndex) { + match parent_cci { + Some(parent_cci) => self + .private_key_tree + .generate_new_node(&parent_cci) + .expect("Parent must be present in a tree"), + None => self + .private_key_tree + .generate_new_node_capped() + .expect("No slots left"), + } } /// Returns the signing key for public transaction signatures @@ -169,10 +187,9 @@ mod tests { fn test_new_account() { let mut user_data = NSSAUserData::default(); - let account_id_pub = - user_data.generate_new_public_transaction_private_key(ChainIndex::root()); - let account_id_private = - user_data.generate_new_privacy_preserving_transaction_key_chain(ChainIndex::root()); + let (account_id_pub, _) = user_data.generate_new_public_transaction_private_key(None); + let (account_id_private, _) = + user_data.generate_new_privacy_preserving_transaction_key_chain(None); let is_private_key_generated = user_data .get_pub_account_signing_key(&account_id_pub) diff --git a/wallet/src/cli/account.rs b/wallet/src/cli/account.rs index 5b23b2b..b3c8d5c 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -98,13 +98,13 @@ pub enum NewSubcommand { Public { #[arg(long)] /// Chain index of a parent node - cci: ChainIndex, + cci: Option, }, /// Register new private account Private { #[arg(long)] /// Chain index of a parent node - cci: ChainIndex, + cci: Option, }, } @@ -115,9 +115,11 @@ impl WalletSubcommand for NewSubcommand { ) -> Result { match self { NewSubcommand::Public { cci } => { - let account_id = wallet_core.create_new_account_public(cci); + let (account_id, chain_index) = wallet_core.create_new_account_public(cci); - println!("Generated new account with account_id Public/{account_id}"); + println!( + "Generated new account with account_id Public/{account_id} at path {chain_index}" + ); let path = wallet_core.store_persistent_data().await?; @@ -126,7 +128,7 @@ impl WalletSubcommand for NewSubcommand { Ok(SubcommandReturnValue::RegisterAccount { account_id }) } NewSubcommand::Private { cci } => { - let account_id = wallet_core.create_new_account_private(cci); + let (account_id, chain_index) = wallet_core.create_new_account_private(cci); let (key, _) = wallet_core .storage @@ -135,7 +137,7 @@ impl WalletSubcommand for NewSubcommand { .unwrap(); println!( - "Generated new account with account_id Private/{}", + "Generated new account with account_id Private/{} at path {chain_index}", account_id.to_bytes().to_base58() ); println!("With npk {}", hex::encode(key.nullifer_public_key.0)); diff --git a/wallet/src/cli/config.rs b/wallet/src/cli/config.rs index 68670af..3026c29 100644 --- a/wallet/src/cli/config.rs +++ b/wallet/src/cli/config.rs @@ -9,10 +9,6 @@ use crate::{ /// Represents generic config CLI subcommand #[derive(Subcommand, Debug, Clone)] pub enum ConfigSubcommand { - /// Command to explicitly setup config and storage - /// - /// Does nothing in case if both already present - Setup {}, /// Getter of config fields Get { key: String }, /// Setter of config fields @@ -27,11 +23,6 @@ impl WalletSubcommand for ConfigSubcommand { wallet_core: &mut WalletCore, ) -> Result { match self { - ConfigSubcommand::Setup {} => { - let path = wallet_core.store_persistent_data().await?; - - println!("Stored persistent accounts at {path:#?}"); - } ConfigSubcommand::Get { key } => match key.as_str() { "all" => { let config_str = diff --git a/wallet/src/cli/mod.rs b/wallet/src/cli/mod.rs index 1f0e53b..e9f76bd 100644 --- a/wallet/src/cli/mod.rs +++ b/wallet/src/cli/mod.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::{io::Write, sync::Arc}; use anyhow::Result; use clap::{Parser, Subcommand}; @@ -16,7 +16,7 @@ use crate::{ token::TokenProgramAgnosticSubcommand, }, }, - helperfunctions::{fetch_config, parse_block_range}, + helperfunctions::{fetch_config, fetch_persistent_storage, parse_block_range}, }; pub mod account; @@ -54,25 +54,12 @@ pub enum Command { /// Command to setup config, get and set config fields #[command(subcommand)] Config(ConfigSubcommand), -} - -/// Represents overarching CLI command for a wallet with setup included -#[derive(Debug, Subcommand, Clone)] -#[clap(about)] -pub enum OverCommand { - /// Represents CLI command for a wallet - #[command(subcommand)] - Command(Command), - /// Setup of a storage. Initializes rots for public and private trees from `password`. - Setup { - #[arg(short, long)] - password: String, - }, + /// Restoring keys from given password at given `depth` + /// /// !!!WARNING!!! will rewrite current storage RestoreKeys { #[arg(short, long)] - password: String, - #[arg(short, long)] + /// Indicates, how deep in tree accounts may be. Affects command complexity. depth: u32, }, } @@ -91,7 +78,7 @@ pub struct Args { pub continuous_run: bool, /// Wallet command #[command(subcommand)] - pub command: Option, + pub command: Option, } #[derive(Debug, Clone)] @@ -104,6 +91,13 @@ pub enum SubcommandReturnValue { } pub async fn execute_subcommand(command: Command) -> Result { + if fetch_persistent_storage().await.is_err() { + println!("Persistent storage not found, need to execute setup"); + + let password = read_password_from_stdin()?; + execute_setup(password).await?; + } + let wallet_config = fetch_config().await?; let mut wallet_core = WalletCore::start_from_config_update_chain(wallet_config).await?; @@ -164,6 +158,12 @@ pub async fn execute_subcommand(command: Command) -> Result { + let password = read_password_from_stdin()?; + execute_keys_restoration(password, depth).await?; + + SubcommandReturnValue::Empty + } }; Ok(subcommand_ret) @@ -197,6 +197,16 @@ pub async fn execute_continuous_run() -> Result<()> { } } +pub fn read_password_from_stdin() -> Result { + let mut password = String::new(); + + print!("Input password: "); + std::io::stdout().flush()?; + std::io::stdin().read_line(&mut password)?; + + Ok(password.trim().to_string()) +} + pub async fn execute_setup(password: String) -> Result<()> { let config = fetch_config().await?; let wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password).await?; diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index f79d947..b2797ed 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -108,13 +108,19 @@ impl WalletCore { Ok(config_path) } - pub fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId { + pub fn create_new_account_public( + &mut self, + chain_index: Option, + ) -> (AccountId, ChainIndex) { self.storage .user_data .generate_new_public_transaction_private_key(chain_index) } - pub fn create_new_account_private(&mut self, chain_index: ChainIndex) -> AccountId { + pub fn create_new_account_private( + &mut self, + chain_index: Option, + ) -> (AccountId, ChainIndex) { self.storage .user_data .generate_new_privacy_preserving_transaction_key_chain(chain_index) diff --git a/wallet/src/main.rs b/wallet/src/main.rs index 00aa6d0..3fcac30 100644 --- a/wallet/src/main.rs +++ b/wallet/src/main.rs @@ -1,10 +1,7 @@ use anyhow::Result; use clap::{CommandFactory as _, Parser as _}; use tokio::runtime::Builder; -use wallet::cli::{ - Args, OverCommand, execute_continuous_run, execute_keys_restoration, execute_setup, - execute_subcommand, -}; +use wallet::cli::{Args, execute_continuous_run, execute_subcommand}; pub const NUM_THREADS: usize = 2; @@ -26,17 +23,9 @@ fn main() -> Result<()> { env_logger::init(); runtime.block_on(async move { - if let Some(over_command) = args.command { - match over_command { - OverCommand::Command(command) => { - let _output = execute_subcommand(command).await?; - Ok(()) - } - OverCommand::RestoreKeys { password, depth } => { - execute_keys_restoration(password, depth).await - } - OverCommand::Setup { password } => execute_setup(password).await, - } + if let Some(command) = args.command { + let _output = execute_subcommand(command).await?; + Ok(()) } else if args.continuous_run { execute_continuous_run().await } else {