feat: keys restoration from mnemonic

integration_tests/src/test_suite_map.rs

@@ -3,6 +3,7 @@ use std::{
     collections::HashMap,
     path::PathBuf,
     pin::Pin,
+    str::FromStr,
     time::{Duration, Instant},
 };
 
@@ -1647,6 +1648,136 @@ 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: Address = ACC_SENDER_PRIVATE.parse().unwrap();
+
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+            cci: ChainIndex::root(),
+        }));
+
+        let sub_ret = wallet::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount { addr: to_addr1 } = 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::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount { addr: to_addr2 } = 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_addr1.to_string())),
+            to_npk: None,
+            to_ipk: None,
+            amount: 100,
+        });
+
+        wallet::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_addr2.to_string())),
+            to_npk: None,
+            to_ipk: None,
+            amount: 100,
+        });
+
+        wallet::execute_subcommand(command).await.unwrap();
+
+        let from: Address = ACC_SENDER.parse().unwrap();
+
+        let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
+            cci: ChainIndex::root(),
+        }));
+
+        let sub_ret = wallet::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount { addr: to_addr3 } = 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::execute_subcommand(command).await.unwrap();
+        let SubcommandReturnValue::RegisterAccount { addr: to_addr4 } = 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_addr3.to_string())),
+            to_npk: None,
+            to_ipk: None,
+            amount: 100,
+        });
+
+        wallet::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_addr4.to_string())),
+            to_npk: None,
+            to_ipk: None,
+            amount: 100,
+        });
+
+        wallet::execute_subcommand(command).await.unwrap();
+
+        info!("########## PREPARATION END ##########");
+
+        wallet::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();
+
+        assert!(
+            wallet_storage
+                .storage
+                .user_data
+                .private_key_tree
+                .get_node(to_addr1)
+                .is_some()
+        );
+        assert!(
+            wallet_storage
+                .storage
+                .user_data
+                .private_key_tree
+                .get_node(to_addr2)
+                .is_some()
+        );
+        assert!(
+            wallet_storage
+                .storage
+                .user_data
+                .public_key_tree
+                .get_node(to_addr3)
+                .is_some()
+        );
+        assert!(
+            wallet_storage
+                .storage
+                .user_data
+                .public_key_tree
+                .get_node(to_addr4)
+                .is_some()
+        );
+
+        info!("Success!");
+    }
+
     println!("{function_map:#?}");
 
     function_map

key_protocol/src/key_management/key_tree/chain_index.rs

@@ -42,10 +42,13 @@ impl FromStr for ChainIndex {
 impl Display for ChainIndex {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "/")?;
-        for cci in &self.0[..(self.0.len() - 1)] {
-            write!(f, "{cci}/")?;
+        if *self != Self::root() {
+            for cci in &self.0[..(self.0.len() - 1)] {
+                write!(f, "{cci}/")?;
+            }
+            write!(f, "{}", self.0.last().unwrap())?;
         }
-        write!(f, "{}", self.0.last().unwrap())
+        Ok(())
     }
 }
 
@@ -74,6 +77,16 @@ impl ChainIndex {
 
         ChainIndex(chain)
     }
+
+    pub fn depth(&self) -> u32 {
+        let mut res = 0;
+
+        for cci in &self.0 {
+            res += cci + 1;
+        }
+
+        res
+    }
 }
 
 #[cfg(test)]

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::{
@@ -128,24 +133,80 @@ impl<Node: KeyNode> KeyTree<Node> {
         self.key_map.insert(chain_index, node);
     }
 
+    pub fn remove(&mut self, addr: nssa::Address) -> Option<Node> {
+        let chain_index = self.addr_map.remove(&addr).unwrap();
+        self.key_map.remove(&chain_index)
+    }
+
     pub fn generate_tree_for_depth(&mut self, depth: u32) {
         let mut id_stack = vec![ChainIndex::root()];
 
-        while !id_stack.is_empty() {
-            let curr_id = id_stack.pop().unwrap();
-
+        while let Some(curr_id) = id_stack.pop() {
             self.generate_new_node(curr_id.clone());
 
             let mut next_id = curr_id.n_th_child(0);
 
-            while (next_id.chain().iter().sum::<u32>()) < depth - 1 {
+            while (next_id.depth()) < depth - 1 {
                 id_stack.push(next_id.clone());
                 next_id = next_id.next_in_line();
-            } 
+            }
         }
     }
 }
 
+impl KeyTree<ChildKeysPrivate> {
+    pub fn cleanup_tree_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.address();
+                self.remove(addr);
+            }
+
+            let mut next_id = curr_id.n_th_child(0);
+
+            while (next_id.depth()) < depth - 1 {
+                id_stack.push(next_id.clone());
+                next_id = next_id.next_in_line();
+            }
+        }
+    }
+}
+
+impl KeyTree<ChildKeysPublic> {
+    pub async fn cleanup_tree_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.address();
+                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.n_th_child(0);
+
+            while (next_id.depth()) < depth - 1 {
+                id_stack.push(next_id.clone());
+                next_id = next_id.next_in_line();
+            }
+        }
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;

wallet/src/lib.rs

@@ -265,7 +265,7 @@ pub enum OverCommand {
         password: String,
         #[arg(short, long)]
         depth: u32,
-    }
+    },
 }
 
 ///To execute commands, env var NSSA_WALLET_HOME_DIR must be set into directory with config
@@ -521,10 +521,59 @@ pub async fn execute_setup(password: String) -> Result<()> {
 
 pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> {
     let config = fetch_config().await?;
-    let mut wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?;
+    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);
-    wallet_core.storage.user_data.private_key_tree.generate_tree_for_depth(depth); 
+    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_for_depth(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}");
+
+    parse_block_range(
+        1,
+        last_block,
+        wallet_core.sequencer_client.clone(),
+        &mut wallet_core,
+    )
+    .await?;
+
+    println!("Private tree clean up start");
+
+    wallet_core
+        .storage
+        .user_data
+        .private_key_tree
+        .cleanup_tree_for_depth(depth);
+
+    println!("Private tree cleaned up");
 
     wallet_core.store_persistent_data().await?;
 

wallet/src/main.rs

@@ -1,7 +1,10 @@
 use anyhow::Result;
 use clap::{CommandFactory, Parser};
 use tokio::runtime::Builder;
-use wallet::{Args, OverCommand, execute_continious_run, execute_keys_restoration, execute_setup, execute_subcommand};
+use wallet::{
+    Args, OverCommand, execute_continious_run, execute_keys_restoration, execute_setup,
+    execute_subcommand,
+};
 
 pub const NUM_THREADS: usize = 2;