Merge pull request #215 from logos-blockchain/Pravdyvy/tree-ux-updates

KeyTree UX updates

integration_tests/src/test_suite_map.rs

@@ -87,9 +87,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
     #[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 {
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None }));
-            cci: ChainIndex::root(),
-        }));
 
         let wallet_config = fetch_config().await.unwrap();
 
@@ -293,9 +291,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -306,9 +302,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -319,9 +313,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -456,9 +448,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -469,9 +459,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {
+            NewSubcommand::Private { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -482,9 +470,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {
+            NewSubcommand::Private { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -618,7 +604,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             account_id: definition_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
             NewSubcommand::Private {
-                cci: ChainIndex::root(),
+                cci: Some(ChainIndex::root()),
             },
         )))
         .await
@@ -631,7 +617,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             account_id: supply_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
             NewSubcommand::Public {
-                cci: ChainIndex::root(),
+                cci: Some(ChainIndex::root()),
             },
         )))
         .await
@@ -680,8 +666,8 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         // The data of a token definition account has the following layout:
         // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ]
         assert_eq!(
-            supply_acc.data,
+            supply_acc.data.as_ref(),
-            vec![
+            &[
                 1, 128, 101, 5, 31, 43, 36, 97, 108, 164, 92, 25, 157, 173, 5, 14, 194, 121, 239,
                 84, 19, 160, 243, 47, 193, 2, 250, 247, 232, 253, 191, 232, 173, 37, 0, 0, 0, 0, 0,
                 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
@@ -703,7 +689,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             account_id: definition_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
             NewSubcommand::Private {
-                cci: ChainIndex::root(),
+                cci: Some(ChainIndex::root()),
             },
         )))
         .await
@@ -716,7 +702,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             account_id: supply_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
             NewSubcommand::Private {
-                cci: ChainIndex::root(),
+                cci: Some(ChainIndex::root()),
             },
         )))
         .await
@@ -770,8 +756,8 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         // The data of a token definition account has the following layout:
         // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ]
         assert_eq!(
-            definition_acc.data,
+            definition_acc.data.as_ref(),
-            vec![
+            &[
                 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
             ]
         );
@@ -780,8 +766,8 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         // The data of a token definition account has the following layout:
         // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ]
         assert_eq!(
-            supply_acc.data,
+            supply_acc.data.as_ref(),
-            vec![
+            &[
                 1, 128, 101, 5, 31, 43, 36, 97, 108, 164, 92, 25, 157, 173, 5, 14, 194, 121, 239,
                 84, 19, 160, 243, 47, 193, 2, 250, 247, 232, 253, 191, 232, 173, 37, 0, 0, 0, 0, 0,
                 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
@@ -800,9 +786,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -813,9 +797,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {
+            NewSubcommand::Private { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -826,9 +808,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {
+            NewSubcommand::Private { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -942,9 +922,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -955,9 +933,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -968,9 +944,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {
+            NewSubcommand::Private { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -1084,9 +1058,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -1097,9 +1069,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {
+            NewSubcommand::Private { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -1110,9 +1080,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {
+            NewSubcommand::Public { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -1313,9 +1281,8 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         );
         let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap();
 
-        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+        let command =
-            cci: ChainIndex::root(),
+            Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None }));
-        }));
 
         let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap();
         let SubcommandReturnValue::RegisterAccount {
@@ -1676,9 +1643,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
     #[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 {
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None }));
-            cci: ChainIndex::root(),
-        }));
         let SubcommandReturnValue::RegisterAccount { account_id } =
             wallet::cli::execute_subcommand(command).await.unwrap()
         else {
@@ -1776,9 +1741,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: winner_account_id,
         } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {
+            NewSubcommand::Private { cci: None },
-                cci: ChainIndex::root(),
-            },
         )))
         .await
         .unwrap()
@@ -1862,7 +1825,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         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();
@@ -1874,7 +1837,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         };
 
         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();
@@ -1912,7 +1875,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         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();
@@ -1924,7 +1887,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         };
 
         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();

key_protocol/src/key_management/key_tree/chain_index.rs

@@ -138,6 +138,22 @@ impl ChainIndex {
 
         cumulative_stack.into_iter().unique()
     }
+
+    pub fn chain_ids_at_depth_rev(depth: usize) -> impl Iterator<Item = ChainIndex> {
+        let mut stack = vec![ChainIndex(vec![0; depth])];
+        let mut cumulative_stack = vec![ChainIndex(vec![0; depth])];
+
+        while let Some(id) = stack.pop() {
+            if let Some(collapsed_id) = id.collapse_back() {
+                for id in collapsed_id.shuffle_iter() {
+                    stack.push(id.clone());
+                    cumulative_stack.push(id);
+                }
+            }
+        }
+
+        cumulative_stack.into_iter().rev().unique()
+    }
 }
 
 #[cfg(test)]

key_protocol/src/key_management/key_tree/mod.rs

@@ -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<N: KeyNode> {
     pub key_map: BTreeMap<ChainIndex, N>,
@@ -120,6 +122,36 @@ impl<N: KeyNode> KeyTree<N> {
         Some((account_id, next_cci))
     }
 
+    fn find_next_slot_layered(&self) -> ChainIndex {
+        let mut depth = 1;
+
+        'outer: loop {
+            for chain_id in ChainIndex::chain_ids_at_depth_rev(depth) {
+                if !self.key_map.contains_key(&chain_id) {
+                    break 'outer chain_id;
+                }
+            }
+            depth += 1;
+        }
+    }
+
+    pub fn fill_node(&mut self, chain_index: &ChainIndex) -> Option<(nssa::AccountId, ChainIndex)> {
+        let parent_keys = self.key_map.get(&chain_index.parent()?)?;
+        let child_id = *chain_index.chain().last()?;
+
+        let child_keys = parent_keys.nth_child(child_id);
+        let account_id = child_keys.account_id();
+
+        self.key_map.insert(chain_index.clone(), child_keys);
+        self.account_id_map.insert(account_id, chain_index.clone());
+
+        Some((account_id, chain_index.clone()))
+    }
+
+    pub fn generate_new_node_layered(&mut self) -> Option<(nssa::AccountId, ChainIndex)> {
+        self.fill_node(&self.find_next_slot_layered())
+    }
+
     pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> {
         self.account_id_map
             .get(&account_id)
@@ -482,6 +514,21 @@ 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..100 {
+            tree.generate_new_node_layered().unwrap();
+        }
+
+        let next_slot = tree.find_next_slot_layered();
+
+        assert_eq!(next_slot, ChainIndex::from_str("/0/0/2/1").unwrap());
+    }
+
     #[test]
     fn test_cleanup() {
         let seed_holder = seed_holder_for_tests();

key_protocol/src/key_protocol_core/mod.rs

@@ -89,12 +89,18 @@ impl NSSAUserData {
     /// Returns the account_id of new account
     pub fn generate_new_public_transaction_private_key(
         &mut self,
-        parent_cci: ChainIndex,
+        parent_cci: Option<ChainIndex>,
-    ) -> nssa::AccountId {
+    ) -> (nssa::AccountId, ChainIndex) {
-        self.public_key_tree
+        match parent_cci {
-            .generate_new_node(&parent_cci)
+            Some(parent_cci) => self
-            .unwrap()
+                .public_key_tree
-            .0
+                .generate_new_node(&parent_cci)
+                .expect("Parent must be present in a tree"),
+            None => self
+                .public_key_tree
+                .generate_new_node_layered()
+                .expect("Search for new node slot failed"),
+        }
     }
 
     /// Returns the signing key for public transaction signatures
@@ -116,12 +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,
+        parent_cci: Option<ChainIndex>,
-    ) -> nssa::AccountId {
+    ) -> (nssa::AccountId, ChainIndex) {
-        self.private_key_tree
+        match parent_cci {
-            .generate_new_node(&parent_cci)
+            Some(parent_cci) => self
-            .unwrap()
+                .private_key_tree
-            .0
+                .generate_new_node(&parent_cci)
+                .expect("Parent must be present in a tree"),
+            None => self
+                .private_key_tree
+                .generate_new_node_layered()
+                .expect("Search for new node slot failed"),
+        }
     }
 
     /// Returns the signing key for public transaction signatures
@@ -175,8 +187,8 @@ mod tests {
     fn test_new_account() {
         let mut user_data = NSSAUserData::default();
 
-        let account_id_private =
+        let (account_id_private, _) = user_data
-            user_data.generate_new_privacy_preserving_transaction_key_chain(ChainIndex::root());
+            .generate_new_privacy_preserving_transaction_key_chain(Some(ChainIndex::root()));
 
         let is_key_chain_generated = user_data.get_private_account(&account_id_private).is_some();
 

wallet/src/cli/account.rs

@@ -96,13 +96,13 @@ pub enum NewSubcommand {
     Public {
         #[arg(long)]
         /// Chain index of a parent node
-        cci: ChainIndex,
+        cci: Option<ChainIndex>,
     },
     /// Register new private account
     Private {
         #[arg(long)]
         /// Chain index of a parent node
-        cci: ChainIndex,
+        cci: Option<ChainIndex>,
     },
 }
 
@@ -113,9 +113,11 @@ impl WalletSubcommand for NewSubcommand {
     ) -> Result<SubcommandReturnValue> {
         match self {
             NewSubcommand::Public { cci } => {
-                let account_id = wallet_core.create_new_account_public(cci).await;
+                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?;
 
@@ -124,7 +126,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
@@ -133,7 +135,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));

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<SubcommandReturnValue> {
         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 =

wallet/src/cli/mod.rs

@@ -1,3 +1,5 @@
+use std::io::Write;
+
 use anyhow::Result;
 use clap::{Parser, Subcommand};
 use nssa::program::Program;
@@ -13,7 +15,7 @@ use crate::{
             token::TokenProgramAgnosticSubcommand,
         },
     },
-    helperfunctions::{fetch_config, merge_auth_config},
+    helperfunctions::{fetch_config, fetch_persistent_storage, merge_auth_config},
 };
 
 pub mod account;
@@ -51,25 +53,12 @@ pub enum Command {
     /// Command to setup config, get and set config fields
     #[command(subcommand)]
     Config(ConfigSubcommand),
-}
+    /// Restoring keys from given password at given `depth`
-
+    ///
-/// 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,
-    },
     /// !!!WARNING!!! will rewrite current storage
     RestoreKeys {
         #[arg(short, long)]
-        password: String,
+        /// Indicates, how deep in tree accounts may be. Affects command complexity.
-        #[arg(short, long)]
         depth: u32,
     },
 }
@@ -91,7 +80,7 @@ pub struct Args {
     pub auth: Option<String>,
     /// Wallet command
     #[command(subcommand)]
-    pub command: Option<OverCommand>,
+    pub command: Option<Command>,
 }
 
 #[derive(Debug, Clone)]
@@ -111,8 +100,15 @@ pub async fn execute_subcommand_with_auth(
     command: Command,
     auth: Option<String>,
 ) -> Result<SubcommandReturnValue> {
+    if fetch_persistent_storage().await.is_err() {
+        println!("Persistent storage not found, need to execute setup");
+
+        let password = read_password_from_stdin()?;
+        execute_setup_with_auth(password, auth.clone()).await?;
+    }
+
     let wallet_config = fetch_config().await?;
-    let wallet_config = merge_auth_config(wallet_config, auth)?;
+    let wallet_config = merge_auth_config(wallet_config, auth.clone())?;
     let mut wallet_core = WalletCore::start_from_config_update_chain(wallet_config).await?;
 
     let subcommand_ret = match command {
@@ -172,6 +168,12 @@ pub async fn execute_subcommand_with_auth(
                 .handle_subcommand(&mut wallet_core)
                 .await?
         }
+        Command::RestoreKeys { depth } => {
+            let password = read_password_from_stdin()?;
+            execute_keys_restoration_with_auth(password, depth, auth).await?;
+
+            SubcommandReturnValue::Empty
+        }
     };
 
     Ok(subcommand_ret)
@@ -200,6 +202,16 @@ pub async fn execute_continuous_run_with_auth(auth: Option<String>) -> Result<()
     }
 }
 
+pub fn read_password_from_stdin() -> Result<String> {
+    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<()> {
     execute_setup_with_auth(password, None).await
 }

wallet/src/lib.rs

@@ -126,13 +126,19 @@ impl WalletCore {
         Ok(config_path)
     }
 
-    pub async fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId {
+    pub fn create_new_account_public(
+        &mut self,
+        chain_index: Option<ChainIndex>,
+    ) -> (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<ChainIndex>,
+    ) -> (AccountId, ChainIndex) {
         self.storage
             .user_data
             .generate_new_privacy_preserving_transaction_key_chain(chain_index)

wallet/src/main.rs

@@ -1,10 +1,7 @@
 use anyhow::Result;
 use clap::{CommandFactory as _, Parser as _};
 use tokio::runtime::Builder;
-use wallet::cli::{
+use wallet::cli::{Args, execute_continuous_run_with_auth, execute_subcommand_with_auth};
-    Args, OverCommand, execute_continuous_run_with_auth, execute_keys_restoration_with_auth,
-    execute_setup_with_auth, execute_subcommand_with_auth,
-};
 
 pub const NUM_THREADS: usize = 2;
 
@@ -25,19 +22,9 @@ fn main() -> Result<()> {
     env_logger::init();
 
     runtime.block_on(async move {
-        if let Some(over_command) = args.command {
+        if let Some(command) = args.command {
-            match over_command {
+            let _output = execute_subcommand_with_auth(command, args.auth).await?;
-                OverCommand::Command(command) => {
+            Ok(())
-                    let _output = execute_subcommand_with_auth(command, args.auth).await?;
-                    Ok(())
-                }
-                OverCommand::RestoreKeys { password, depth } => {
-                    execute_keys_restoration_with_auth(password, depth, args.auth).await
-                }
-                OverCommand::Setup { password } => {
-                    execute_setup_with_auth(password, args.auth).await
-                }
-            }
         } else if args.continuous_run {
             execute_continuous_run_with_auth(args.auth).await
         } else {