From 6cbc5028cfa786a041e27177e65d7d4596778be6 Mon Sep 17 00:00:00 2001
From: Pravdyvy <awxrvtb@gmail.com>
Date: Tue, 11 Nov 2025 17:25:08 +0200
Subject: [PATCH] feat: tree construction

---
 .../src/key_management/key_tree/mod.rs        | 17 +++++++++++++++++
 wallet/src/lib.rs                             | 19 +++++++++++++++++++
 wallet/src/main.rs                            |  5 ++++-
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs
index dcc027b..7a25c81 100644
--- a/key_protocol/src/key_management/key_tree/mod.rs
+++ b/key_protocol/src/key_management/key_tree/mod.rs
@@ -127,6 +127,23 @@ impl<Node: KeyNode> KeyTree<Node> {
         self.addr_map.insert(addr, chain_index.clone());
         self.key_map.insert(chain_index, node);
     }
+
+    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();
+
+            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 {
+                id_stack.push(next_id.clone());
+                next_id = next_id.next_in_line();
+            } 
+        }
+    }
 }
 
 #[cfg(test)]
diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs
index 3b494a9..c39669f 100644
--- a/wallet/src/lib.rs
+++ b/wallet/src/lib.rs
@@ -259,6 +259,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
@@ -511,3 +518,15 @@ pub async fn execute_setup(password: String) -> Result<()> {
 
     Ok(())
 }
+
+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?;
+
+    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.store_persistent_data().await?;
+
+    Ok(())
+}
diff --git a/wallet/src/main.rs b/wallet/src/main.rs
index 1fe52b3..986b27e 100644
--- a/wallet/src/main.rs
+++ b/wallet/src/main.rs
@@ -1,7 +1,7 @@
 use anyhow::Result;
 use clap::{CommandFactory, Parser};
 use tokio::runtime::Builder;
-use wallet::{Args, OverCommand, execute_continious_run, execute_setup, execute_subcommand};
+use wallet::{Args, OverCommand, execute_continious_run, execute_keys_restoration, execute_setup, execute_subcommand};
 
 pub const NUM_THREADS: usize = 2;
 
@@ -25,6 +25,9 @@ fn main() -> Result<()> {
                 OverCommand::Setup { password } => {
                     execute_setup(password).await.unwrap();
                 }
+                OverCommand::RestoreKeys { password, depth } => {
+                    execute_keys_restoration(password, depth).await.unwrap();
+                }
             }
         } else if args.continious_run {
             execute_continious_run().await.unwrap();