diff --git a/src/backend_task/identity/load_identity.rs b/src/backend_task/identity/load_identity.rs index b82da7ad9..c5f63dd43 100644 --- a/src/backend_task/identity/load_identity.rs +++ b/src/backend_task/identity/load_identity.rs @@ -4,26 +4,37 @@ use crate::context::AppContext; use crate::model::qualified_identity::PrivateKeyTarget::{ self, PrivateKeyOnMainIdentity, PrivateKeyOnVoterIdentity, }; -use crate::model::qualified_identity::encrypted_key_storage::PrivateKeyData; +use crate::model::qualified_identity::encrypted_key_storage::{ + PrivateKeyData, WalletDerivationPath, +}; use crate::model::qualified_identity::qualified_identity_public_key::QualifiedIdentityPublicKey; use crate::model::qualified_identity::{ DPNSNameInfo, IdentityStatus, IdentityType, QualifiedIdentity, }; +use crate::model::wallet::{Wallet, WalletSeedHash}; +use crate::ui::identities::add_new_identity_screen::MAX_IDENTITY_INDEX; use dash_sdk::Sdk; use dash_sdk::dashcore_rpc::dashcore::PrivateKey; use dash_sdk::dashcore_rpc::dashcore::key::Secp256k1; use dash_sdk::dpp::dashcore::hashes::Hash; use dash_sdk::dpp::document::DocumentV0Getters; use dash_sdk::dpp::identifier::MasternodeIdentifiers; +use dash_sdk::dpp::identity::KeyType; use dash_sdk::dpp::identity::SecurityLevel; use dash_sdk::dpp::identity::accessors::IdentityGettersV0; use dash_sdk::dpp::identity::identity_public_key::accessors::v0::IdentityPublicKeyGettersV0; +use dash_sdk::dpp::key_wallet::bip32::{DerivationPath, KeyDerivationType}; use dash_sdk::dpp::platform_value::Value; use dash_sdk::dpp::platform_value::string_encoding::Encoding; use dash_sdk::drive::query::{WhereClause, WhereOperator}; use dash_sdk::platform::{Document, DocumentQuery, Fetch, FetchMany, Identifier, Identity}; use egui::ahash::HashMap; use std::collections::BTreeMap; +use std::convert::TryInto; +use std::sync::{Arc, RwLock}; + +type WalletKeyMap = BTreeMap<(PrivateKeyTarget, u32), (QualifiedIdentityPublicKey, PrivateKeyData)>; +type WalletMatchResult = Option<(WalletSeedHash, u32, WalletKeyMap)>; impl AppContext { pub(super) async fn load_identity( @@ -39,6 +50,8 @@ impl AppContext { owner_private_key_input, payout_address_private_key_input, keys_input, + derive_keys_from_wallets, + selected_wallet_seed_hash, } = input; // Verify the voting private key @@ -69,6 +82,17 @@ impl AppContext { let wallets = self.wallets.read().unwrap().clone(); + if identity_type == IdentityType::User + && derive_keys_from_wallets + && let Some((_, _, wallet_private_keys)) = self.match_user_identity_keys_with_wallet( + &identity, + &wallets, + selected_wallet_seed_hash, + )? + { + encrypted_private_keys.extend(wallet_private_keys); + } + if identity_type != IdentityType::User && let Some(owner_private_key_bytes) = owner_private_key_bytes { @@ -198,44 +222,50 @@ impl AppContext { .unzip(); for (&key_id, public_key) in identity.public_keys().iter() { + let key_map_key = (PrivateKeyTarget::PrivateKeyOnMainIdentity, key_id); let qualified_key = QualifiedIdentityPublicKey::from_identity_public_key_with_wallets_check( public_key.clone(), self.network, &wallets.values().collect::>(), ); - - if let Some(wallet_derivation_path) = - qualified_key.in_wallet_at_derivation_path.clone() - { - encrypted_private_keys.insert( - (PrivateKeyTarget::PrivateKeyOnMainIdentity, key_id), - ( - qualified_key, - PrivateKeyData::AtWalletDerivationPath(wallet_derivation_path), - ), - ); - } else if let Some(private_key_bytes) = + if let Some(private_key_bytes) = public_key_lookup.get(public_key.data().0.as_slice()) { let private_data = match public_key.security_level() { SecurityLevel::MEDIUM => PrivateKeyData::AlwaysClear(*private_key_bytes), _ => PrivateKeyData::Clear(*private_key_bytes), }; - encrypted_private_keys.insert( - (PrivateKeyTarget::PrivateKeyOnMainIdentity, key_id), - (qualified_key, private_data), - ); - } else if let Some(private_key_bytes) = + encrypted_private_keys + .insert(key_map_key, (qualified_key.clone(), private_data)); + continue; + } + + if let Some(private_key_bytes) = public_key_hash_lookup.get(public_key.data().0.as_slice()) { let private_data = match public_key.security_level() { SecurityLevel::MEDIUM => PrivateKeyData::AlwaysClear(*private_key_bytes), _ => PrivateKeyData::Clear(*private_key_bytes), }; + encrypted_private_keys + .insert(key_map_key, (qualified_key.clone(), private_data)); + continue; + } + + if encrypted_private_keys.contains_key(&key_map_key) { + continue; + } + + if let Some(wallet_derivation_path) = + qualified_key.in_wallet_at_derivation_path.clone() + { encrypted_private_keys.insert( - (PrivateKeyTarget::PrivateKeyOnMainIdentity, key_id), - (qualified_key, private_data), + key_map_key, + ( + qualified_key, + PrivateKeyData::AtWalletDerivationPath(wallet_derivation_path), + ), ); } } @@ -317,8 +347,218 @@ impl AppContext { self.insert_local_qualified_identity(&qualified_identity, &wallet_info) .map_err(|e| format!("Database error: {}", e))?; + if let Some((wallet_seed_hash, identity_index)) = wallet_info + && let Some(wallet_arc) = wallets.get(&wallet_seed_hash) + { + let mut wallet = wallet_arc.write().unwrap(); + wallet + .identities + .insert(identity_index, qualified_identity.identity.clone()); + } + Ok(BackendTaskSuccessResult::Message( "Successfully loaded identity".to_string(), )) } + + fn match_user_identity_keys_with_wallet( + &self, + identity: &Identity, + wallets: &BTreeMap>>, + wallet_filter: Option, + ) -> Result { + let highest_identity_key_id = identity.public_keys().keys().copied().max().unwrap_or(0); + let top_bound = highest_identity_key_id.saturating_add(6).max(1); + + for (&wallet_seed_hash, wallet_arc) in wallets.iter() { + if wallet_filter.is_some_and(|filter| filter != wallet_seed_hash) { + continue; + } + let mut wallet = wallet_arc.write().unwrap(); + if !wallet.is_open() { + continue; + } + + if let Some((identity_index, wallet_private_keys)) = self + .attempt_match_identity_with_wallet( + identity, + &mut wallet, + wallet_seed_hash, + top_bound, + )? + { + drop(wallet); + return Ok(Some(( + wallet_seed_hash, + identity_index, + wallet_private_keys, + ))); + } + } + + Ok(None) + } + + fn attempt_match_identity_with_wallet( + &self, + identity: &Identity, + wallet: &mut Wallet, + wallet_seed_hash: WalletSeedHash, + top_bound: u32, + ) -> Result, String> { + let identity_id = identity.id(); + + if let Some((&identity_index, _)) = wallet + .identities + .iter() + .find(|(_, existing)| existing.id() == identity_id) + { + let (public_key_map, public_key_hash_map) = wallet + .identity_authentication_ecdsa_public_keys_data_map( + self.network, + identity_index, + 0..top_bound, + Some(self), + )?; + let wallet_private_keys = self.build_wallet_private_key_map( + identity, + wallet_seed_hash, + identity_index, + &public_key_map, + &public_key_hash_map, + ); + + if !wallet_private_keys.is_empty() { + return Ok(Some((identity_index, wallet_private_keys))); + } + } + + for candidate_index in 0..MAX_IDENTITY_INDEX { + let (public_key_map, public_key_hash_map) = wallet + .identity_authentication_ecdsa_public_keys_data_map( + self.network, + candidate_index, + 0..top_bound, + None, + )?; + + if !Self::identity_matches_wallet_key_material( + identity, + &public_key_map, + &public_key_hash_map, + ) { + continue; + } + + let (public_key_map, public_key_hash_map) = wallet + .identity_authentication_ecdsa_public_keys_data_map( + self.network, + candidate_index, + 0..top_bound, + Some(self), + )?; + + let wallet_private_keys = self.build_wallet_private_key_map( + identity, + wallet_seed_hash, + candidate_index, + &public_key_map, + &public_key_hash_map, + ); + + if wallet_private_keys.is_empty() { + continue; + } + + return Ok(Some((candidate_index, wallet_private_keys))); + } + + Ok(None) + } + + fn identity_matches_wallet_key_material( + identity: &Identity, + public_key_map: &BTreeMap, u32>, + public_key_hash_map: &BTreeMap<[u8; 20], u32>, + ) -> bool { + identity + .public_keys() + .values() + .any(|public_key| match public_key.key_type() { + KeyType::ECDSA_SECP256K1 => { + if public_key_map.contains_key(public_key.data().as_slice()) { + true + } else if let Ok(hash) = <[u8; 20]>::try_from(public_key.data().as_slice()) { + public_key_hash_map.contains_key(&hash) + } else { + false + } + } + KeyType::ECDSA_HASH160 => { + if let Ok(hash) = <[u8; 20]>::try_from(public_key.data().as_slice()) { + public_key_hash_map.contains_key(&hash) + } else { + false + } + } + _ => false, + }) + } + + fn build_wallet_private_key_map( + &self, + identity: &Identity, + wallet_seed_hash: WalletSeedHash, + identity_index: u32, + public_key_map: &BTreeMap, u32>, + public_key_hash_map: &BTreeMap<[u8; 20], u32>, + ) -> WalletKeyMap { + identity + .public_keys() + .values() + .filter_map(|public_key| { + let index = + match public_key.key_type() { + KeyType::ECDSA_SECP256K1 => public_key_map + .get(public_key.data().as_slice()) + .copied() + .or_else(|| { + public_key.data().as_slice().try_into().ok().and_then( + |hash: [u8; 20]| public_key_hash_map.get(&hash).copied(), + ) + }), + KeyType::ECDSA_HASH160 => public_key + .data() + .as_slice() + .try_into() + .ok() + .and_then(|hash: [u8; 20]| public_key_hash_map.get(&hash).copied()), + _ => None, + }?; + + let derivation_path = DerivationPath::identity_authentication_path( + self.network, + KeyDerivationType::ECDSA, + identity_index, + index, + ); + + let wallet_derivation_path = WalletDerivationPath { + wallet_seed_hash, + derivation_path, + }; + + Some(( + (PrivateKeyTarget::PrivateKeyOnMainIdentity, public_key.id()), + ( + QualifiedIdentityPublicKey::from_identity_public_key_in_wallet( + public_key.clone(), + Some(wallet_derivation_path.clone()), + ), + PrivateKeyData::AtWalletDerivationPath(wallet_derivation_path), + ), + )) + }) + .collect() + } } diff --git a/src/backend_task/identity/load_identity_from_wallet.rs b/src/backend_task/identity/load_identity_from_wallet.rs index 4d1b589c4..0acf3ae15 100644 --- a/src/backend_task/identity/load_identity_from_wallet.rs +++ b/src/backend_task/identity/load_identity_from_wallet.rs @@ -55,8 +55,8 @@ impl AppContext { sender .send(TaskResult::Success(Box::new( BackendTaskSuccessResult::Message(format!( - "Searching for identity using key at index {}...", - key_index + "Searching for identity at index {} using key at index {}...", + identity_index, key_index )), ))) .await diff --git a/src/backend_task/identity/mod.rs b/src/backend_task/identity/mod.rs index ddf53bd97..e9611f899 100644 --- a/src/backend_task/identity/mod.rs +++ b/src/backend_task/identity/mod.rs @@ -43,6 +43,8 @@ pub struct IdentityInputToLoad { pub owner_private_key_input: String, pub payout_address_private_key_input: String, pub keys_input: Vec, + pub derive_keys_from_wallets: bool, + pub selected_wallet_seed_hash: Option, } #[derive(Debug, Clone, PartialEq)] diff --git a/src/ui/identities/add_existing_identity_screen.rs b/src/ui/identities/add_existing_identity_screen.rs index 8348fcf8f..b1db1db2f 100644 --- a/src/ui/identities/add_existing_identity_screen.rs +++ b/src/ui/identities/add_existing_identity_screen.rs @@ -85,6 +85,7 @@ pub struct AddExistingIdentityScreen { add_identity_status: AddIdentityStatus, testnet_loaded_nodes: Option, selected_wallet: Option>>, + identity_associated_with_wallet: bool, show_password: bool, wallet_password: String, error_message: Option, @@ -112,10 +113,11 @@ impl AddExistingIdentityScreen { voting_private_key_input: String::new(), owner_private_key_input: String::new(), payout_address_private_key_input: String::new(), - keys_input: vec![String::new(), String::new(), String::new()], + keys_input: vec![], add_identity_status: AddIdentityStatus::NotStarted, testnet_loaded_nodes, selected_wallet, + identity_associated_with_wallet: true, show_password: false, wallet_password: "".to_string(), error_message: None, @@ -142,6 +144,132 @@ impl AddExistingIdentityScreen { ui.add_space(10.0); } + let wallets_snapshot: Vec<(String, Arc>)> = { + let wallets_guard = self.app_context.wallets.read().unwrap(); + wallets_guard + .values() + .map(|wallet| { + let alias = wallet + .read() + .unwrap() + .alias + .clone() + .unwrap_or_else(|| "Unnamed Wallet".to_string()); + (alias, wallet.clone()) + }) + .collect() + }; + let has_wallets = !wallets_snapshot.is_empty(); + let mut should_return_early = false; + + ui.add_space(10.0); + + ui.vertical(|ui| { + ui.horizontal(|ui| { + let checkbox_response = ui.checkbox( + &mut self.identity_associated_with_wallet, + "Try to automatically derive private keys from loaded wallet", + ); + let response = crate::ui::helpers::info_icon_button( + ui, + "When enabled, Dash Evo Tool scans the selected unlocked wallet (or all unlocked wallets) right now to find matching keys.", + ); + if response.clicked() { + self.show_pop_up_info = Some( + "When enabled, Dash Evo Tool scans the selected unlocked wallet (or all unlocked wallets) right now to find matching keys." + .to_string(), + ); + } + + if checkbox_response.changed() && !self.identity_associated_with_wallet { + self.selected_wallet = None; + } + }); + + if self.identity_associated_with_wallet { + if has_wallets { + let selected_label = self + .selected_wallet + .as_ref() + .and_then(|selected| { + wallets_snapshot.iter().find_map(|(alias, wallet)| { + if Arc::ptr_eq(selected, wallet) { + Some(alias.clone()) + } else { + None + } + }) + }) + .unwrap_or_else(|| "All unlocked wallets".to_string()); + + ComboBox::from_id_salt("identity_wallet_selector") + .selected_text(selected_label) + .show_ui(ui, |ui| { + if ui + .selectable_label( + self.selected_wallet.is_none(), + "All unlocked wallets", + ) + .clicked() + { + self.selected_wallet = None; + } + + for (alias, wallet) in &wallets_snapshot { + let is_selected = self + .selected_wallet + .as_ref() + .is_some_and(|selected| Arc::ptr_eq(selected, wallet)); + + if ui.selectable_label(is_selected, alias).clicked() { + self.selected_wallet = Some(wallet.clone()); + } + } + }); + + if let Some(selected_wallet) = &self.selected_wallet { + let wallet_still_loaded = wallets_snapshot + .iter() + .any(|(_, wallet)| Arc::ptr_eq(wallet, selected_wallet)); + + if wallet_still_loaded { + let (needed_unlock, just_unlocked) = + self.render_wallet_unlock_if_needed(ui); + if needed_unlock && !just_unlocked { + should_return_early = true; + ui.colored_label( + Color32::DARK_RED, + "Press return/enter after typing the password.", + ); + } else if just_unlocked { + ui.colored_label( + Color32::GREEN, + "Wallet unlocked. We'll pull any matching keys automatically.", + ); + } + } else { + self.selected_wallet = None; + ui.colored_label( + Color32::RED, + "Selected wallet is no longer loaded. We'll search unlocked wallets instead.", + ); + } + } + } else { + ui.colored_label( + Color32::GRAY, + "No wallets are currently loaded. Import one to scan for keys.", + ); + } + } + }); + + if should_return_early { + return action; + } + + ui.add_space(10.0); + egui::Grid::new("add_existing_identity_grid") .num_columns(2) .spacing([10.0, 10.0]) @@ -153,22 +281,25 @@ impl AddExistingIdentityScreen { ui.end_row(); ui.label("Identity Type:"); - egui::ComboBox::from_id_salt("identity_type_selector") - .selected_text(format!("{:?}", self.identity_type)) - // .width(350.0) // This sets the entire row's width - .show_ui(ui, |ui| { - ui.selectable_value(&mut self.identity_type, IdentityType::User, "User"); - ui.selectable_value( - &mut self.identity_type, - IdentityType::Masternode, - "Masternode", - ); - ui.selectable_value( - &mut self.identity_type, - IdentityType::Evonode, - "Evonode", - ); - }); + + ui.with_layout(egui::Layout::top_down(egui::Align::LEFT), |ui| { + egui::ComboBox::from_id_salt("identity_type_selector") + .selected_text(format!("{:?}", self.identity_type)) + // .width(350.0) // This sets the entire row's width + .show_ui(ui, |ui| { + ui.selectable_value(&mut self.identity_type, IdentityType::User, "User"); + ui.selectable_value( + &mut self.identity_type, + IdentityType::Masternode, + "Masternode", + ); + ui.selectable_value( + &mut self.identity_type, + IdentityType::Evonode, + "Evonode", + ); + }); + }); ui.label(""); ui.end_row(); @@ -246,10 +377,11 @@ impl AddExistingIdentityScreen { } } }); + ui.add_space(10.0); // Add button to add more keys - if ui.button("+ Add Key").clicked() { + if ui.button("+ Add key manually").clicked() { self.keys_input.push(String::new()); } ui.add_space(10.0); @@ -331,7 +463,7 @@ impl AddExistingIdentityScreen { if wallets_len == 0 { ui.colored_label( Color32::GRAY, - "No wallets available. Import or create a wallet to search by derivation path.", + "No wallets available. Import a wallet to search by derivation path.", ); return action; } @@ -380,7 +512,9 @@ impl AddExistingIdentityScreen { let identity_index_label = match self.wallet_search_mode { WalletIdentitySearchMode::SpecificIndex => "Identity index:", - WalletIdentitySearchMode::UpToIndex => "Highest identity index to search (inclusive):", + WalletIdentitySearchMode::UpToIndex => { + "Highest identity index to search (inclusive, max 29):" + } }; ui.horizontal(|ui| { @@ -436,6 +570,14 @@ impl AddExistingIdentityScreen { } fn load_identity_clicked(&mut self) -> AppAction { + let selected_wallet_seed_hash = if self.identity_associated_with_wallet { + self.selected_wallet + .as_ref() + .map(|wallet| wallet.read().unwrap().seed_hash()) + } else { + None + }; + let identity_input = IdentityInputToLoad { identity_id_input: self.identity_id_input.trim().to_string(), identity_type: self.identity_type, @@ -444,6 +586,8 @@ impl AddExistingIdentityScreen { owner_private_key_input: self.owner_private_key_input.clone(), payout_address_private_key_input: self.payout_address_private_key_input.clone(), keys_input: self.keys_input.clone(), + derive_keys_from_wallets: self.identity_associated_with_wallet, + selected_wallet_seed_hash, }; AppAction::BackendTask(BackendTask::IdentityTask(IdentityTask::LoadIdentity( diff --git a/src/ui/identities/add_new_identity_screen/mod.rs b/src/ui/identities/add_new_identity_screen/mod.rs index 040b1a540..833b40dbd 100644 --- a/src/ui/identities/add_new_identity_screen/mod.rs +++ b/src/ui/identities/add_new_identity_screen/mod.rs @@ -34,6 +34,8 @@ use std::fmt; use std::sync::atomic::Ordering; use std::sync::{Arc, RwLock}; +pub const MAX_IDENTITY_INDEX: u32 = 30; + #[derive(Debug, PartialEq, Eq, Copy, Clone)] pub enum FundingMethod { NoSelection, @@ -240,8 +242,8 @@ impl AddNewIdentityScreen { ComboBox::from_id_salt("identity_index") .selected_text(selected_text) .show_ui(ui, |ui| { - // Provide up to 30 entries for selection (0 to 29) - for i in 0..30 { + // Provide up to 30 entries for selection + for i in 0..MAX_IDENTITY_INDEX { let is_used = used_indices.contains(&i); let label = if is_used { format!("{} (used)", i)