Merge pull request #149 from vacp2p/Pravdyvy/deterministic-key-derivation

Deterministic key derivation

integration_tests/Cargo.toml

@@ -36,9 +36,9 @@ path = "../wallet"
 [dependencies.common]
 path = "../common"
 
+[dependencies.key_protocol]
+path = "../key_protocol"
+
 [dependencies.nssa]
 path = "../nssa"
 features = ["no_docker"]
-
-[dependencies.key_protocol]
-path = "../key_protocol"

integration_tests/src/lib.rs

@@ -56,6 +56,8 @@ fn make_private_account_input_from_str(account_id: &str) -> String {
 pub async fn pre_test(
     home_dir: PathBuf,
 ) -> Result<(ServerHandle, JoinHandle<Result<()>>, TempDir)> {
+    wallet::execute_setup("test_pass".to_owned()).await?;
+
     let home_dir_sequencer = home_dir.join("sequencer");
 
     let mut sequencer_config =

integration_tests/src/test_suite_map.rs

@@ -8,6 +8,7 @@ use std::{
 use actix_web::dev::ServerHandle;
 use anyhow::Result;
 use common::{PINATA_BASE58, sequencer_client::SequencerClient};
+use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use log::info;
 use nssa::{AccountId, ProgramDeploymentTransaction, program::Program};
 use nssa_core::{NullifierPublicKey, encryption::shared_key_derivation::Secp256k1Point};
@@ -23,7 +24,7 @@ use wallet::{
         pinata_program::PinataProgramAgnosticSubcommand,
         token_program::TokenProgramAgnosticSubcommand,
     },
-    config::{PersistentAccountData, PersistentStorage},
+    config::PersistentStorage,
     helperfunctions::{fetch_config, fetch_persistent_storage},
 };
 
@@ -83,7 +84,9 @@ 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: ChainIndex::root(),
+        }));
 
         let wallet_config = fetch_config().await.unwrap();
 
@@ -284,51 +287,44 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let wallet_config = fetch_config().await.unwrap();
 
         // Create new account for the token definition
-        wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
+        let SubcommandReturnValue::RegisterAccount {
-            NewSubcommand::Public {},
+            account_id: definition_account_id,
+        } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
-        .unwrap();
+        .unwrap()
+        else {
+            panic!("invalid subcommand return value");
+        };
         // Create new account for the token supply holder
-        wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
+        let SubcommandReturnValue::RegisterAccount {
-            NewSubcommand::Public {},
+            account_id: supply_account_id,
+        } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
-        .unwrap();
+        .unwrap()
+        else {
+            panic!("invalid subcommand return value");
+        };
         // Create new account for receiving a token transaction
-        wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
+        let SubcommandReturnValue::RegisterAccount {
-            NewSubcommand::Public {},
+            account_id: recipient_account_id,
+        } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
-        .unwrap();
+        .unwrap()
-
+        else {
-        let PersistentStorage {
+            panic!("invalid subcommand return value");
-            accounts: persistent_accounts,
+        };
-            last_synced_block: _,
-        } = fetch_persistent_storage().await.unwrap();
-
-        let mut new_persistent_accounts_account_id = Vec::new();
-
-        for per_acc in persistent_accounts {
-            match per_acc {
-                PersistentAccountData::Public(per_acc) => {
-                    if (per_acc.account_id.to_string() != ACC_RECEIVER)
-                        && (per_acc.account_id.to_string() != ACC_SENDER)
-                    {
-                        new_persistent_accounts_account_id.push(per_acc.account_id);
-                    }
-                }
-                _ => continue,
-            }
-        }
-
-        let [
-            definition_account_id,
-            supply_account_id,
-            recipient_account_id,
-        ] = new_persistent_accounts_account_id
-            .try_into()
-            .expect("Failed to produce new account, not present in persistent accounts");
 
         // Create new token
         let subcommand = TokenProgramAgnosticSubcommand::New {
@@ -454,7 +450,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {},
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -465,7 +463,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {},
+            NewSubcommand::Private {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -476,7 +476,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {},
+            NewSubcommand::Private {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -609,7 +611,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {},
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -620,7 +624,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {},
+            NewSubcommand::Private {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -631,7 +637,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {},
+            NewSubcommand::Private {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -745,7 +753,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {},
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -756,7 +766,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {},
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -767,7 +779,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {},
+            NewSubcommand::Private {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -881,7 +895,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: definition_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {},
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -892,7 +908,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: supply_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {},
+            NewSubcommand::Private {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -903,7 +921,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: recipient_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Public {},
+            NewSubcommand::Public {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()
@@ -1104,7 +1124,9 @@ 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 = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+            cci: ChainIndex::root(),
+        }));
 
         let sub_ret = wallet::execute_subcommand(command).await.unwrap();
         let SubcommandReturnValue::RegisterAccount {
@@ -1123,8 +1145,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let (to_keys, _) = wallet_storage
             .storage
             .user_data
-            .user_private_accounts
+            .get_private_account(&to_account_id)
-            .get(&to_account_id)
             .cloned()
             .unwrap();
 
@@ -1468,7 +1489,9 @@ 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: ChainIndex::root(),
+        }));
         let SubcommandReturnValue::RegisterAccount { account_id } =
             wallet::execute_subcommand(command).await.unwrap()
         else {
@@ -1560,7 +1583,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
 
     #[nssa_integration_test]
     pub async fn test_pinata_private_receiver_new_account() {
-        info!("########## test_pinata_private_receiver ##########");
+        info!("########## test_pinata_private_receiver_new_account ##########");
         let pinata_account_id = PINATA_BASE58;
         let pinata_prize = 150;
         let solution = 989106;
@@ -1569,7 +1592,9 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
         let SubcommandReturnValue::RegisterAccount {
             account_id: winner_account_id,
         } = wallet::execute_subcommand(Command::Account(AccountSubcommand::New(
-            NewSubcommand::Private {},
+            NewSubcommand::Private {
+                cci: ChainIndex::root(),
+            },
         )))
         .await
         .unwrap()

key_protocol/Cargo.toml

@@ -14,6 +14,7 @@ hex = "0.4.3"
 aes-gcm.workspace = true
 bip39.workspace = true
 hmac-sha512.workspace = true
+thiserror.workspace = true
 nssa-core = { path = "../nssa/core", features = ["host"] }
 
 [dependencies.common]

key_protocol/src/key_management/key_tree/chain_index.rs

@@ -0,0 +1,148 @@
+use std::{fmt::Display, str::FromStr};
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize)]
+pub struct ChainIndex(Vec<u32>);
+
+#[derive(thiserror::Error, Debug)]
+pub enum ChainIndexError {
+    #[error("No root found")]
+    NoRootFound,
+    #[error("Failed to parse segment into a number")]
+    ParseIntError(#[from] std::num::ParseIntError),
+}
+
+impl FromStr for ChainIndex {
+    type Err = ChainIndexError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if !s.starts_with('/') {
+            return Err(ChainIndexError::NoRootFound);
+        }
+
+        if s == "/" {
+            return Ok(ChainIndex(vec![]));
+        }
+
+        let uprooted_substring = s.strip_prefix("/").unwrap();
+
+        let splitted_chain: Vec<&str> = uprooted_substring.split("/").collect();
+        let mut res = vec![];
+
+        for split_ch in splitted_chain {
+            let cci = split_ch.parse()?;
+            res.push(cci);
+        }
+
+        Ok(Self(res))
+    }
+}
+
+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().saturating_sub(1))] {
+            write!(f, "{cci}/")?;
+        }
+        if let Some(last) = self.0.last() {
+            write!(f, "{}", last)?;
+        }
+        Ok(())
+    }
+}
+
+impl Default for ChainIndex {
+    fn default() -> Self {
+        ChainIndex::from_str("/").expect("Root parsing failure")
+    }
+}
+
+impl ChainIndex {
+    pub fn root() -> Self {
+        ChainIndex::default()
+    }
+
+    pub fn chain(&self) -> &[u32] {
+        &self.0
+    }
+
+    pub fn next_in_line(&self) -> ChainIndex {
+        let mut chain = self.0.clone();
+        // ToDo: Add overflow check
+        if let Some(last_p) = chain.last_mut() {
+            *last_p += 1
+        }
+
+        ChainIndex(chain)
+    }
+
+    pub fn nth_child(&self, child_id: u32) -> ChainIndex {
+        let mut chain = self.0.clone();
+        chain.push(child_id);
+
+        ChainIndex(chain)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_chain_id_root_correct() {
+        let chain_id = ChainIndex::root();
+        let chain_id_2 = ChainIndex::from_str("/").unwrap();
+
+        assert_eq!(chain_id, chain_id_2);
+    }
+
+    #[test]
+    fn test_chain_id_deser_correct() {
+        let chain_id = ChainIndex::from_str("/257").unwrap();
+
+        assert_eq!(chain_id.chain(), &[257]);
+    }
+
+    #[test]
+    fn test_chain_id_deser_failure_no_root() {
+        let chain_index_error = ChainIndex::from_str("257").err().unwrap();
+
+        assert!(matches!(chain_index_error, ChainIndexError::NoRootFound));
+    }
+
+    #[test]
+    fn test_chain_id_deser_failure_int_parsing_failure() {
+        let chain_index_error = ChainIndex::from_str("/hello").err().unwrap();
+
+        assert!(matches!(
+            chain_index_error,
+            ChainIndexError::ParseIntError(_)
+        ));
+    }
+
+    #[test]
+    fn test_chain_id_next_in_line_correct() {
+        let chain_id = ChainIndex::from_str("/257").unwrap();
+        let next_in_line = chain_id.next_in_line();
+
+        assert_eq!(next_in_line, ChainIndex::from_str("/258").unwrap());
+    }
+
+    #[test]
+    fn test_chain_id_child_correct() {
+        let chain_id = ChainIndex::from_str("/257").unwrap();
+        let child = chain_id.nth_child(3);
+
+        assert_eq!(child, ChainIndex::from_str("/257/3").unwrap());
+    }
+
+    #[test]
+    fn test_correct_display() {
+        let chainid = ChainIndex(vec![5, 7, 8]);
+
+        let string_index = format!("{chainid}");
+
+        assert_eq!(string_index, "/5/7/8".to_string());
+    }
+}

key_protocol/src/key_management/key_tree/keys_private.rs

@@ -0,0 +1,263 @@
+use k256::{Scalar, elliptic_curve::PrimeField};
+use nssa_core::encryption::IncomingViewingPublicKey;
+use serde::{Deserialize, Serialize};
+
+use crate::key_management::{
+    KeyChain,
+    key_tree::traits::KeyNode,
+    secret_holders::{PrivateKeyHolder, SecretSpendingKey},
+};
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ChildKeysPrivate {
+    pub value: (KeyChain, nssa::Account),
+    pub ccc: [u8; 32],
+    /// Can be [`None`] if root
+    pub cci: Option<u32>,
+}
+
+impl KeyNode for ChildKeysPrivate {
+    fn root(seed: [u8; 64]) -> Self {
+        let hash_value = hmac_sha512::HMAC::mac(seed, "NSSA_master_priv");
+
+        let ssk = SecretSpendingKey(
+            *hash_value
+                .first_chunk::<32>()
+                .expect("hash_value is 64 bytes, must be safe to get first 32"),
+        );
+        let ccc = *hash_value
+            .last_chunk::<32>()
+            .expect("hash_value is 64 bytes, must be safe to get last 32");
+
+        let nsk = ssk.generate_nullifier_secret_key();
+        let isk = ssk.generate_incoming_viewing_secret_key();
+        let ovk = ssk.generate_outgoing_viewing_secret_key();
+
+        let npk = (&nsk).into();
+        let ipk = IncomingViewingPublicKey::from_scalar(isk);
+
+        Self {
+            value: (
+                KeyChain {
+                    secret_spending_key: ssk,
+                    nullifer_public_key: npk,
+                    incoming_viewing_public_key: ipk,
+                    private_key_holder: PrivateKeyHolder {
+                        nullifier_secret_key: nsk,
+                        incoming_viewing_secret_key: isk,
+                        outgoing_viewing_secret_key: ovk,
+                    },
+                },
+                nssa::Account::default(),
+            ),
+            ccc,
+            cci: None,
+        }
+    }
+
+    fn nth_child(&self, cci: u32) -> Self {
+        let parent_pt = Scalar::from_repr(
+            self.value
+                .0
+                .private_key_holder
+                .outgoing_viewing_secret_key
+                .into(),
+        )
+        .expect("Key generated as scalar, must be valid representation")
+            + Scalar::from_repr(self.value.0.private_key_holder.nullifier_secret_key.into())
+                .expect("Key generated as scalar, must be valid representation")
+                * Scalar::from_repr(
+                    self.value
+                        .0
+                        .private_key_holder
+                        .incoming_viewing_secret_key
+                        .into(),
+                )
+                .expect("Key generated as scalar, must be valid representation");
+        let mut input = vec![];
+
+        input.extend_from_slice(b"NSSA_seed_priv");
+        input.extend_from_slice(&parent_pt.to_bytes());
+        input.extend_from_slice(&cci.to_le_bytes());
+
+        let hash_value = hmac_sha512::HMAC::mac(input, self.ccc);
+
+        let ssk = SecretSpendingKey(
+            *hash_value
+                .first_chunk::<32>()
+                .expect("hash_value is 64 bytes, must be safe to get first 32"),
+        );
+        let ccc = *hash_value
+            .last_chunk::<32>()
+            .expect("hash_value is 64 bytes, must be safe to get last 32");
+
+        let nsk = ssk.generate_nullifier_secret_key();
+        let isk = ssk.generate_incoming_viewing_secret_key();
+        let ovk = ssk.generate_outgoing_viewing_secret_key();
+
+        let npk = (&nsk).into();
+        let ipk = IncomingViewingPublicKey::from_scalar(isk);
+
+        Self {
+            value: (
+                KeyChain {
+                    secret_spending_key: ssk,
+                    nullifer_public_key: npk,
+                    incoming_viewing_public_key: ipk,
+                    private_key_holder: PrivateKeyHolder {
+                        nullifier_secret_key: nsk,
+                        incoming_viewing_secret_key: isk,
+                        outgoing_viewing_secret_key: ovk,
+                    },
+                },
+                nssa::Account::default(),
+            ),
+            ccc,
+            cci: Some(cci),
+        }
+    }
+
+    fn chain_code(&self) -> &[u8; 32] {
+        &self.ccc
+    }
+
+    fn child_index(&self) -> Option<u32> {
+        self.cci
+    }
+
+    fn account_id(&self) -> nssa::AccountId {
+        nssa::AccountId::from(&self.value.0.nullifer_public_key)
+    }
+}
+
+impl<'a> From<&'a ChildKeysPrivate> for &'a (KeyChain, nssa::Account) {
+    fn from(value: &'a ChildKeysPrivate) -> Self {
+        &value.value
+    }
+}
+
+impl<'a> From<&'a mut ChildKeysPrivate> for &'a mut (KeyChain, nssa::Account) {
+    fn from(value: &'a mut ChildKeysPrivate) -> Self {
+        &mut value.value
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_keys_deterministic_generation() {
+        let root_keys = ChildKeysPrivate::root([42; 64]);
+        let child_keys = root_keys.nth_child(5);
+
+        assert_eq!(root_keys.cci, None);
+        assert_eq!(child_keys.cci, Some(5));
+
+        assert_eq!(
+            root_keys.value.0.secret_spending_key.0,
+            [
+                249, 83, 253, 32, 174, 204, 185, 44, 253, 167, 61, 92, 128, 5, 152, 4, 220, 21, 88,
+                84, 167, 180, 154, 249, 44, 77, 33, 136, 59, 131, 203, 152
+            ]
+        );
+        assert_eq!(
+            child_keys.value.0.secret_spending_key.0,
+            [
+                16, 242, 229, 242, 252, 158, 153, 210, 234, 120, 70, 85, 83, 196, 5, 53, 28, 26,
+                187, 230, 22, 193, 146, 232, 237, 3, 166, 184, 122, 1, 233, 93
+            ]
+        );
+
+        assert_eq!(
+            root_keys.value.0.private_key_holder.nullifier_secret_key,
+            [
+                38, 195, 52, 182, 16, 66, 167, 156, 9, 14, 65, 100, 17, 93, 166, 71, 27, 148, 93,
+                85, 116, 109, 130, 8, 195, 222, 159, 214, 141, 41, 124, 57
+            ]
+        );
+        assert_eq!(
+            child_keys.value.0.private_key_holder.nullifier_secret_key,
+            [
+                215, 46, 2, 151, 174, 60, 86, 154, 5, 3, 175, 245, 12, 176, 220, 58, 250, 118, 236,
+                49, 254, 221, 229, 58, 40, 1, 170, 145, 175, 108, 23, 170
+            ]
+        );
+
+        assert_eq!(
+            root_keys
+                .value
+                .0
+                .private_key_holder
+                .incoming_viewing_secret_key,
+            [
+                153, 161, 15, 34, 96, 184, 165, 165, 27, 244, 155, 40, 70, 5, 241, 133, 78, 40, 61,
+                118, 48, 148, 226, 5, 97, 18, 201, 128, 82, 248, 163, 72
+            ]
+        );
+        assert_eq!(
+            child_keys
+                .value
+                .0
+                .private_key_holder
+                .incoming_viewing_secret_key,
+            [
+                192, 155, 55, 43, 164, 115, 71, 145, 227, 225, 21, 57, 55, 12, 226, 44, 10, 103,
+                39, 73, 230, 173, 60, 69, 69, 122, 110, 241, 164, 3, 192, 57
+            ]
+        );
+
+        assert_eq!(
+            root_keys
+                .value
+                .0
+                .private_key_holder
+                .outgoing_viewing_secret_key,
+            [
+                205, 87, 71, 129, 90, 242, 217, 200, 140, 252, 124, 46, 207, 7, 33, 156, 83, 166,
+                150, 81, 98, 131, 182, 156, 110, 92, 78, 140, 125, 218, 152, 154
+            ]
+        );
+        assert_eq!(
+            child_keys
+                .value
+                .0
+                .private_key_holder
+                .outgoing_viewing_secret_key,
+            [
+                131, 202, 219, 172, 219, 29, 48, 120, 226, 209, 209, 10, 216, 173, 48, 167, 233,
+                17, 35, 155, 30, 217, 176, 120, 72, 146, 250, 226, 165, 178, 255, 90
+            ]
+        );
+
+        assert_eq!(
+            root_keys.value.0.nullifer_public_key.0,
+            [
+                65, 176, 149, 243, 192, 45, 216, 177, 169, 56, 229, 7, 28, 66, 204, 87, 109, 83,
+                152, 64, 14, 188, 179, 210, 147, 60, 22, 251, 203, 70, 89, 215
+            ]
+        );
+        assert_eq!(
+            child_keys.value.0.nullifer_public_key.0,
+            [
+                69, 104, 130, 115, 48, 134, 19, 188, 67, 148, 163, 54, 155, 237, 57, 27, 136, 228,
+                111, 233, 205, 158, 149, 31, 84, 11, 241, 176, 243, 12, 138, 249
+            ]
+        );
+
+        assert_eq!(
+            root_keys.value.0.incoming_viewing_public_key.0,
+            &[
+                3, 174, 56, 136, 244, 179, 18, 122, 38, 220, 36, 50, 200, 41, 104, 167, 70, 18, 60,
+                202, 93, 193, 29, 16, 125, 252, 96, 51, 199, 152, 47, 233, 178
+            ]
+        );
+        assert_eq!(
+            child_keys.value.0.incoming_viewing_public_key.0,
+            &[
+                3, 18, 202, 246, 79, 141, 169, 51, 55, 202, 120, 169, 244, 201, 156, 162, 216, 115,
+                126, 53, 46, 94, 235, 125, 114, 178, 215, 81, 171, 93, 93, 88, 117
+            ]
+        );
+    }
+}

key_protocol/src/key_management/key_tree/keys_public.rs

@@ -0,0 +1,132 @@
+use serde::{Deserialize, Serialize};
+
+use crate::key_management::key_tree::traits::KeyNode;
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ChildKeysPublic {
+    pub csk: nssa::PrivateKey,
+    pub cpk: nssa::PublicKey,
+    pub ccc: [u8; 32],
+    /// Can be [`None`] if root
+    pub cci: Option<u32>,
+}
+
+impl KeyNode for ChildKeysPublic {
+    fn root(seed: [u8; 64]) -> Self {
+        let hash_value = hmac_sha512::HMAC::mac(seed, "NSSA_master_pub");
+
+        let csk = nssa::PrivateKey::try_new(*hash_value.first_chunk::<32>().unwrap()).unwrap();
+        let ccc = *hash_value.last_chunk::<32>().unwrap();
+        let cpk = nssa::PublicKey::new_from_private_key(&csk);
+
+        Self {
+            csk,
+            cpk,
+            ccc,
+            cci: None,
+        }
+    }
+
+    fn nth_child(&self, cci: u32) -> Self {
+        let mut hash_input = vec![];
+        hash_input.extend_from_slice(self.csk.value());
+        hash_input.extend_from_slice(&cci.to_le_bytes());
+
+        let hash_value = hmac_sha512::HMAC::mac(&hash_input, self.ccc);
+
+        let csk = nssa::PrivateKey::try_new(
+            *hash_value
+                .first_chunk::<32>()
+                .expect("hash_value is 64 bytes, must be safe to get first 32"),
+        )
+        .unwrap();
+        let ccc = *hash_value
+            .last_chunk::<32>()
+            .expect("hash_value is 64 bytes, must be safe to get last 32");
+        let cpk = nssa::PublicKey::new_from_private_key(&csk);
+
+        Self {
+            csk,
+            cpk,
+            ccc,
+            cci: Some(cci),
+        }
+    }
+
+    fn chain_code(&self) -> &[u8; 32] {
+        &self.ccc
+    }
+
+    fn child_index(&self) -> Option<u32> {
+        self.cci
+    }
+
+    fn account_id(&self) -> nssa::AccountId {
+        nssa::AccountId::from(&self.cpk)
+    }
+}
+
+impl<'a> From<&'a ChildKeysPublic> for &'a nssa::PrivateKey {
+    fn from(value: &'a ChildKeysPublic) -> Self {
+        &value.csk
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_keys_deterministic_generation() {
+        let root_keys = ChildKeysPublic::root([42; 64]);
+        let child_keys = root_keys.nth_child(5);
+
+        assert_eq!(root_keys.cci, None);
+        assert_eq!(child_keys.cci, Some(5));
+
+        assert_eq!(
+            root_keys.ccc,
+            [
+                61, 30, 91, 26, 133, 91, 236, 192, 231, 53, 186, 139, 11, 221, 202, 11, 178, 215,
+                254, 103, 191, 60, 117, 112, 1, 226, 31, 156, 83, 104, 150, 224
+            ]
+        );
+        assert_eq!(
+            child_keys.ccc,
+            [
+                67, 26, 102, 68, 189, 155, 102, 80, 199, 188, 112, 142, 207, 157, 36, 210, 48, 224,
+                35, 6, 112, 180, 11, 190, 135, 218, 9, 14, 84, 231, 58, 98
+            ]
+        );
+
+        assert_eq!(
+            root_keys.csk.value(),
+            &[
+                241, 82, 246, 237, 62, 130, 116, 47, 189, 112, 99, 67, 178, 40, 115, 245, 141, 193,
+                77, 164, 243, 76, 222, 64, 50, 146, 23, 145, 91, 164, 92, 116
+            ]
+        );
+        assert_eq!(
+            child_keys.csk.value(),
+            &[
+                11, 151, 27, 212, 167, 26, 77, 234, 103, 145, 53, 191, 184, 25, 240, 191, 156, 25,
+                60, 144, 65, 22, 193, 163, 246, 227, 212, 81, 49, 170, 33, 158
+            ]
+        );
+
+        assert_eq!(
+            root_keys.cpk.value(),
+            &[
+                220, 170, 95, 177, 121, 37, 86, 166, 56, 238, 232, 72, 21, 106, 107, 217, 158, 74,
+                133, 91, 143, 244, 155, 15, 2, 230, 223, 169, 13, 20, 163, 138
+            ]
+        );
+        assert_eq!(
+            child_keys.cpk.value(),
+            &[
+                152, 249, 236, 111, 132, 96, 184, 122, 21, 179, 240, 15, 234, 155, 164, 144, 108,
+                110, 120, 74, 176, 147, 196, 168, 243, 186, 203, 79, 97, 17, 194, 52
+            ]
+        );
+    }
+}

key_protocol/src/key_management/key_tree/mod.rs

@@ -0,0 +1,324 @@
+use std::collections::{BTreeMap, HashMap};
+
+use serde::{Deserialize, Serialize};
+
+use crate::key_management::{
+    key_tree::{
+        chain_index::ChainIndex, keys_private::ChildKeysPrivate, keys_public::ChildKeysPublic,
+        traits::KeyNode,
+    },
+    secret_holders::SeedHolder,
+};
+
+pub mod chain_index;
+pub mod keys_private;
+pub mod keys_public;
+pub mod traits;
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct KeyTree<N: KeyNode> {
+    pub key_map: BTreeMap<ChainIndex, N>,
+    pub account_id_map: HashMap<nssa::AccountId, ChainIndex>,
+}
+
+pub type KeyTreePublic = KeyTree<ChildKeysPublic>;
+pub type KeyTreePrivate = KeyTree<ChildKeysPrivate>;
+
+impl<N: KeyNode> KeyTree<N> {
+    pub fn new(seed: &SeedHolder) -> Self {
+        let seed_fit: [u8; 64] = seed
+            .seed
+            .clone()
+            .try_into()
+            .expect("SeedHolder seed is 64 bytes long");
+
+        let root_keys = N::root(seed_fit);
+        let account_id = root_keys.account_id();
+
+        let key_map = BTreeMap::from_iter([(ChainIndex::root(), root_keys)]);
+        let account_id_map = HashMap::from_iter([(account_id, ChainIndex::root())]);
+
+        Self {
+            key_map,
+            account_id_map,
+        }
+    }
+
+    pub fn new_from_root(root: N) -> Self {
+        let account_id_map = HashMap::from_iter([(root.account_id(), ChainIndex::root())]);
+        let key_map = BTreeMap::from_iter([(ChainIndex::root(), root)]);
+
+        Self {
+            key_map,
+            account_id_map,
+        }
+    }
+
+    // ToDo: Add function to create a tree from list of nodes with consistency check.
+
+    pub fn find_next_last_child_of_id(&self, parent_id: &ChainIndex) -> Option<u32> {
+        if !self.key_map.contains_key(parent_id) {
+            return None;
+        }
+
+        let leftmost_child = parent_id.nth_child(u32::MIN);
+
+        if !self.key_map.contains_key(&leftmost_child) {
+            return Some(0);
+        }
+
+        let mut right = u32::MAX - 1;
+        let mut left_border = u32::MIN;
+        let mut right_border = u32::MAX;
+
+        loop {
+            let rightmost_child = parent_id.nth_child(right);
+
+            let rightmost_ref = self.key_map.get(&rightmost_child);
+            let rightmost_ref_next = self.key_map.get(&rightmost_child.next_in_line());
+
+            match (&rightmost_ref, &rightmost_ref_next) {
+                (Some(_), Some(_)) => {
+                    left_border = right;
+                    right = (right + right_border) / 2;
+                }
+                (Some(_), None) => {
+                    break Some(right + 1);
+                }
+                (None, None) => {
+                    right_border = right;
+                    right = (left_border + right) / 2;
+                }
+                (None, Some(_)) => {
+                    unreachable!();
+                }
+            }
+        }
+    }
+
+    pub fn generate_new_node(&mut self, parent_cci: ChainIndex) -> Option<nssa::AccountId> {
+        let father_keys = self.key_map.get(&parent_cci)?;
+        let next_child_id = self
+            .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 account_id = child_keys.account_id();
+
+        self.key_map.insert(next_cci.clone(), child_keys);
+        self.account_id_map.insert(account_id, next_cci);
+
+        Some(account_id)
+    }
+
+    pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> {
+        self.account_id_map
+            .get(&account_id)
+            .and_then(|chain_id| self.key_map.get(chain_id))
+    }
+
+    pub fn get_node_mut(&mut self, account_id: nssa::AccountId) -> Option<&mut N> {
+        self.account_id_map
+            .get(&account_id)
+            .and_then(|chain_id| self.key_map.get_mut(chain_id))
+    }
+
+    pub fn insert(&mut self, account_id: nssa::AccountId, chain_index: ChainIndex, node: N) {
+        self.account_id_map.insert(account_id, chain_index.clone());
+        self.key_map.insert(chain_index, node);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use nssa::AccountId;
+
+    use super::*;
+
+    fn seed_holder_for_tests() -> SeedHolder {
+        SeedHolder {
+            seed: [42; 64].to_vec(),
+        }
+    }
+
+    #[test]
+    fn test_simple_key_tree() {
+        let seed_holder = seed_holder_for_tests();
+
+        let tree = KeyTreePublic::new(&seed_holder);
+
+        assert!(tree.key_map.contains_key(&ChainIndex::root()));
+        assert!(tree.account_id_map.contains_key(&AccountId::new([
+            46, 223, 229, 177, 59, 18, 189, 219, 153, 31, 249, 90, 112, 230, 180, 164, 80, 25, 106,
+            159, 14, 238, 1, 192, 91, 8, 210, 165, 199, 41, 60, 104,
+        ])));
+    }
+
+    #[test]
+    fn test_small_key_tree() {
+        let seed_holder = seed_holder_for_tests();
+
+        let mut tree = KeyTreePublic::new(&seed_holder);
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 0);
+
+        tree.generate_new_node(ChainIndex::root()).unwrap();
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/0").unwrap())
+        );
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        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();
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 7);
+    }
+
+    #[test]
+    fn test_key_tree_can_not_make_child_keys() {
+        let seed_holder = seed_holder_for_tests();
+
+        let mut tree = KeyTreePublic::new(&seed_holder);
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 0);
+
+        tree.generate_new_node(ChainIndex::root()).unwrap();
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/0").unwrap())
+        );
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 1);
+
+        let key_opt = tree.generate_new_node(ChainIndex::from_str("/3").unwrap());
+
+        assert_eq!(key_opt, None);
+    }
+
+    #[test]
+    fn test_key_tree_complex_structure() {
+        let seed_holder = seed_holder_for_tests();
+
+        let mut tree = KeyTreePublic::new(&seed_holder);
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 0);
+
+        tree.generate_new_node(ChainIndex::root()).unwrap();
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/0").unwrap())
+        );
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 1);
+
+        tree.generate_new_node(ChainIndex::root()).unwrap();
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/1").unwrap())
+        );
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::root())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 2);
+
+        tree.generate_new_node(ChainIndex::from_str("/0").unwrap())
+            .unwrap();
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::from_str("/0").unwrap())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 1);
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/0/0").unwrap())
+        );
+
+        tree.generate_new_node(ChainIndex::from_str("/0").unwrap())
+            .unwrap();
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::from_str("/0").unwrap())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 2);
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/0/1").unwrap())
+        );
+
+        tree.generate_new_node(ChainIndex::from_str("/0").unwrap())
+            .unwrap();
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::from_str("/0").unwrap())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 3);
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/0/2").unwrap())
+        );
+
+        tree.generate_new_node(ChainIndex::from_str("/0/1").unwrap())
+            .unwrap();
+
+        assert!(
+            tree.key_map
+                .contains_key(&ChainIndex::from_str("/0/1/0").unwrap())
+        );
+
+        let next_last_child_for_parent_id = tree
+            .find_next_last_child_of_id(&ChainIndex::from_str("/0/1").unwrap())
+            .unwrap();
+
+        assert_eq!(next_last_child_for_parent_id, 1);
+    }
+}

key_protocol/src/key_management/key_tree/traits.rs

@@ -0,0 +1,14 @@
+/// Trait, that reperesents a Node in hierarchical key tree
+pub trait KeyNode {
+    /// Tree root node
+    fn root(seed: [u8; 64]) -> Self;
+
+    /// `cci`'s child of node
+    fn nth_child(&self, cci: u32) -> Self;
+
+    fn chain_code(&self) -> &[u8; 32];
+
+    fn child_index(&self) -> Option<u32>;
+
+    fn account_id(&self) -> nssa::AccountId;
+}

key_protocol/src/key_management/mod.rs

@@ -8,12 +8,13 @@ use serde::{Deserialize, Serialize};
 pub type PublicAccountSigningKey = [u8; 32];
 
 pub mod ephemeral_key_holder;
+pub mod key_tree;
 pub mod secret_holders;
 
 #[derive(Serialize, Deserialize, Clone, Debug)]
 /// Entrypoint to key management
 pub struct KeyChain {
-    secret_spending_key: SecretSpendingKey,
+    pub secret_spending_key: SecretSpendingKey,
     pub private_key_holder: PrivateKeyHolder,
     pub nullifer_public_key: NullifierPublicKey,
     pub incoming_viewing_public_key: IncomingViewingPublicKey,
@@ -39,6 +40,25 @@ impl KeyChain {
         }
     }
 
+    pub fn new_mnemonic(passphrase: String) -> Self {
+        // Currently dropping SeedHolder at the end of initialization.
+        // Not entirely sure if we need it in the future.
+        let seed_holder = SeedHolder::new_mnemonic(passphrase);
+        let secret_spending_key = seed_holder.produce_top_secret_key_holder();
+
+        let private_key_holder = secret_spending_key.produce_private_key_holder();
+
+        let nullifer_public_key = private_key_holder.generate_nullifier_public_key();
+        let incoming_viewing_public_key = private_key_holder.generate_incoming_viewing_public_key();
+
+        Self {
+            secret_spending_key,
+            private_key_holder,
+            nullifer_public_key,
+            incoming_viewing_public_key,
+        }
+    }
+
     pub fn calculate_shared_secret_receiver(
         &self,
         ephemeral_public_key_sender: EphemeralPublicKey,

key_protocol/src/key_management/secret_holders.rs

@@ -8,6 +8,8 @@ use rand::{RngCore, rngs::OsRng};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, digest::FixedOutput};
 
+const NSSA_ENTROPY_BYTES: [u8; 32] = [0; 32];
+
 #[derive(Debug)]
 /// Seed holder. Non-clonable to ensure that different holders use different seeds.
 /// Produces `TopSecretKeyHolder` objects.
@@ -37,7 +39,8 @@ impl SeedHolder {
         let mut enthopy_bytes: [u8; 32] = [0; 32];
         OsRng.fill_bytes(&mut enthopy_bytes);
 
-        let mnemonic = Mnemonic::from_entropy(&enthopy_bytes).unwrap();
+        let mnemonic = Mnemonic::from_entropy(&enthopy_bytes)
+            .expect("Enthropy must be a multiple of 32 bytes");
         let seed_wide = mnemonic.to_seed("mnemonic");
 
         Self {
@@ -45,6 +48,16 @@ impl SeedHolder {
         }
     }
 
+    pub fn new_mnemonic(passphrase: String) -> Self {
+        let mnemonic = Mnemonic::from_entropy(&NSSA_ENTROPY_BYTES)
+            .expect("Enthropy must be a multiple of 32 bytes");
+        let seed_wide = mnemonic.to_seed(passphrase);
+
+        Self {
+            seed: seed_wide.to_vec(),
+        }
+    }
+
     pub fn generate_secret_spending_key_hash(&self) -> HashType {
         let mut hash = hmac_sha512::HMAC::mac(&self.seed, "NSSA_seed");
 
@@ -155,4 +168,14 @@ mod tests {
 
         let _ = top_secret_key_holder.generate_outgoing_viewing_secret_key();
     }
+
+    #[test]
+    fn two_seeds_generated_same_from_same_mnemonic() {
+        let mnemonic = "test_pass";
+
+        let seed_holder1 = SeedHolder::new_mnemonic(mnemonic.to_string());
+        let seed_holder2 = SeedHolder::new_mnemonic(mnemonic.to_string());
+
+        assert_eq!(seed_holder1.seed, seed_holder2.seed);
+    }
 }

key_protocol/src/key_protocol_core/mod.rs

@@ -4,16 +4,25 @@ use anyhow::Result;
 use k256::AffinePoint;
 use serde::{Deserialize, Serialize};
 
-use crate::key_management::KeyChain;
+use crate::key_management::{
+    KeyChain,
+    key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex},
+    secret_holders::SeedHolder,
+};
 
 pub type PublicKey = AffinePoint;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct NSSAUserData {
-    /// Map for all user public accounts
+    /// Default public accounts
-    pub pub_account_signing_keys: HashMap<nssa::AccountId, nssa::PrivateKey>,
+    pub default_pub_account_signing_keys: HashMap<nssa::AccountId, nssa::PrivateKey>,
-    /// Map for all user private accounts
+    /// Default private accounts
-    pub user_private_accounts: HashMap<nssa::AccountId, (KeyChain, nssa_core::account::Account)>,
+    pub default_user_private_accounts:
+        HashMap<nssa::AccountId, (KeyChain, nssa_core::account::Account)>,
+    /// Tree of public keys
+    pub public_key_tree: KeyTreePublic,
+    /// Tree of private keys
+    pub private_key_tree: KeyTreePrivate,
 }
 
 impl NSSAUserData {
@@ -47,39 +56,42 @@ impl NSSAUserData {
     }
 
     pub fn new_with_accounts(
-        accounts_keys: HashMap<nssa::AccountId, nssa::PrivateKey>,
+        default_accounts_keys: HashMap<nssa::AccountId, nssa::PrivateKey>,
-        accounts_key_chains: HashMap<nssa::AccountId, (KeyChain, nssa_core::account::Account)>,
+        default_accounts_key_chains: HashMap<
+            nssa::AccountId,
+            (KeyChain, nssa_core::account::Account),
+        >,
+        public_key_tree: KeyTreePublic,
+        private_key_tree: KeyTreePrivate,
     ) -> Result<Self> {
-        if !Self::valid_public_key_transaction_pairing_check(&accounts_keys) {
+        if !Self::valid_public_key_transaction_pairing_check(&default_accounts_keys) {
             anyhow::bail!(
                 "Key transaction pairing check not satisfied, there is account_ids, which is not derived from keys"
             );
         }
 
-        if !Self::valid_private_key_transaction_pairing_check(&accounts_key_chains) {
+        if !Self::valid_private_key_transaction_pairing_check(&default_accounts_key_chains) {
             anyhow::bail!(
                 "Key transaction pairing check not satisfied, there is account_ids, which is not derived from keys"
             );
         }
 
         Ok(Self {
-            pub_account_signing_keys: accounts_keys,
+            default_pub_account_signing_keys: default_accounts_keys,
-            user_private_accounts: accounts_key_chains,
+            default_user_private_accounts: default_accounts_key_chains,
+            public_key_tree,
+            private_key_tree,
         })
     }
 
     /// Generated new private key for public transaction signatures
     ///
     /// Returns the account_id of new account
-    pub fn generate_new_public_transaction_private_key(&mut self) -> nssa::AccountId {
+    pub fn generate_new_public_transaction_private_key(
-        let private_key = nssa::PrivateKey::new_os_random();
+        &mut self,
-        let account_id =
+        parent_cci: ChainIndex,
-            nssa::AccountId::from(&nssa::PublicKey::new_from_private_key(&private_key));
+    ) -> nssa::AccountId {
-
+        self.public_key_tree.generate_new_node(parent_cci).unwrap()
-        self.pub_account_signing_keys
-            .insert(account_id, private_key);
-
-        account_id
     }
 
     /// Returns the signing key for public transaction signatures
@@ -87,22 +99,23 @@ impl NSSAUserData {
         &self,
         account_id: &nssa::AccountId,
     ) -> Option<&nssa::PrivateKey> {
-        self.pub_account_signing_keys.get(account_id)
+        // First seek in defaults
+        if let Some(key) = self.default_pub_account_signing_keys.get(account_id) {
+            Some(key)
+        // Then seek in tree
+        } else {
+            self.public_key_tree.get_node(*account_id).map(Into::into)
+        }
     }
 
     /// Generated new private key for privacy preserving transactions
     ///
     /// Returns the account_id of new account
-    pub fn generate_new_privacy_preserving_transaction_key_chain(&mut self) -> nssa::AccountId {
+    pub fn generate_new_privacy_preserving_transaction_key_chain(
-        let key_chain = KeyChain::new_os_random();
+        &mut self,
-        let account_id = nssa::AccountId::from(&key_chain.nullifer_public_key);
+        parent_cci: ChainIndex,
-
+    ) -> nssa::AccountId {
-        self.user_private_accounts.insert(
+        self.private_key_tree.generate_new_node(parent_cci).unwrap()
-            account_id,
-            (key_chain, nssa_core::account::Account::default()),
-        );
-
-        account_id
     }
 
     /// Returns the signing key for public transaction signatures
@@ -110,7 +123,13 @@ impl NSSAUserData {
         &self,
         account_id: &nssa::AccountId,
     ) -> Option<&(KeyChain, nssa_core::account::Account)> {
-        self.user_private_accounts.get(account_id)
+        // First seek in defaults
+        if let Some(key) = self.default_user_private_accounts.get(account_id) {
+            Some(key)
+        // Then seek in tree
+        } else {
+            self.private_key_tree.get_node(*account_id).map(Into::into)
+        }
     }
 
     /// Returns the signing key for public transaction signatures
@@ -118,14 +137,27 @@ impl NSSAUserData {
         &mut self,
         account_id: &nssa::AccountId,
     ) -> Option<&mut (KeyChain, nssa_core::account::Account)> {
-        self.user_private_accounts.get_mut(account_id)
+        // First seek in defaults
+        if let Some(key) = self.default_user_private_accounts.get_mut(account_id) {
+            Some(key)
+        // Then seek in tree
+        } else {
+            self.private_key_tree
+                .get_node_mut(*account_id)
+                .map(Into::into)
+        }
     }
 }
 
 impl Default for NSSAUserData {
     fn default() -> Self {
-        // Safe unwrap as maps are empty
+        Self::new_with_accounts(
-        Self::new_with_accounts(HashMap::default(), HashMap::default()).unwrap()
+            HashMap::new(),
+            HashMap::new(),
+            KeyTreePublic::new(&SeedHolder::new_mnemonic("default".to_string())),
+            KeyTreePrivate::new(&SeedHolder::new_mnemonic("default".to_string())),
+        )
+        .unwrap()
     }
 }
 
@@ -137,20 +169,27 @@ mod tests {
     fn test_new_account() {
         let mut user_data = NSSAUserData::default();
 
-        let addr_pub = user_data.generate_new_public_transaction_private_key();
+        let account_id_pub =
-        let addr_private = user_data.generate_new_privacy_preserving_transaction_key_chain();
+            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(&addr_pub).is_some();
+        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(&addr_private).is_some();
+        let is_key_chain_generated = user_data.get_private_account(&account_id_private).is_some();
 
         assert!(is_key_chain_generated);
 
-        let addr_private_str = addr_private.to_string();
+        let account_id_private_str = account_id_private.to_string();
-        println!("{addr_private_str:#?}");
+        println!("{account_id_private_str:#?}");
-        let key_chain = &user_data.get_private_account(&addr_private).unwrap().0;
+        let key_chain = &user_data
+            .get_private_account(&account_id_private)
+            .unwrap()
+            .0;
         println!("{key_chain:#?}");
     }
 }

nssa/src/signature/public_key.rs

@@ -1,10 +1,11 @@
 use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::account::AccountId;
+use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 
 use crate::{PrivateKey, error::NssaError};
 
-#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize)]
+#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, Serialize, Deserialize)]
 pub struct PublicKey([u8; 32]);
 
 impl BorshDeserialize for PublicKey {

wallet/src/chain_storage/mod.rs

@@ -1,7 +1,13 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, hash_map::Entry};
 
 use anyhow::Result;
-use key_protocol::key_protocol_core::NSSAUserData;
+use key_protocol::{
+    key_management::{
+        key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex},
+        secret_holders::SeedHolder,
+    },
+    key_protocol_core::NSSAUserData,
+};
 use nssa::program::Program;
 
 use crate::config::{InitialAccountData, PersistentAccountData, WalletConfig};
@@ -12,7 +18,76 @@ pub struct WalletChainStore {
 }
 
 impl WalletChainStore {
-    pub fn new(config: WalletConfig) -> Result<Self> {
+    pub fn new(
+        config: WalletConfig,
+        persistent_accounts: Vec<PersistentAccountData>,
+    ) -> Result<Self> {
+        if persistent_accounts.is_empty() {
+            anyhow::bail!("Roots not found; please run setup beforehand");
+        }
+
+        let mut public_init_acc_map = HashMap::new();
+        let mut private_init_acc_map = HashMap::new();
+
+        let public_root = persistent_accounts
+            .iter()
+            .find(|data| match data {
+                &PersistentAccountData::Public(data) => data.chain_index == ChainIndex::root(),
+                _ => false,
+            })
+            .cloned()
+            .expect("Malformed persistent account data, must have public root");
+
+        let private_root = persistent_accounts
+            .iter()
+            .find(|data| match data {
+                &PersistentAccountData::Private(data) => data.chain_index == ChainIndex::root(),
+                _ => false,
+            })
+            .cloned()
+            .expect("Malformed persistent account data, must have private root");
+
+        let mut public_tree = KeyTreePublic::new_from_root(match public_root {
+            PersistentAccountData::Public(data) => data.data,
+            _ => unreachable!(),
+        });
+        let mut private_tree = KeyTreePrivate::new_from_root(match private_root {
+            PersistentAccountData::Private(data) => data.data,
+            _ => unreachable!(),
+        });
+
+        for pers_acc_data in persistent_accounts {
+            match pers_acc_data {
+                PersistentAccountData::Public(data) => {
+                    public_tree.insert(data.account_id, data.chain_index, data.data);
+                }
+                PersistentAccountData::Private(data) => {
+                    private_tree.insert(data.account_id, data.chain_index, data.data);
+                }
+                PersistentAccountData::Preconfigured(acc_data) => match acc_data {
+                    InitialAccountData::Public(data) => {
+                        public_init_acc_map.insert(data.account_id.parse()?, data.pub_sign_key);
+                    }
+                    InitialAccountData::Private(data) => {
+                        private_init_acc_map
+                            .insert(data.account_id.parse()?, (data.key_chain, data.account));
+                    }
+                },
+            }
+        }
+
+        Ok(Self {
+            user_data: NSSAUserData::new_with_accounts(
+                public_init_acc_map,
+                private_init_acc_map,
+                public_tree,
+                private_tree,
+            )?,
+            wallet_config: config,
+        })
+    }
+
+    pub fn new_storage(config: WalletConfig, password: String) -> Result<Self> {
         let mut public_init_acc_map = HashMap::new();
         let mut private_init_acc_map = HashMap::new();
 
@@ -33,8 +108,16 @@ impl WalletChainStore {
             }
         }
 
+        let public_tree = KeyTreePublic::new(&SeedHolder::new_mnemonic(password.clone()));
+        let private_tree = KeyTreePrivate::new(&SeedHolder::new_mnemonic(password));
+
         Ok(Self {
-            user_data: NSSAUserData::new_with_accounts(public_init_acc_map, private_init_acc_map)?,
+            user_data: NSSAUserData::new_with_accounts(
+                public_init_acc_map,
+                private_init_acc_map,
+                public_tree,
+                private_tree,
+            )?,
             wallet_config: config,
         })
     }
@@ -44,36 +127,40 @@ impl WalletChainStore {
         account_id: nssa::AccountId,
         account: nssa_core::account::Account,
     ) {
-        println!(
+        println!("inserting at address {account_id}, this account {account:?}");
-            "inserting at addres {}, this account {:?}",
-            account_id, account
-        );
-        self.user_data
-            .user_private_accounts
-            .entry(account_id)
-            .and_modify(|(_, acc)| *acc = account);
-    }
 
-    pub(crate) fn insert_account_data(&mut self, acc_data: PersistentAccountData) {
+        let entry = self
-        match acc_data {
+            .user_data
-            PersistentAccountData::Public(acc_data) => {
+            .default_user_private_accounts
+            .entry(account_id)
+            .and_modify(|data| data.1 = account.clone());
+
+        if matches!(entry, Entry::Vacant(_)) {
             self.user_data
-                    .pub_account_signing_keys
+                .private_key_tree
-                    .insert(acc_data.account_id, acc_data.pub_sign_key);
+                .account_id_map
-            }
+                .get(&account_id)
-            PersistentAccountData::Private(acc_data) => {
+                .map(|chain_index| {
                     self.user_data
-                    .user_private_accounts
+                        .private_key_tree
-                    .insert(acc_data.account_id, (acc_data.key_chain, acc_data.account));
+                        .key_map
-            }
+                        .entry(chain_index.clone())
+                        .and_modify(|data| data.value.1 = account)
+                });
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use key_protocol::key_management::key_tree::{
+        keys_private::ChildKeysPrivate, keys_public::ChildKeysPublic, traits::KeyNode,
+    };
+
     use super::*;
-    use crate::config::InitialAccountData;
+    use crate::config::{
+        InitialAccountData, PersistentAccountDataPrivate, PersistentAccountDataPublic,
+    };
 
     fn create_initial_accounts() -> Vec<InitialAccountData> {
         let initial_acc1 = serde_json::from_str(
@@ -179,10 +266,29 @@ mod tests {
         }
     }
 
+    fn create_sample_persistent_accounts() -> Vec<PersistentAccountData> {
+        let public_data = ChildKeysPublic::root([42; 64]);
+        let private_data = ChildKeysPrivate::root([47; 64]);
+
+        vec![
+            PersistentAccountData::Public(PersistentAccountDataPublic {
+                account_id: public_data.account_id(),
+                chain_index: ChainIndex::root(),
+                data: public_data,
+            }),
+            PersistentAccountData::Private(PersistentAccountDataPrivate {
+                account_id: private_data.account_id(),
+                chain_index: ChainIndex::root(),
+                data: private_data,
+            }),
+        ]
+    }
+
     #[test]
     fn test_new_initializes_correctly() {
         let config = create_sample_wallet_config();
+        let accs = create_sample_persistent_accounts();
 
-        let _ = WalletChainStore::new(config.clone()).unwrap();
+        let _ = WalletChainStore::new(config.clone(), accs).unwrap();
     }
 }

wallet/src/cli/account.rs

@@ -2,6 +2,7 @@ use anyhow::Result;
 use base58::ToBase58;
 use clap::Subcommand;
 use itertools::Itertools as _;
+use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use nssa::{Account, AccountId, program::Program};
 use serde::Serialize;
 
@@ -93,9 +94,17 @@ pub enum AccountSubcommand {
 #[derive(Subcommand, Debug, Clone)]
 pub enum NewSubcommand {
     /// Register new public account
-    Public {},
+    Public {
+        #[arg(long)]
+        /// Chain index of a parent node
+        cci: ChainIndex,
+    },
     /// Register new private account
-    Private {},
+    Private {
+        #[arg(long)]
+        /// Chain index of a parent node
+        cci: ChainIndex,
+    },
 }
 
 impl WalletSubcommand for NewSubcommand {
@@ -104,8 +113,8 @@ impl WalletSubcommand for NewSubcommand {
         wallet_core: &mut WalletCore,
     ) -> Result<SubcommandReturnValue> {
         match self {
-            NewSubcommand::Public {} => {
+            NewSubcommand::Public { cci } => {
-                let account_id = wallet_core.create_new_account_public();
+                let account_id = wallet_core.create_new_account_public(cci);
 
                 println!("Generated new account with account_id Public/{account_id}");
 
@@ -115,8 +124,8 @@ impl WalletSubcommand for NewSubcommand {
 
                 Ok(SubcommandReturnValue::RegisterAccount { account_id })
             }
-            NewSubcommand::Private {} => {
+            NewSubcommand::Private { cci } => {
-                let account_id = wallet_core.create_new_account_private();
+                let account_id = wallet_core.create_new_account_private(cci);
 
                 let (key, _) = wallet_core
                     .storage
@@ -275,12 +284,19 @@ impl WalletSubcommand for AccountSubcommand {
                     .await?
                     .last_block;
 
-                if !wallet_core
+                if wallet_core
                     .storage
                     .user_data
-                    .user_private_accounts
+                    .private_key_tree
+                    .account_id_map
                     .is_empty()
                 {
+                    wallet_core.last_synced_block = curr_last_block;
+
+                    let path = wallet_core.store_persistent_data().await?;
+
+                    println!("Stored persistent data at {path:#?}");
+                } else {
                     parse_block_range(
                         last_synced_block + 1,
                         curr_last_block,
@@ -288,12 +304,6 @@ impl WalletSubcommand for AccountSubcommand {
                         wallet_core,
                     )
                     .await?;
-                } else {
-                    wallet_core.last_synced_block = curr_last_block;
-
-                    let path = wallet_core.store_persistent_data().await?;
-
-                    println!("Stored persistent data at {path:#?}");
                 }
 
                 Ok(SubcommandReturnValue::SyncedToBlock(curr_last_block))
@@ -301,14 +311,28 @@ impl WalletSubcommand for AccountSubcommand {
             AccountSubcommand::List {} => {
                 let user_data = &wallet_core.storage.user_data;
                 let accounts = user_data
-                    .pub_account_signing_keys
+                    .default_pub_account_signing_keys
                     .keys()
-                    .map(|id| format!("Public/{id}"))
+                    .map(|id| format!("Preconfigured Public/{id}"))
                     .chain(
                         user_data
-                            .user_private_accounts
+                            .default_user_private_accounts
                             .keys()
-                            .map(|id| format!("Private/{id}")),
+                            .map(|id| format!("Preconfigured Private/{id}")),
+                    )
+                    .chain(
+                        user_data
+                            .public_key_tree
+                            .account_id_map
+                            .iter()
+                            .map(|(id, chain_index)| format!("{chain_index} Public/{id}")),
+                    )
+                    .chain(
+                        user_data
+                            .private_key_tree
+                            .account_id_map
+                            .iter()
+                            .map(|(id, chain_index)| format!("{chain_index} Private/{id}")),
                     )
                     .format(",\n");
 

wallet/src/config.rs

@@ -1,4 +1,9 @@
-use key_protocol::key_management::KeyChain;
+use key_protocol::key_management::{
+    KeyChain,
+    key_tree::{
+        chain_index::ChainIndex, keys_private::ChildKeysPrivate, keys_public::ChildKeysPublic,
+    },
+};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -10,7 +15,8 @@ pub struct InitialAccountDataPublic {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PersistentAccountDataPublic {
     pub account_id: nssa::AccountId,
-    pub pub_sign_key: nssa::PrivateKey,
+    pub chain_index: ChainIndex,
+    pub data: ChildKeysPublic,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -23,8 +29,8 @@ pub struct InitialAccountDataPrivate {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PersistentAccountDataPrivate {
     pub account_id: nssa::AccountId,
-    pub account: nssa_core::account::Account,
+    pub chain_index: ChainIndex,
-    pub key_chain: KeyChain,
+    pub data: ChildKeysPrivate,
 }
 
 // Big difference in enum variants sizes
@@ -45,6 +51,7 @@ pub enum InitialAccountData {
 pub enum PersistentAccountData {
     Public(PersistentAccountDataPublic),
     Private(PersistentAccountDataPrivate),
+    Preconfigured(InitialAccountData),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -67,6 +74,7 @@ impl PersistentAccountData {
         match &self {
             Self::Public(acc) => acc.account_id,
             Self::Private(acc) => acc.account_id,
+            Self::Preconfigured(acc) => acc.account_id(),
         }
     }
 }
@@ -95,6 +103,12 @@ impl From<PersistentAccountDataPrivate> for PersistentAccountData {
     }
 }
 
+impl From<InitialAccountData> for PersistentAccountData {
+    fn from(value: InitialAccountData) -> Self {
+        Self::Preconfigured(value)
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GasConfig {
     /// Gas spent per deploying one byte of data

wallet/src/helperfunctions.rs

@@ -12,6 +12,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use crate::{
     HOME_DIR_ENV_VAR,
     config::{
+        InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic,
         PersistentAccountDataPrivate, PersistentAccountDataPublic, PersistentStorage, WalletConfig,
     },
 };
@@ -90,7 +91,7 @@ pub async fn fetch_config() -> Result<WalletConfig> {
 
 /// Fetch data stored at home
 ///
-/// If file not present, it is considered as empty list of persistent accounts
+/// File must be created through setup beforehand.
 pub async fn fetch_persistent_storage() -> Result<PersistentStorage> {
     let home = get_home()?;
     let accs_path = home.join("storage.json");
@@ -102,10 +103,9 @@ pub async fn fetch_persistent_storage() -> Result<PersistentStorage> {
             Ok(serde_json::from_slice(&storage_content)?)
         }
         Err(err) => match err.kind() {
-            std::io::ErrorKind::NotFound => Ok(PersistentStorage {
+            std::io::ErrorKind::NotFound => {
-                accounts: vec![],
+                anyhow::bail!("Not found, please setup roots from config command beforehand");
-                last_synced_block: 0,
+            }
-            }),
             _ => {
                 anyhow::bail!("IO error {err:#?}");
             }
@@ -120,26 +120,52 @@ pub fn produce_data_for_storage(
 ) -> PersistentStorage {
     let mut vec_for_storage = vec![];
 
-    for (account_id, key) in &user_data.pub_account_signing_keys {
+    for (account_id, key) in &user_data.public_key_tree.account_id_map {
+        if let Some(data) = user_data.public_key_tree.key_map.get(key) {
             vec_for_storage.push(
                 PersistentAccountDataPublic {
                     account_id: *account_id,
-                pub_sign_key: key.clone(),
+                    chain_index: key.clone(),
+                    data: data.clone(),
                 }
                 .into(),
             );
         }
+    }
 
-    for (account_id, (key, acc)) in &user_data.user_private_accounts {
+    for (account_id, key) in &user_data.private_key_tree.account_id_map {
+        if let Some(data) = user_data.private_key_tree.key_map.get(key) {
             vec_for_storage.push(
                 PersistentAccountDataPrivate {
                     account_id: *account_id,
-                account: acc.clone(),
+                    chain_index: key.clone(),
-                key_chain: key.clone(),
+                    data: data.clone(),
                 }
                 .into(),
             );
         }
+    }
+
+    for (account_id, key) in &user_data.default_pub_account_signing_keys {
+        vec_for_storage.push(
+            InitialAccountData::Public(InitialAccountDataPublic {
+                account_id: account_id.to_string(),
+                pub_sign_key: key.clone(),
+            })
+            .into(),
+        )
+    }
+
+    for (account_id, (key_chain, account)) in &user_data.default_user_private_accounts {
+        vec_for_storage.push(
+            InitialAccountData::Private(InitialAccountDataPrivate {
+                account_id: account_id.to_string(),
+                account: account.clone(),
+                key_chain: key_chain.clone(),
+            })
+            .into(),
+        )
+    }
 
     PersistentStorage {
         accounts: vec_for_storage,

wallet/src/lib.rs

@@ -10,6 +10,7 @@ use common::{
     transaction::{EncodedTransaction, NSSATransaction},
 };
 use config::WalletConfig;
+use key_protocol::key_management::key_tree::{chain_index::ChainIndex, traits::KeyNode};
 use log::info;
 use nssa::{
     Account, AccountId, privacy_preserving_transaction::message::EncryptedAccountData,
@@ -54,15 +55,12 @@ impl WalletCore {
         let client = Arc::new(SequencerClient::new(config.sequencer_addr.clone())?);
         let tx_poller = TxPoller::new(config.clone(), client.clone());
 
-        let mut storage = WalletChainStore::new(config)?;
-
         let PersistentStorage {
             accounts: persistent_accounts,
             last_synced_block,
         } = fetch_persistent_storage().await?;
-        for pers_acc_data in persistent_accounts {
+
-            storage.insert_account_data(pers_acc_data);
+        let storage = WalletChainStore::new(config, persistent_accounts)?;
-        }
 
         Ok(Self {
             storage,
@@ -72,6 +70,23 @@ impl WalletCore {
         })
     }
 
+    pub async fn start_from_config_new_storage(
+        config: WalletConfig,
+        password: String,
+    ) -> Result<Self> {
+        let client = Arc::new(SequencerClient::new(config.sequencer_addr.clone())?);
+        let tx_poller = TxPoller::new(config.clone(), client.clone());
+
+        let storage = WalletChainStore::new_storage(config, password)?;
+
+        Ok(Self {
+            storage,
+            poller: tx_poller,
+            sequencer_client: client.clone(),
+            last_synced_block: 0,
+        })
+    }
+
     /// Store persistent data at home
     pub async fn store_persistent_data(&self) -> Result<PathBuf> {
         let home = get_home()?;
@@ -102,16 +117,16 @@ impl WalletCore {
         Ok(config_path)
     }
 
-    pub fn create_new_account_public(&mut self) -> AccountId {
+    pub fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId {
         self.storage
             .user_data
-            .generate_new_public_transaction_private_key()
+            .generate_new_public_transaction_private_key(chain_index)
     }
 
-    pub fn create_new_account_private(&mut self) -> AccountId {
+    pub fn create_new_account_private(&mut self, chain_index: ChainIndex) -> AccountId {
         self.storage
             .user_data
-            .generate_new_privacy_preserving_transaction_key_chain()
+            .generate_new_privacy_preserving_transaction_key_chain(chain_index)
     }
 
     /// Get account balance
@@ -144,17 +159,12 @@ impl WalletCore {
     pub fn get_account_private(&self, account_id: &AccountId) -> Option<Account> {
         self.storage
             .user_data
-            .user_private_accounts
+            .get_private_account(account_id)
-            .get(account_id)
             .map(|value| value.1.clone())
     }
 
     pub fn get_private_account_commitment(&self, account_id: &AccountId) -> Option<Commitment> {
-        let (keys, account) = self
+        let (keys, account) = self.storage.user_data.get_private_account(account_id)?;
-            .storage
-            .user_data
-            .user_private_accounts
-            .get(account_id)?;
         Some(Commitment::new(&keys.nullifer_public_key, account))
     }
 
@@ -237,6 +247,20 @@ pub enum Command {
     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,
+    },
+}
+
 /// To execute commands, env var NSSA_WALLET_HOME_DIR must be set into directory with config
 ///
 /// All account adresses must be valid 32 byte base58 strings.
@@ -251,7 +275,7 @@ pub struct Args {
     pub continious_run: bool,
     /// Wallet command
     #[command(subcommand)]
-    pub command: Option<Command>,
+    pub command: Option<OverCommand>,
 }
 
 #[derive(Debug, Clone)]
@@ -346,7 +370,7 @@ pub async fn parse_block_range(
                 let mut affected_accounts = vec![];
 
                 for (acc_account_id, (key_chain, _)) in
-                    &wallet_core.storage.user_data.user_private_accounts
+                    &wallet_core.storage.user_data.default_user_private_accounts
                 {
                     let view_tag = EncryptedAccountData::compute_view_tag(
                         key_chain.nullifer_public_key.clone(),
@@ -383,6 +407,51 @@ pub async fn parse_block_range(
                     }
                 }
 
+                for keys_node in wallet_core
+                    .storage
+                    .user_data
+                    .private_key_tree
+                    .key_map
+                    .values()
+                {
+                    let acc_account_id = keys_node.account_id();
+                    let key_chain = &keys_node.value.0;
+
+                    let view_tag = EncryptedAccountData::compute_view_tag(
+                        key_chain.nullifer_public_key.clone(),
+                        key_chain.incoming_viewing_public_key.clone(),
+                    );
+
+                    for (ciph_id, encrypted_data) in tx
+                        .message()
+                        .encrypted_private_post_states
+                        .iter()
+                        .enumerate()
+                    {
+                        if encrypted_data.view_tag == view_tag {
+                            let ciphertext = &encrypted_data.ciphertext;
+                            let commitment = &tx.message.new_commitments[ciph_id];
+                            let shared_secret = key_chain
+                                .calculate_shared_secret_receiver(encrypted_data.epk.clone());
+
+                            let res_acc = nssa_core::EncryptionScheme::decrypt(
+                                ciphertext,
+                                &shared_secret,
+                                commitment,
+                                ciph_id as u32,
+                            );
+
+                            if let Some(res_acc) = res_acc {
+                                println!(
+                                    "Received new account for account_id {acc_account_id:#?} with account object {res_acc:#?}"
+                                );
+
+                                affected_accounts.push((acc_account_id, res_acc));
+                            }
+                        }
+                    }
+                }
+
                 for (affected_account_id, new_acc) in affected_accounts {
                     wallet_core
                         .storage
@@ -430,3 +499,12 @@ pub async fn execute_continious_run() -> Result<()> {
         latest_block_num = seq_client.get_last_block().await?.last_block;
     }
 }
+
+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_core.store_persistent_data().await?;
+
+    Ok(())
+}

wallet/src/main.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use clap::{CommandFactory, Parser};
 use tokio::runtime::Builder;
-use wallet::{Args, execute_continious_run, execute_subcommand};
+use wallet::{Args, OverCommand, execute_continious_run, execute_setup, execute_subcommand};
 
 pub const NUM_THREADS: usize = 2;
 
@@ -21,9 +21,15 @@ fn main() -> Result<()> {
     env_logger::init();
 
     runtime.block_on(async move {
-        if let Some(command) = args.command {
+        if let Some(overcommand) = args.command {
-            // TODO: It should return error, not panic
+            match overcommand {
+                OverCommand::Command(command) => {
                     execute_subcommand(command).await.unwrap();
+                }
+                OverCommand::Setup { password } => {
+                    execute_setup(password).await.unwrap();
+                }
+            }
         } else if args.continious_run {
             execute_continious_run().await.unwrap();
         } else {