Merge pull request #153 from logos-blockchain/Pravdyvy/keys-restoration-from-mnemonic

Keys restoration from mnemonic

integration_tests/src/test_suite_map.rs

@@ -2,6 +2,7 @@ use std::{
     collections::HashMap,
     path::PathBuf,
     pin::Pin,
+    str::FromStr,
     time::{Duration, Instant},
 };
 
@@ -1672,6 +1673,217 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         info!("Success!");
     }
 
+    #[nssa_integration_test]
+    pub async fn test_keys_restoration() {
+        info!("########## test_keys_restoration ##########");
+        let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap();
+
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+            cci: ChainIndex::root(),
+        }));
+
+        let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount {
+            account_id: to_account_id1,
+        } = sub_ret
+        else {
+            panic!("FAILED TO REGISTER ACCOUNT");
+        };
+
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+            cci: ChainIndex::from_str("/0").unwrap(),
+        }));
+
+        let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount {
+            account_id: to_account_id2,
+        } = sub_ret
+        else {
+            panic!("FAILED TO REGISTER ACCOUNT");
+        };
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_private_account_input_from_str(&from.to_string()),
+            to: Some(make_private_account_input_from_str(
+                &to_account_id1.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 100,
+        });
+
+        wallet::cli::execute_subcommand(command).await.unwrap();
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_private_account_input_from_str(&from.to_string()),
+            to: Some(make_private_account_input_from_str(
+                &to_account_id2.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 101,
+        });
+
+        wallet::cli::execute_subcommand(command).await.unwrap();
+
+        let from: AccountId = ACC_SENDER.parse().unwrap();
+
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
+            cci: ChainIndex::root(),
+        }));
+
+        let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount {
+            account_id: to_account_id3,
+        } = sub_ret
+        else {
+            panic!("FAILED TO REGISTER ACCOUNT");
+        };
+
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
+            cci: ChainIndex::from_str("/0").unwrap(),
+        }));
+
+        let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount {
+            account_id: to_account_id4,
+        } = sub_ret
+        else {
+            panic!("FAILED TO REGISTER ACCOUNT");
+        };
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_public_account_input_from_str(&from.to_string()),
+            to: Some(make_public_account_input_from_str(
+                &to_account_id3.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 102,
+        });
+
+        wallet::cli::execute_subcommand(command).await.unwrap();
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_public_account_input_from_str(&from.to_string()),
+            to: Some(make_public_account_input_from_str(
+                &to_account_id4.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 103,
+        });
+
+        wallet::cli::execute_subcommand(command).await.unwrap();
+
+        info!("########## PREPARATION END ##########");
+
+        wallet::cli::execute_keys_restoration("test_pass".to_string(), 10)
+            .await
+            .unwrap();
+
+        let wallet_config = fetch_config().await.unwrap();
+        let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config.clone())
+            .await
+            .unwrap();
+
+        let acc1 = wallet_storage
+            .storage
+            .user_data
+            .private_key_tree
+            .get_node(to_account_id1)
+            .expect("Acc 1 should be restored");
+
+        let acc2 = wallet_storage
+            .storage
+            .user_data
+            .private_key_tree
+            .get_node(to_account_id2)
+            .expect("Acc 2 should be restored");
+
+        let _ = wallet_storage
+            .storage
+            .user_data
+            .public_key_tree
+            .get_node(to_account_id3)
+            .expect("Acc 3 should be restored");
+
+        let _ = wallet_storage
+            .storage
+            .user_data
+            .public_key_tree
+            .get_node(to_account_id4)
+            .expect("Acc 4 should be restored");
+
+        assert_eq!(
+            acc1.value.1.program_owner,
+            Program::authenticated_transfer_program().id()
+        );
+        assert_eq!(
+            acc2.value.1.program_owner,
+            Program::authenticated_transfer_program().id()
+        );
+
+        assert_eq!(acc1.value.1.balance, 100);
+        assert_eq!(acc2.value.1.balance, 101);
+
+        info!("########## TREE CHECKS END ##########");
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_private_account_input_from_str(&to_account_id1.to_string()),
+            to: Some(make_private_account_input_from_str(
+                &to_account_id2.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 10,
+        });
+
+        wallet::cli::execute_subcommand(command).await.unwrap();
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_public_account_input_from_str(&to_account_id3.to_string()),
+            to: Some(make_public_account_input_from_str(
+                &to_account_id4.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 11,
+        });
+
+        wallet::cli::execute_subcommand(command).await.unwrap();
+
+        let wallet_config = fetch_config().await.unwrap();
+        let seq_client = SequencerClient::new(wallet_config.sequencer_addr.clone()).unwrap();
+        let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config.clone())
+            .await
+            .unwrap();
+
+        let comm1 = wallet_storage
+            .get_private_account_commitment(&to_account_id1)
+            .expect("Acc 1 commitment should exist");
+        let comm2 = wallet_storage
+            .get_private_account_commitment(&to_account_id2)
+            .expect("Acc 2 commitment should exist");
+
+        assert!(verify_commitment_is_in_state(comm1, &seq_client).await);
+        assert!(verify_commitment_is_in_state(comm2, &seq_client).await);
+
+        let acc3 = seq_client
+            .get_account_balance(to_account_id3.to_string())
+            .await
+            .expect("Acc 3 must be present in public state");
+        let acc4 = seq_client
+            .get_account_balance(to_account_id4.to_string())
+            .await
+            .expect("Acc 4 must be present in public state");
+
+        assert_eq!(acc3.balance, 91);
+        assert_eq!(acc4.balance, 114);
+
+        info!("Success!");
+    }
+
     println!("{function_map:#?}");
 
     function_map

key_protocol/Cargo.toml

@@ -16,6 +16,7 @@ bip39.workspace = true
 hmac-sha512.workspace = true
 thiserror.workspace = true
 nssa-core = { path = "../nssa/core", features = ["host"] }
+itertools.workspace = true
 
 [dependencies.common]
 path = "../common"

key_protocol/src/key_management/key_tree/chain_index.rs

@@ -1,8 +1,9 @@
 use std::{fmt::Display, str::FromStr};
 
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, Hash)]
 pub struct ChainIndex(Vec<u32>);
 
 #[derive(thiserror::Error, Debug)]
@@ -77,12 +78,66 @@ impl ChainIndex {
         ChainIndex(chain)
     }
 
+    pub fn previous_in_line(&self) -> Option<ChainIndex> {
+        let mut chain = self.0.clone();
+        if let Some(last_p) = chain.last_mut() {
+            *last_p = last_p.checked_sub(1)?;
+        }
+
+        Some(ChainIndex(chain))
+    }
+
+    pub fn parent(&self) -> Option<ChainIndex> {
+        if self.0.is_empty() {
+            None
+        } else {
+            Some(ChainIndex(self.0[..(self.0.len() - 1)].to_vec()))
+        }
+    }
+
     pub fn nth_child(&self, child_id: u32) -> ChainIndex {
         let mut chain = self.0.clone();
         chain.push(child_id);
 
         ChainIndex(chain)
     }
+
+    pub fn depth(&self) -> u32 {
+        self.0.iter().map(|cci| cci + 1).sum()
+    }
+
+    fn collapse_back(&self) -> Option<Self> {
+        let mut res = self.parent()?;
+
+        let last_mut = res.0.last_mut()?;
+        *last_mut += *(self.0.last()?) + 1;
+
+        Some(res)
+    }
+
+    fn shuffle_iter(&self) -> impl Iterator<Item = ChainIndex> {
+        self.0
+            .iter()
+            .permutations(self.0.len())
+            .unique()
+            .map(|item| ChainIndex(item.into_iter().cloned().collect()))
+    }
+
+    pub fn chain_ids_at_depth(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().unique()
+    }
 }
 
 #[cfg(test)]
@@ -145,4 +200,83 @@ mod tests {
 
         assert_eq!(string_index, "/5/7/8".to_string());
     }
+
+    #[test]
+    fn test_prev_in_line() {
+        let chain_id = ChainIndex(vec![1, 7, 3]);
+
+        let prev_chain_id = chain_id.previous_in_line().unwrap();
+
+        assert_eq!(prev_chain_id, ChainIndex(vec![1, 7, 2]))
+    }
+
+    #[test]
+    fn test_prev_in_line_no_prev() {
+        let chain_id = ChainIndex(vec![1, 7, 0]);
+
+        let prev_chain_id = chain_id.previous_in_line();
+
+        assert_eq!(prev_chain_id, None)
+    }
+
+    #[test]
+    fn test_parent() {
+        let chain_id = ChainIndex(vec![1, 7, 3]);
+
+        let parent_chain_id = chain_id.parent().unwrap();
+
+        assert_eq!(parent_chain_id, ChainIndex(vec![1, 7]))
+    }
+
+    #[test]
+    fn test_parent_no_parent() {
+        let chain_id = ChainIndex(vec![]);
+
+        let parent_chain_id = chain_id.parent();
+
+        assert_eq!(parent_chain_id, None)
+    }
+
+    #[test]
+    fn test_parent_root() {
+        let chain_id = ChainIndex(vec![1]);
+
+        let parent_chain_id = chain_id.parent().unwrap();
+
+        assert_eq!(parent_chain_id, ChainIndex::root())
+    }
+
+    #[test]
+    fn test_collapse_back() {
+        let chain_id = ChainIndex(vec![1, 1]);
+
+        let collapsed = chain_id.collapse_back().unwrap();
+
+        assert_eq!(collapsed, ChainIndex(vec![3]))
+    }
+
+    #[test]
+    fn test_collapse_back_one() {
+        let chain_id = ChainIndex(vec![1]);
+
+        let collapsed = chain_id.collapse_back();
+
+        assert_eq!(collapsed, None)
+    }
+
+    #[test]
+    fn test_collapse_back_root() {
+        let chain_id = ChainIndex(vec![]);
+
+        let collapsed = chain_id.collapse_back();
+
+        assert_eq!(collapsed, None)
+    }
+
+    #[test]
+    fn test_shuffle() {
+        for id in ChainIndex::chain_ids_at_depth(5) {
+            println!("{id}");
+        }
+    }
 }

key_protocol/src/key_management/key_tree/mod.rs

@@ -1,5 +1,10 @@
-use std::collections::{BTreeMap, HashMap};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+};
 
+use anyhow::Result;
+use common::sequencer_client::SequencerClient;
 use serde::{Deserialize, Serialize};
 
 use crate::key_management::{
@@ -96,21 +101,23 @@ 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 parent_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);
 
-        let child_keys = father_keys.nth_child(next_child_id);
-
+        let child_keys = parent_keys.nth_child(next_child_id);
         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))
     }
 
     pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> {
@@ -129,11 +136,164 @@ impl<N: KeyNode> KeyTree<N> {
         self.account_id_map.insert(account_id, chain_index.clone());
         self.key_map.insert(chain_index, node);
     }
+
+    pub fn remove(&mut self, addr: nssa::AccountId) -> Option<N> {
+        let chain_index = self.account_id_map.remove(&addr).unwrap();
+        self.key_map.remove(&chain_index)
+    }
+
+    /// Populates tree with children.
+    ///
+    /// For given `depth` adds children to a tree such that their `ChainIndex::depth(&self) <
+    /// depth`.
+    ///
+    /// Tree must be empty before start
+    pub fn generate_tree_for_depth(&mut self, depth: u32) {
+        let mut id_stack = vec![ChainIndex::root()];
+
+        while let Some(curr_id) = id_stack.pop() {
+            let mut next_id = curr_id.nth_child(0);
+
+            while (next_id.depth()) < depth {
+                self.generate_new_node(&curr_id);
+                id_stack.push(next_id.clone());
+                next_id = next_id.next_in_line();
+            }
+        }
+    }
+}
+
+impl KeyTree<ChildKeysPrivate> {
+    /// Cleanup of all non-initialized accounts in a private tree
+    ///
+    /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) <
+    /// depth`.
+    ///
+    /// If account is default, removes them.
+    ///
+    /// Chain must be parsed for accounts beforehand
+    ///
+    /// Fast, leaves gaps between accounts
+    pub fn cleanup_tree_remove_uninit_for_depth(&mut self, depth: u32) {
+        let mut id_stack = vec![ChainIndex::root()];
+
+        while let Some(curr_id) = id_stack.pop() {
+            if let Some(node) = self.key_map.get(&curr_id)
+                && node.value.1 == nssa::Account::default()
+                && curr_id != ChainIndex::root()
+            {
+                let addr = node.account_id();
+                self.remove(addr);
+            }
+
+            let mut next_id = curr_id.nth_child(0);
+
+            while (next_id.depth()) < depth {
+                id_stack.push(next_id.clone());
+                next_id = next_id.next_in_line();
+            }
+        }
+    }
+
+    /// Cleanup of non-initialized accounts in a private tree
+    ///
+    /// If account is default, removes them, stops at first non-default account.
+    ///
+    /// Walks through tree in lairs of same depth using `ChainIndex::chain_ids_at_depth()`
+    ///
+    /// Chain must be parsed for accounts beforehand
+    ///
+    /// Slow, maintains tree consistency.
+    pub fn cleanup_tree_remove_uninit_layered(&mut self, depth: u32) {
+        'outer: for i in (1..(depth as usize)).rev() {
+            println!("Cleanup of tree at depth {i}");
+            for id in ChainIndex::chain_ids_at_depth(i) {
+                if let Some(node) = self.key_map.get(&id) {
+                    if node.value.1 == nssa::Account::default() {
+                        let addr = node.account_id();
+                        self.remove(addr);
+                    } else {
+                        break 'outer;
+                    }
+                }
+            }
+        }
+    }
+}
+
+impl KeyTree<ChildKeysPublic> {
+    /// Cleanup of all non-initialized accounts in a public tree
+    ///
+    /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) <
+    /// depth`.
+    ///
+    /// If account is default, removes them.
+    ///
+    /// Fast, leaves gaps between accounts
+    pub async fn cleanup_tree_remove_ininit_for_depth(
+        &mut self,
+        depth: u32,
+        client: Arc<SequencerClient>,
+    ) -> Result<()> {
+        let mut id_stack = vec![ChainIndex::root()];
+
+        while let Some(curr_id) = id_stack.pop() {
+            if let Some(node) = self.key_map.get(&curr_id) {
+                let address = node.account_id();
+                let node_acc = client.get_account(address.to_string()).await?.account;
+
+                if node_acc == nssa::Account::default() && curr_id != ChainIndex::root() {
+                    self.remove(address);
+                }
+            }
+
+            let mut next_id = curr_id.nth_child(0);
+
+            while (next_id.depth()) < depth {
+                id_stack.push(next_id.clone());
+                next_id = next_id.next_in_line();
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Cleanup of non-initialized accounts in a public tree
+    ///
+    /// If account is default, removes them, stops at first non-default account.
+    ///
+    /// Walks through tree in lairs of same depth using `ChainIndex::chain_ids_at_depth()`
+    ///
+    /// Slow, maintains tree consistency.
+    pub async fn cleanup_tree_remove_uninit_layered(
+        &mut self,
+        depth: u32,
+        client: Arc<SequencerClient>,
+    ) -> Result<()> {
+        'outer: for i in (1..(depth as usize)).rev() {
+            println!("Cleanup of tree at depth {i}");
+            for id in ChainIndex::chain_ids_at_depth(i) {
+                if let Some(node) = self.key_map.get(&id) {
+                    let address = node.account_id();
+                    let node_acc = client.get_account(address.to_string()).await?.account;
+
+                    if node_acc == nssa::Account::default() {
+                        let addr = node.account_id();
+                        self.remove(addr);
+                    } else {
+                        break 'outer;
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::str::FromStr;
+    use std::{collections::HashSet, str::FromStr};
 
     use nssa::AccountId;
 
@@ -162,7 +322,7 @@ mod tests {
     fn test_small_key_tree() {
         let seed_holder = seed_holder_for_tests();
 
-        let mut tree = KeyTreePublic::new(&seed_holder);
+        let mut tree = KeyTreePrivate::new(&seed_holder);
 
         let next_last_child_for_parent_id = tree
             .find_next_last_child_of_id(&ChainIndex::root())
@@ -170,7 +330,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
@@ -183,12 +343,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())
@@ -201,7 +361,7 @@ mod tests {
     fn test_key_tree_can_not_make_child_keys() {
         let seed_holder = seed_holder_for_tests();
 
-        let mut tree = KeyTreePublic::new(&seed_holder);
+        let mut tree = KeyTreePrivate::new(&seed_holder);
 
         let next_last_child_for_parent_id = tree
             .find_next_last_child_of_id(&ChainIndex::root())
@@ -209,7 +369,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
@@ -222,7 +382,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);
     }
@@ -239,7 +399,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
@@ -252,7 +412,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
@@ -265,7 +425,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
@@ -279,7 +439,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
@@ -293,7 +453,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
@@ -307,7 +467,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!(
@@ -321,4 +481,79 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 1);
     }
+
+    #[test]
+    fn test_cleanup() {
+        let seed_holder = seed_holder_for_tests();
+
+        let mut tree = KeyTreePrivate::new(&seed_holder);
+        tree.generate_tree_for_depth(10);
+
+        let acc = tree
+            .key_map
+            .get_mut(&ChainIndex::from_str("/1").unwrap())
+            .unwrap();
+        acc.value.1.balance = 2;
+
+        let acc = tree
+            .key_map
+            .get_mut(&ChainIndex::from_str("/2").unwrap())
+            .unwrap();
+        acc.value.1.balance = 3;
+
+        let acc = tree
+            .key_map
+            .get_mut(&ChainIndex::from_str("/0/1").unwrap())
+            .unwrap();
+        acc.value.1.balance = 5;
+
+        let acc = tree
+            .key_map
+            .get_mut(&ChainIndex::from_str("/1/0").unwrap())
+            .unwrap();
+        acc.value.1.balance = 6;
+
+        tree.cleanup_tree_remove_uninit_layered(10);
+
+        let mut key_set_res = HashSet::new();
+        key_set_res.insert("/0".to_string());
+        key_set_res.insert("/1".to_string());
+        key_set_res.insert("/2".to_string());
+        key_set_res.insert("/".to_string());
+        key_set_res.insert("/0/0".to_string());
+        key_set_res.insert("/0/1".to_string());
+        key_set_res.insert("/1/0".to_string());
+
+        let mut key_set = HashSet::new();
+
+        for key in tree.key_map.keys() {
+            key_set.insert(key.to_string());
+        }
+
+        assert_eq!(key_set, key_set_res);
+
+        let acc = tree
+            .key_map
+            .get(&ChainIndex::from_str("/1").unwrap())
+            .unwrap();
+        assert_eq!(acc.value.1.balance, 2);
+
+        let acc = tree
+            .key_map
+            .get(&ChainIndex::from_str("/2").unwrap())
+            .unwrap();
+        assert_eq!(acc.value.1.balance, 3);
+
+        let acc = tree
+            .key_map
+            .get(&ChainIndex::from_str("/0/1").unwrap())
+            .unwrap();
+        assert_eq!(acc.value.1.balance, 5);
+
+        let acc = tree
+            .key_map
+            .get(&ChainIndex::from_str("/1/0").unwrap())
+            .unwrap();
+        assert_eq!(acc.value.1.balance, 6);
+    }
 }

key_protocol/src/key_protocol_core/mod.rs

@@ -91,7 +91,10 @@ impl NSSAUserData {
         &mut self,
         parent_cci: ChainIndex,
     ) -> nssa::AccountId {
-        self.public_key_tree.generate_new_node(parent_cci).unwrap()
+        self.public_key_tree
+            .generate_new_node(&parent_cci)
+            .unwrap()
+            .0
     }
 
     /// Returns the signing key for public transaction signatures
@@ -115,7 +118,10 @@ impl NSSAUserData {
         &mut self,
         parent_cci: ChainIndex,
     ) -> nssa::AccountId {
-        self.private_key_tree.generate_new_node(parent_cci).unwrap()
+        self.private_key_tree
+            .generate_new_node(&parent_cci)
+            .unwrap()
+            .0
     }
 
     /// Returns the signing key for public transaction signatures
@@ -169,17 +175,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 is_private_key_generated = user_data
-            .get_pub_account_signing_key(&account_id_pub)
-            .is_some();
-
-        assert!(is_private_key_generated);
-
         let is_key_chain_generated = user_data.get_private_account(&account_id_private).is_some();
 
         assert!(is_key_chain_generated);

wallet/src/cli/account.rs

@@ -113,7 +113,7 @@ impl WalletSubcommand for NewSubcommand {
     ) -> Result<SubcommandReturnValue> {
         match self {
             NewSubcommand::Public { cci } => {
-                let account_id = wallet_core.create_new_account_public(cci);
+                let account_id = wallet_core.create_new_account_public(cci).await;
 
                 println!("Generated new account with account_id Public/{account_id}");
 

wallet/src/cli/mod.rs

@@ -65,6 +65,13 @@ pub enum OverCommand {
         #[arg(short, long)]
         password: String,
     },
+    /// !!!WARNING!!! will rewrite current storage
+    RestoreKeys {
+        #[arg(short, long)]
+        password: String,
+        #[arg(short, long)]
+        depth: u32,
+    },
 }
 
 /// To execute commands, env var NSSA_WALLET_HOME_DIR must be set into directory with config
@@ -206,3 +213,67 @@ pub async fn execute_setup_with_auth(password: String, auth: Option<String>) ->
 
     Ok(())
 }
+
+pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> {
+    execute_keys_restoration_with_auth(password, depth, None).await
+}
+
+pub async fn execute_keys_restoration_with_auth(
+    password: String,
+    depth: u32,
+    auth: Option<String>,
+) -> Result<()> {
+    let config = fetch_config().await?;
+    let config = merge_auth_config(config, auth)?;
+    let mut wallet_core =
+        WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?;
+
+    wallet_core
+        .storage
+        .user_data
+        .public_key_tree
+        .generate_tree_for_depth(depth);
+
+    println!("Public tree generated");
+
+    wallet_core
+        .storage
+        .user_data
+        .private_key_tree
+        .generate_tree_for_depth(depth);
+
+    println!("Private tree generated");
+
+    wallet_core
+        .storage
+        .user_data
+        .public_key_tree
+        .cleanup_tree_remove_uninit_layered(depth, wallet_core.sequencer_client.clone())
+        .await?;
+
+    println!("Public tree cleaned up");
+
+    let last_block = wallet_core
+        .sequencer_client
+        .get_last_block()
+        .await?
+        .last_block;
+
+    println!("Last block is {last_block}");
+
+    wallet_core.sync_to_block(last_block).await?;
+
+    println!("Private tree clean up start");
+
+    wallet_core
+        .storage
+        .user_data
+        .private_key_tree
+        .cleanup_tree_remove_uninit_layered(depth);
+
+    println!("Private tree cleaned up");
+
+    wallet_core.store_persistent_data().await?;
+
+    Ok(())
+}

wallet/src/lib.rs

@@ -126,7 +126,7 @@ impl WalletCore {
         Ok(config_path)
     }
 
-    pub fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId {
+    pub async fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId {
         self.storage
             .user_data
             .generate_new_public_transaction_private_key(chain_index)

wallet/src/main.rs

@@ -2,8 +2,8 @@ use anyhow::Result;
 use clap::{CommandFactory as _, Parser as _};
 use tokio::runtime::Builder;
 use wallet::cli::{
-    Args, OverCommand, execute_continuous_run_with_auth, execute_setup_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;
@@ -31,6 +31,9 @@ fn main() -> Result<()> {
                     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
                 }