feat: UX improvements

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 {
-            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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         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<String, TestFunction> {
         );
         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<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 {
-            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<String, TestFunction> {
         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<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();
@@ -1688,7 +1651,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();
@@ -1726,7 +1689,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();
@@ -1738,7 +1701,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/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<N: KeyNode> {
     pub key_map: BTreeMap<ChainIndex, N>,
@@ -101,10 +103,13 @@ impl<N: KeyNode> KeyTree<N> {
         }
     }
 
-    pub fn generate_new_node(&mut self, parent_cci: ChainIndex) -> Option<nssa::AccountId> {
-        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<N: KeyNode> KeyTree<N> {
         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<ChainIndex> {
+        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<N: KeyNode> KeyTree<N> {
         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());
+    }
 }

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<ChainIndex>,
+    ) -> (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<ChainIndex>,
+    ) -> (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)

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<ChainIndex>,
     },
     /// Register new private account
     Private {
         #[arg(long)]
         /// Chain index of a parent node
-        cci: ChainIndex,
+        cci: Option<ChainIndex>,
     },
 }
 
@@ -115,9 +115,11 @@ impl WalletSubcommand for NewSubcommand {
     ) -> Result<SubcommandReturnValue> {
         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));

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,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<OverCommand>,
+    pub command: Option<Command>,
 }
 
 #[derive(Debug, Clone)]
@@ -104,6 +91,13 @@ pub enum SubcommandReturnValue {
 }
 
 pub async fn execute_subcommand(command: Command) -> 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(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<SubcommandReturnValu
                 .handle_subcommand(&mut wallet_core)
                 .await?
         }
+        Command::RestoreKeys { depth } => {
+            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<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<()> {
     let config = fetch_config().await?;
     let wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password).await?;

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<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::{
-    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 {