fix: continuation of integration

key_protocol/src/key_management/key_tree/keys_private.rs

@@ -1,24 +1,19 @@
 use k256::{Scalar, elliptic_curve::PrimeField};
-use nssa_core::{NullifierPublicKey, NullifierSecretKey, encryption::IncomingViewingPublicKey};
+use nssa_core::encryption::IncomingViewingPublicKey;
 use serde::{Deserialize, Serialize};
 
 use crate::key_management::{
+    KeyChain,
     key_tree::traits::KeyNode,
-    secret_holders::{IncomingViewingSecretKey, OutgoingViewingSecretKey, SecretSpendingKey},
+    secret_holders::{PrivateKeyHolder, SecretSpendingKey},
 };
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ChildKeysPrivate {
-    pub ssk: SecretSpendingKey,
+    pub value: (KeyChain, nssa::Account),
-    pub nsk: NullifierSecretKey,
-    pub isk: IncomingViewingSecretKey,
-    pub ovk: OutgoingViewingSecretKey,
-    pub npk: NullifierPublicKey,
-    pub ipk: IncomingViewingPublicKey,
     pub ccc: [u8; 32],
     ///Can be None if root
     pub cci: Option<u32>,
-    pub account: nssa::Account,
 }
 
 impl KeyNode for ChildKeysPrivate {
@@ -36,22 +31,43 @@ impl KeyNode for ChildKeysPrivate {
         let ipk = IncomingViewingPublicKey::from_scalar(isk);
 
         Self {
-            ssk,
+            value: (
-            nsk,
+                KeyChain {
-            isk,
+                    secret_spending_key: ssk,
-            ovk,
+                    nullifer_public_key: npk,
-            npk,
+                    incoming_viewing_public_key: ipk,
-            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,
-            account: nssa::Account::default(),
         }
     }
 
     fn n_th_child(&self, cci: u32) -> Self {
-        let parent_pt = Scalar::from_repr(self.ovk.into()).unwrap()
+        let parent_pt = Scalar::from_repr(
-            + Scalar::from_repr(self.nsk.into()).unwrap()
+            self.value
-                * Scalar::from_repr(self.isk.into()).unwrap();
+                .0
+                .private_key_holder
+                .outgoing_viewing_secret_key
+                .into(),
+        )
+        .unwrap()
+            + Scalar::from_repr(self.value.0.private_key_holder.nullifier_secret_key.into())
+                .unwrap()
+                * Scalar::from_repr(
+                    self.value
+                        .0
+                        .private_key_holder
+                        .incoming_viewing_secret_key
+                        .into(),
+                )
+                .unwrap();
         let mut input = vec![];
 
         input.extend_from_slice(b"NSSA_seed_priv");
@@ -71,15 +87,21 @@ impl KeyNode for ChildKeysPrivate {
         let ipk = IncomingViewingPublicKey::from_scalar(isk);
 
         Self {
-            ssk,
+            value: (
-            nsk,
+                KeyChain {
-            isk,
+                    secret_spending_key: ssk,
-            ovk,
+                    nullifer_public_key: npk,
-            npk,
+                    incoming_viewing_public_key: ipk,
-            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),
-            account: nssa::Account::default(),
         }
     }
 
@@ -92,7 +114,19 @@ impl KeyNode for ChildKeysPrivate {
     }
 
     fn address(&self) -> nssa::Address {
-        nssa::Address::from(&self.npk)
+        nssa::Address::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
     }
 }
 
@@ -109,14 +143,14 @@ mod tests {
         assert_eq!(child_keys.cci, Some(5));
 
         assert_eq!(
-            root_keys.ssk.0,
+            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.ssk.0,
+            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
@@ -124,14 +158,14 @@ mod tests {
         );
 
         assert_eq!(
-            root_keys.nsk,
+            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.nsk,
+            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
@@ -139,14 +173,22 @@ mod tests {
         );
 
         assert_eq!(
-            root_keys.isk,
+            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.isk,
+            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
@@ -154,14 +196,22 @@ mod tests {
         );
 
         assert_eq!(
-            root_keys.ovk,
+            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.ovk,
+            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
@@ -169,14 +219,14 @@ mod tests {
         );
 
         assert_eq!(
-            root_keys.npk.0,
+            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.npk.0,
+            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
@@ -184,14 +234,14 @@ mod tests {
         );
 
         assert_eq!(
-            root_keys.ipk.0,
+            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.ipk.0,
+            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

@@ -59,6 +59,12 @@ impl KeyNode for ChildKeysPublic {
     }
 }
 
+impl<'a> From<&'a ChildKeysPublic> for &'a nssa::PrivateKey {
+    fn from(value: &'a ChildKeysPublic) -> Self {
+        &value.csk
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

key_protocol/src/key_management/key_tree/mod.rs

@@ -80,7 +80,7 @@ impl<Node: KeyNode> KeyTree<Node> {
         }
     }
 
-    pub fn generate_new_pub_keys(&mut self, parent_cci: ChainIndex) -> Option<nssa::Address> {
+    pub fn generate_new_node(&mut self, parent_cci: ChainIndex) -> Option<nssa::Address> {
         if !self.key_map.contains_key(&parent_cci) {
             return None;
         }
@@ -99,11 +99,22 @@ impl<Node: KeyNode> KeyTree<Node> {
         Some(address)
     }
 
-    pub fn get_pub_keys(&self, addr: nssa::Address) -> Option<&Node> {
+    pub fn get_node(&self, addr: nssa::Address) -> Option<&Node> {
         self.addr_map
             .get(&addr)
             .and_then(|chain_id| self.key_map.get(chain_id))
     }
+
+    pub fn get_node_mut(&mut self, addr: nssa::Address) -> Option<&mut Node> {
+        self.addr_map
+            .get(&addr)
+            .and_then(|chain_id| self.key_map.get_mut(chain_id))
+    }
+
+    pub fn insert(&mut self, addr: nssa::Address, chain_index: ChainIndex, node: Node) {
+        self.addr_map.insert(addr, chain_index.clone());
+        self.key_map.insert(chain_index, node);
+    }
 }
 
 #[cfg(test)]
@@ -145,7 +156,7 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 0);
 
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
 
         assert!(
             tree.key_map
@@ -158,12 +169,12 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 1);
 
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
-        tree.generate_new_pub_keys(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())
@@ -184,7 +195,7 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 0);
 
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
 
         assert!(
             tree.key_map
@@ -197,7 +208,7 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 1);
 
-        let key_opt = tree.generate_new_pub_keys(ChainIndex::from_str("03000000").unwrap());
+        let key_opt = tree.generate_new_node(ChainIndex::from_str("03000000").unwrap());
 
         assert_eq!(key_opt, None);
     }
@@ -214,7 +225,7 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 0);
 
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
 
         assert!(
             tree.key_map
@@ -227,7 +238,7 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 1);
 
-        tree.generate_new_pub_keys(ChainIndex::root()).unwrap();
+        tree.generate_new_node(ChainIndex::root()).unwrap();
 
         assert!(
             tree.key_map
@@ -240,7 +251,7 @@ mod tests {
 
         assert_eq!(next_last_child_for_parent_id, 2);
 
-        tree.generate_new_pub_keys(ChainIndex::from_str("00000000").unwrap())
+        tree.generate_new_node(ChainIndex::from_str("00000000").unwrap())
             .unwrap();
 
         let next_last_child_for_parent_id = tree
@@ -254,7 +265,7 @@ mod tests {
                 .contains_key(&ChainIndex::from_str("0000000000000000").unwrap())
         );
 
-        tree.generate_new_pub_keys(ChainIndex::from_str("00000000").unwrap())
+        tree.generate_new_node(ChainIndex::from_str("00000000").unwrap())
             .unwrap();
 
         let next_last_child_for_parent_id = tree
@@ -268,7 +279,7 @@ mod tests {
                 .contains_key(&ChainIndex::from_str("0000000001000000").unwrap())
         );
 
-        tree.generate_new_pub_keys(ChainIndex::from_str("00000000").unwrap())
+        tree.generate_new_node(ChainIndex::from_str("00000000").unwrap())
             .unwrap();
 
         let next_last_child_for_parent_id = tree
@@ -282,7 +293,7 @@ mod tests {
                 .contains_key(&ChainIndex::from_str("0000000002000000").unwrap())
         );
 
-        tree.generate_new_pub_keys(ChainIndex::from_str("0000000001000000").unwrap())
+        tree.generate_new_node(ChainIndex::from_str("0000000001000000").unwrap())
             .unwrap();
 
         assert!(

key_protocol/src/key_management/mod.rs

@@ -14,7 +14,7 @@ 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,

key_protocol/src/key_protocol_core/mod.rs

@@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize};
 use crate::key_management::{
     KeyChain,
     key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex},
+    secret_holders::SeedHolder,
 };
 
 pub type PublicKey = AffinePoint;
@@ -56,48 +57,62 @@ impl NSSAUserData {
     }
 
     pub fn new_with_accounts(
-        accounts_keys: HashMap<nssa::Address, nssa::PrivateKey>,
+        default_accounts_keys: HashMap<nssa::Address, nssa::PrivateKey>,
-        accounts_key_chains: HashMap<nssa::Address, (KeyChain, nssa_core::account::Account)>,
+        default_accounts_key_chains: HashMap<
+            nssa::Address,
+            (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 addresses, 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 addresses, 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,
             password: "mnemonic".to_string(),
         })
     }
 
     pub fn new_with_accounts_and_password(
-        accounts_keys: HashMap<nssa::Address, nssa::PrivateKey>,
+        default_accounts_keys: HashMap<nssa::Address, nssa::PrivateKey>,
-        accounts_key_chains: HashMap<nssa::Address, (KeyChain, nssa_core::account::Account)>,
+        default_accounts_key_chains: HashMap<
+            nssa::Address,
+            (KeyChain, nssa_core::account::Account),
+        >,
+        public_key_tree: KeyTreePublic,
+        private_key_tree: KeyTreePrivate,
         password: String,
     ) -> 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 addresses, 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 addresses, 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,
             password,
         })
     }
@@ -109,9 +124,7 @@ impl NSSAUserData {
         &mut self,
         parent_cci: ChainIndex,
     ) -> nssa::Address {
-        self.public_key_tree
+        self.public_key_tree.generate_new_node(parent_cci).unwrap()
-            .generate_new_pub_keys(parent_cci)
-            .unwrap()
     }
 
     /// Returns the signing key for public transaction signatures
@@ -119,35 +132,25 @@ impl NSSAUserData {
         &self,
         address: &nssa::Address,
     ) -> Option<&nssa::PrivateKey> {
-        self.public_key_tree.get_pub_keys(address)
+        //First seek in defaults
+        if let Some(key) = self.default_pub_account_signing_keys.get(address) {
+            Some(key)
+        //Then seek in tree
+        } else {
+            self.public_key_tree
+                .get_node(*address)
+                .and_then(|chain_keys| Some(chain_keys.into()))
+        }
     }
 
     /// Generated new private key for privacy preserving transactions
     ///
     /// Returns the address of new account
-    pub fn generate_new_privacy_preserving_transaction_key_chain(&mut self) -> nssa::Address {
+    pub fn generate_new_privacy_preserving_transaction_key_chain(
-        let key_chain = KeyChain::new_os_random();
-        let address = nssa::Address::from(&key_chain.nullifer_public_key);
-
-        self.user_private_accounts
-            .insert(address, (key_chain, nssa_core::account::Account::default()));
-
-        address
-    }
-
-    /// Generated new private key for privacy preserving transactions
-    ///
-    /// Returns the address of new account
-    pub fn generate_new_privacy_preserving_transaction_key_chain_mnemonic(
         &mut self,
+        parent_cci: ChainIndex,
     ) -> nssa::Address {
-        let key_chain = KeyChain::new_mnemonic(self.password.clone());
+        self.private_key_tree.generate_new_node(parent_cci).unwrap()
-        let address = nssa::Address::from(&key_chain.nullifer_public_key);
-
-        self.user_private_accounts
-            .insert(address, (key_chain, nssa_core::account::Account::default()));
-
-        address
     }
 
     /// Returns the signing key for public transaction signatures
@@ -155,7 +158,15 @@ impl NSSAUserData {
         &self,
         address: &nssa::Address,
     ) -> Option<&(KeyChain, nssa_core::account::Account)> {
-        self.user_private_accounts.get(address)
+        //First seek in defaults
+        if let Some(key) = self.default_user_private_accounts.get(address) {
+            Some(key)
+        //Then seek in tree
+        } else {
+            self.private_key_tree
+                .get_node(*address)
+                .and_then(|chain_keys| Some(chain_keys.into()))
+        }
     }
 
     /// Returns the signing key for public transaction signatures
@@ -163,14 +174,27 @@ impl NSSAUserData {
         &mut self,
         address: &nssa::Address,
     ) -> Option<&mut (KeyChain, nssa_core::account::Account)> {
-        self.user_private_accounts.get_mut(address)
+        //First seek in defaults
+        if let Some(key) = self.default_user_private_accounts.get_mut(address) {
+            Some(key)
+        //Then seek in tree
+        } else {
+            self.private_key_tree
+                .get_node_mut(*address)
+                .and_then(|chain_keys| Some(chain_keys.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()
     }
 }
 
@@ -182,8 +206,9 @@ mod tests {
     fn test_new_account() {
         let mut user_data = NSSAUserData::default();
 
-        let addr_pub = user_data.generate_new_public_transaction_private_key();
+        let addr_pub = user_data.generate_new_public_transaction_private_key(ChainIndex::root());
-        let addr_private = user_data.generate_new_privacy_preserving_transaction_key_chain();
+        let addr_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();
 

wallet/src/chain_storage/mod.rs

@@ -1,7 +1,13 @@
 use std::collections::HashMap;
 
 use anyhow::Result;
-use key_protocol::key_protocol_core::NSSAUserData;
+use key_protocol::{
+    key_management::{
+        key_tree::{KeyTreePrivate, KeyTreePublic},
+        secret_holders::SeedHolder,
+    },
+    key_protocol_core::NSSAUserData,
+};
 use nssa::program::Program;
 
 use crate::config::{InitialAccountData, PersistentAccountData, WalletConfig};
@@ -12,7 +18,11 @@ pub struct WalletChainStore {
 }
 
 impl WalletChainStore {
-    pub fn new(config: WalletConfig) -> Result<Self> {
+    pub fn new(
+        config: WalletConfig,
+        persistent_accounts: Vec<PersistentAccountData>,
+        password: String,
+    ) -> Result<Self> {
         let mut public_init_acc_map = HashMap::new();
         let mut private_init_acc_map = HashMap::new();
 
@@ -32,8 +42,27 @@ impl WalletChainStore {
             }
         }
 
+        let mut public_tree = KeyTreePublic::new(&SeedHolder::new_mnemonic(password.clone()));
+        let mut private_tree = KeyTreePrivate::new(&SeedHolder::new_mnemonic(password));
+
+        for pers_acc_data in persistent_accounts {
+            match pers_acc_data {
+                PersistentAccountData::Public(data) => {
+                    public_tree.insert(data.address, data.chain_index, data.data);
+                }
+                PersistentAccountData::Private(data) => {
+                    private_tree.insert(data.address, data.chain_index, data.data);
+                }
+            }
+        }
+
         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,
         })
     }
@@ -43,26 +72,20 @@ impl WalletChainStore {
         addr: nssa::Address,
         account: nssa_core::account::Account,
     ) {
-        println!("inserting at addres {}, this account {:?}", addr, account);
+        println!("inserting at address {}, this account {:?}", addr, account);
         self.user_data
-            .user_private_accounts
+            .private_key_tree
-            .entry(addr)
+            .addr_map
-            .and_modify(|(_, acc)| *acc = account);
+            .get(&addr)
-    }
+            .and_then(|chain_index| {
-
+                Some(
-    pub(crate) fn insert_account_data(&mut self, acc_data: PersistentAccountData) {
+                    self.user_data
-        match acc_data {
+                        .private_key_tree
-            PersistentAccountData::Public(acc_data) => {
+                        .key_map
-                self.user_data
+                        .entry(chain_index.clone())
-                    .pub_account_signing_keys
+                        .and_modify(|data| data.value.1 = account),
-                    .insert(acc_data.address, acc_data.pub_sign_key);
+                )
-            }
+            });
-            PersistentAccountData::Private(acc_data) => {
-                self.user_data
-                    .user_private_accounts
-                    .insert(acc_data.address, (acc_data.key_chain, acc_data.account));
-            }
-        }
     }
 }
 
@@ -180,6 +203,6 @@ mod tests {
     fn test_new_initializes_correctly() {
         let config = create_sample_wallet_config();
 
-        let _ = WalletChainStore::new(config.clone()).unwrap();
+        let _ = WalletChainStore::new(config.clone(), vec![], "test_pass".to_string()).unwrap();
     }
 }

wallet/src/cli/account.rs

@@ -1,6 +1,7 @@
 use anyhow::Result;
 use base58::ToBase58;
 use clap::Subcommand;
+use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use nssa::{Account, Address, program::Program};
 use serde::Serialize;
 
@@ -89,9 +90,15 @@ pub enum AccountSubcommand {
 #[derive(Subcommand, Debug, Clone)]
 pub enum NewSubcommand {
     ///Register new public account
-    Public {},
+    Public { 
+        #[arg(long)]
+        cci: ChainIndex
+    },
     ///Register new private account
-    Private {},
+    Private {
+        #[arg(long)]
+        cci: ChainIndex
+    },
 }
 
 impl WalletSubcommand for NewSubcommand {
@@ -100,8 +107,8 @@ impl WalletSubcommand for NewSubcommand {
         wallet_core: &mut WalletCore,
     ) -> Result<SubcommandReturnValue> {
         match self {
-            NewSubcommand::Public {} => {
+            NewSubcommand::Public { cci } => {
-                let addr = wallet_core.create_new_account_public();
+                let addr = wallet_core.create_new_account_public(cci);
 
                 println!("Generated new account with addr Public/{addr}");
 
@@ -111,8 +118,8 @@ impl WalletSubcommand for NewSubcommand {
 
                 Ok(SubcommandReturnValue::RegisterAccount { addr })
             }
-            NewSubcommand::Private {} => {
+            NewSubcommand::Private { cci } => {
-                let addr = wallet_core.create_new_account_private();
+                let addr = wallet_core.create_new_account_private(cci);
 
                 let (key, _) = wallet_core
                     .storage
@@ -270,7 +277,8 @@ impl WalletSubcommand for AccountSubcommand {
                 if !wallet_core
                     .storage
                     .user_data
-                    .user_private_accounts
+                    .private_key_tree
+                    .addr_map
                     .is_empty()
                 {
                     parse_block_range(

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 address: nssa::Address,
-    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 address: nssa::Address,
-    pub account: nssa_core::account::Account,
+    pub chain_index: ChainIndex,
-    pub key_chain: KeyChain,
+    pub data: ChildKeysPrivate,
 }
 
 //Big difference in enum variants sizes
@@ -48,6 +54,7 @@ pub enum PersistentAccountData {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PersistentStorage {
     pub accounts: Vec<PersistentAccountData>,
+    pub password: String,
     pub last_synced_block: u64,
 }
 

wallet/src/helperfunctions.rs

@@ -104,6 +104,7 @@ pub async fn fetch_persistent_storage() -> Result<PersistentStorage> {
         Err(err) => match err.kind() {
             std::io::ErrorKind::NotFound => Ok(PersistentStorage {
                 accounts: vec![],
+                password: "default".to_string(),
                 last_synced_block: 0,
             }),
             _ => {
@@ -120,29 +121,29 @@ pub fn produce_data_for_storage(
 ) -> PersistentStorage {
     let mut vec_for_storage = vec![];
 
-    for (addr, key) in &user_data.pub_account_signing_keys {
+    for (addr, key) in &user_data.public_key_tree.addr_map {
-        vec_for_storage.push(
+        if let Some(data) = user_data.public_key_tree.key_map.get(key) {
-            PersistentAccountDataPublic {
+            vec_for_storage.push(PersistentAccountDataPublic {
                 address: *addr,
-                pub_sign_key: key.clone(),
+                chain_index: key.clone(),
-            }
+                data: data.clone(),
-            .into(),
+            }.into());
-        );
+        }
     }
 
-    for (addr, (key, acc)) in &user_data.user_private_accounts {
+    for (addr, key) in &user_data.private_key_tree.addr_map {
-        vec_for_storage.push(
+        if let Some(data) = user_data.private_key_tree.key_map.get(key) {
-            PersistentAccountDataPrivate {
+            vec_for_storage.push(PersistentAccountDataPrivate {
                 address: *addr,
-                account: acc.clone(),
+                chain_index: key.clone(),
-                key_chain: key.clone(),
+                data: data.clone(),
-            }
+            }.into());
-            .into(),
+        }
-        );
     }
 
     PersistentStorage {
         accounts: vec_for_storage,
+        password: user_data.password.clone(),
         last_synced_block,
     }
 }

wallet/src/lib.rs

@@ -10,6 +10,7 @@ use common::{
 use anyhow::Result;
 use chain_storage::WalletChainStore;
 use config::WalletConfig;
+use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use log::info;
 use nssa::{
     Account, Address, privacy_preserving_transaction::message::EncryptedAccountData,
@@ -59,15 +60,13 @@ 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,
+            password,
             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, password)?;
-        }
 
         Ok(Self {
             storage,
@@ -107,16 +106,16 @@ impl WalletCore {
         Ok(config_path)
     }
 
-    pub fn create_new_account_public(&mut self) -> Address {
+    pub fn create_new_account_public(&mut self, chain_index: ChainIndex) -> Address {
         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) -> Address {
+    pub fn create_new_account_private(&mut self, chain_index: ChainIndex) -> Address {
         self.storage
             .user_data
-            .generate_new_privacy_preserving_transaction_key_chain()
+            .generate_new_privacy_preserving_transaction_key_chain(chain_index)
     }
 
     ///Get account balance
@@ -146,13 +145,12 @@ impl WalletCore {
     pub fn get_account_private(&self, addr: &Address) -> Option<Account> {
         self.storage
             .user_data
-            .user_private_accounts
+            .get_private_account(addr)
-            .get(addr)
             .map(|value| value.1.clone())
     }
 
     pub fn get_private_account_commitment(&self, addr: &Address) -> Option<Commitment> {
-        let (keys, account) = self.storage.user_data.user_private_accounts.get(addr)?;
+        let (keys, account) = self.storage.user_data.get_private_account(addr)?;
         Some(Commitment::new(&keys.nullifer_public_key, account))
     }