1use std::{
75 collections::HashMap,
76 fmt::Debug,
77 hash::{Hash, Hasher},
78 str::FromStr,
79 sync::{Arc, Weak},
80};
81
82use crate::neo_crypto::utils::ToHexString;
83use primitive_types::H160;
84use serde_derive::{Deserialize, Serialize};
85use signature::{hazmat::PrehashSigner, Error, SignerMut};
86
87use crate::{
88 neo_builder::VerificationScript,
89 neo_clients::{public_key_to_address, APITrait, JsonRpcProvider, ProviderError, RpcClient},
90 neo_crypto::{private_key_from_wif, KeyPair, Secp256r1PublicKey, Secp256r1Signature},
91 neo_protocol::{get_nep2_from_private_key, get_private_key_from_nep2},
92 neo_types::{
93 deserialize_address_or_script_hash, serialize_address_or_script_hash, Address,
94 AddressOrScriptHash, ContractParameterType, ScriptHash,
95 },
96 neo_wallets::{NEP6Account, NEP6Contract, NEP6Parameter, Wallet},
97 vec_to_array32, Base64Encode, ScriptHashExtension,
98};
99
100pub trait AccountTrait: Sized + PartialEq + Send + Sync + Debug + Clone {
101 type Error: Sync + Send + Debug + Sized;
102
103 fn key_pair(&self) -> &Option<KeyPair>;
105 fn address_or_scripthash(&self) -> &AddressOrScriptHash;
106 fn label(&self) -> &Option<String>;
107 fn verification_script(&self) -> &Option<VerificationScript>;
108 fn is_locked(&self) -> bool;
109 fn encrypted_private_key(&self) -> &Option<String>;
110 fn signing_threshold(&self) -> &Option<u32>;
111 fn nr_of_participants(&self) -> &Option<u32>;
112 fn set_key_pair(&mut self, key_pair: Option<KeyPair>);
113 fn set_address_or_scripthash(&mut self, address_or_scripthash: AddressOrScriptHash);
114 fn set_label(&mut self, label: Option<String>);
115 fn set_verification_script(&mut self, verification_script: Option<VerificationScript>);
116 fn set_locked(&mut self, is_locked: bool);
117 fn set_encrypted_private_key(&mut self, encrypted_private_key: Option<String>);
118
119 fn set_signing_threshold(&mut self, signing_threshold: Option<u32>);
120 fn set_nr_of_participants(&mut self, nr_of_participants: Option<u32>);
121
122 fn new(
123 address: AddressOrScriptHash,
124 label: Option<String>,
125 verification_script: Option<VerificationScript>,
126 signing_threshold: Option<u32>,
127 nr_of_participants: Option<u32>,
128 ) -> Self;
129
130 fn from_key_pair(
131 key_pair: KeyPair,
132 signing_threshold: Option<u32>,
133 nr_of_participants: Option<u32>,
134 ) -> Result<Self, Self::Error>;
135
136 fn from_key_pair_opt(
137 key_pair: Option<KeyPair>,
138 address: AddressOrScriptHash,
139 label: Option<String>,
140 verification_script: Option<VerificationScript>,
141 is_locked: bool,
142 is_default: bool,
143 encrypted_private_key: Option<String>,
144 signing_threshold: Option<u32>,
145 nr_of_participants: Option<u32>,
146 ) -> Self;
147
148 fn from_wif(wif: &str) -> Result<Self, Self::Error>;
149
150 fn decrypt_private_key(&mut self, password: &str) -> Result<(), Self::Error>;
151
152 fn encrypt_private_key(&mut self, password: &str) -> Result<(), Self::Error>;
153
154 fn get_script_hash(&self) -> ScriptHash;
155
156 fn get_signing_threshold(&self) -> Result<u32, Self::Error>;
157
158 fn get_nr_of_participants(&self) -> Result<u32, Self::Error>;
159
160 fn from_verification_script(script: &VerificationScript) -> Result<Self, Self::Error>;
161
162 fn from_public_key(public_key: &Secp256r1PublicKey) -> Result<Self, Self::Error>;
163
164 fn set_wallet(&mut self, wallet: Option<Weak<Wallet>>);
165
166 fn get_wallet(&self) -> Option<Arc<Wallet>>;
167
168 fn multi_sig_from_public_keys(
169 public_keys: &mut [Secp256r1PublicKey],
170 signing_threshold: u32,
171 ) -> Result<Self, Self::Error>;
172 fn multi_sig_from_addr(
173 address: String,
174 signing_threshold: u8,
175 nr_of_participants: u8,
176 ) -> Result<Self, Self::Error>;
177
178 fn from_address(address: &str) -> Result<Self, Self::Error>;
179
180 fn from_script_hash(script_hash: &H160) -> Result<Self, Self::Error>;
181
182 fn create() -> Result<Self, Self::Error>;
183
184 fn is_multi_sig(&self) -> bool;
185}
186
187#[derive(Debug, Clone, Serialize, Deserialize, Default)]
188pub struct Account {
189 #[serde(skip)]
190 pub key_pair: Option<KeyPair>,
191 #[serde(
192 serialize_with = "serialize_address_or_script_hash",
193 deserialize_with = "deserialize_address_or_script_hash"
194 )]
195 pub address_or_scripthash: AddressOrScriptHash,
196 pub label: Option<String>,
197 pub verification_script: Option<VerificationScript>,
198 pub is_default: bool,
199 pub is_locked: bool,
200 pub encrypted_private_key: Option<String>,
201 pub signing_threshold: Option<u32>,
202 pub nr_of_participants: Option<u32>,
203 #[serde(skip)]
204 pub wallet: Option<Weak<Wallet>>,
205}
206
207impl Account {
208 pub fn get_address(&self) -> String {
209 self.address_or_scripthash.address()
210 }
211
212 pub fn get_script_hash(&self) -> H160 {
213 self.address_or_scripthash.script_hash()
214 }
215
216 pub fn get_verification_script(&self) -> Option<VerificationScript> {
217 self.verification_script.clone()
218 }
219 pub fn get_public_key(&self) -> Option<Secp256r1PublicKey> {
220 self.key_pair.as_ref().map(|k| k.public_key.clone())
221 }
222}
223
224impl From<H160> for Account {
225 fn from(script_hash: H160) -> Self {
226 Self {
227 address_or_scripthash: AddressOrScriptHash::ScriptHash(script_hash),
228 ..Default::default()
229 }
230 }
231}
232
233impl From<&H160> for Account {
234 fn from(script_hash: &H160) -> Self {
235 Self {
236 address_or_scripthash: AddressOrScriptHash::ScriptHash(script_hash.clone()),
237 ..Default::default()
238 }
239 }
240}
241
242impl PartialEq for Account {
243 fn eq(&self, other: &Self) -> bool {
244 self.address_or_scripthash == other.address_or_scripthash
245 && self.label == other.label
246 && self.verification_script == other.verification_script
247 && self.is_locked == other.is_locked
248 && self.encrypted_private_key == other.encrypted_private_key
249 && self.signing_threshold == other.signing_threshold
250 && self.nr_of_participants == other.nr_of_participants
251 }
252}
253
254impl Hash for Account {
255 fn hash<H: Hasher>(&self, state: &mut H) {
256 self.address_or_scripthash.hash(state);
257 self.label.hash(state);
258 self.verification_script.hash(state);
259 self.is_locked.hash(state);
260 self.encrypted_private_key.hash(state);
261 self.signing_threshold.hash(state);
262 self.nr_of_participants.hash(state);
263 }
264}
265
266impl AccountTrait for Account {
267 type Error = ProviderError;
268
269 fn key_pair(&self) -> &Option<KeyPair> {
270 &self.key_pair
271 }
272
273 fn address_or_scripthash(&self) -> &AddressOrScriptHash {
274 &self.address_or_scripthash
275 }
276
277 fn label(&self) -> &Option<String> {
278 &self.label
279 }
280
281 fn verification_script(&self) -> &Option<VerificationScript> {
282 &self.verification_script
283 }
284
285 fn is_locked(&self) -> bool {
286 self.is_locked
287 }
288
289 fn encrypted_private_key(&self) -> &Option<String> {
290 &self.encrypted_private_key
291 }
292
293 fn signing_threshold(&self) -> &Option<u32> {
294 &self.signing_threshold
295 }
296
297 fn nr_of_participants(&self) -> &Option<u32> {
298 &self.nr_of_participants
299 }
300
301 fn set_key_pair(&mut self, key_pair: Option<KeyPair>) {
302 self.key_pair = key_pair;
303 }
304
305 fn set_address_or_scripthash(&mut self, address_or_scripthash: AddressOrScriptHash) {
306 self.address_or_scripthash = address_or_scripthash;
307 }
308
309 fn set_label(&mut self, label: Option<String>) {
310 self.label = label;
311 }
312
313 fn set_verification_script(&mut self, verification_script: Option<VerificationScript>) {
314 self.verification_script = verification_script;
315 }
316
317 fn set_locked(&mut self, is_locked: bool) {
318 self.is_locked = is_locked;
319 }
320
321 fn set_encrypted_private_key(&mut self, encrypted_private_key: Option<String>) {
322 self.encrypted_private_key = encrypted_private_key;
323 }
324
325 fn set_signing_threshold(&mut self, signing_threshold: Option<u32>) {
326 self.signing_threshold = signing_threshold;
327 }
328
329 fn set_nr_of_participants(&mut self, nr_of_participants: Option<u32>) {
330 self.nr_of_participants = nr_of_participants;
331 }
332
333 fn new(
334 address: AddressOrScriptHash,
335 label: Option<String>,
336 verification_script: Option<VerificationScript>,
337 signing_threshold: Option<u32>,
338 nr_of_participants: Option<u32>,
339 ) -> Self {
340 Self {
341 key_pair: None,
342 address_or_scripthash: address,
343 label,
344 verification_script,
345 is_default: false,
346 is_locked: false,
347 encrypted_private_key: None,
348 signing_threshold,
349 nr_of_participants,
350 wallet: None,
351 }
352 }
353
354 fn from_key_pair(
355 key_pair: KeyPair,
356 signing_threshold: Option<u32>,
357 nr_of_participants: Option<u32>,
358 ) -> Result<Self, Self::Error> {
359 let address = public_key_to_address(&key_pair.public_key);
360 Ok(Self {
361 key_pair: Some(key_pair.clone()),
362 address_or_scripthash: AddressOrScriptHash::Address(address.clone()),
363 label: Some(address),
364 verification_script: Some(VerificationScript::from_public_key(
365 &key_pair.clone().public_key(),
366 )),
367 is_default: false,
368 is_locked: false,
369 encrypted_private_key: None,
370 signing_threshold,
371 nr_of_participants,
372 wallet: None,
373 })
374 }
375
376 fn from_key_pair_opt(
377 key_pair: Option<KeyPair>,
378 address: AddressOrScriptHash,
379 label: Option<String>,
380 verification_script: Option<VerificationScript>,
381 is_locked: bool,
382 _is_default: bool,
383 encrypted_private_key: Option<String>,
384 signing_threshold: Option<u32>,
385 nr_of_participants: Option<u32>,
386 ) -> Self {
387 Self {
388 key_pair,
389 address_or_scripthash: address,
390 label,
391 verification_script,
392 is_default: false,
393 is_locked,
394 encrypted_private_key,
395 signing_threshold,
396 nr_of_participants,
397 wallet: None,
398 }
399 }
400
401 fn from_wif(wif: &str) -> Result<Self, Self::Error> {
402 let key_pair = KeyPair::from_secret_key(&private_key_from_wif(wif)?);
403 Self::from_key_pair(key_pair, None, None)
404 }
405
406 fn decrypt_private_key(&mut self, password: &str) -> Result<(), Self::Error> {
407 if self.key_pair.is_some() {
408 return Ok(());
409 }
410
411 let encrypted_private_key = self
412 .encrypted_private_key
413 .as_ref()
414 .ok_or(Self::Error::IllegalState("No encrypted private key present".to_string()))?;
415
416 let key_pair = get_private_key_from_nep2(encrypted_private_key, password).map_err(|e| {
417 Self::Error::IllegalState(format!("Failed to decrypt private key: {e}"))
418 })?;
419
420 let key_pair_array = vec_to_array32(key_pair).map_err(|_| {
421 Self::Error::IllegalState("Failed to convert private key to 32-byte array".to_string())
422 })?;
423
424 self.key_pair =
425 Some(KeyPair::from_private_key(&key_pair_array).map_err(|e| {
426 Self::Error::IllegalState(format!("Failed to create key pair: {e}"))
427 })?);
428
429 Ok(())
430 }
431
432 fn encrypt_private_key(&mut self, password: &str) -> Result<(), Self::Error> {
433 let key_pair = self.key_pair.as_ref().ok_or(Self::Error::IllegalState(
434 "The account does not hold a decrypted private key.".to_string(),
435 ))?;
436
437 let encrypted_private_key = get_nep2_from_private_key(
438 &key_pair.private_key.to_raw_bytes().to_hex_string(),
439 password,
440 )
441 .map_err(|e| Self::Error::IllegalState(format!("Failed to encrypt private key: {e}")))?;
442
443 self.encrypted_private_key = Some(encrypted_private_key);
444 self.key_pair = None;
445 Ok(())
446 }
447
448 fn get_script_hash(&self) -> ScriptHash {
449 self.address_or_scripthash.script_hash()
450 }
451
452 fn get_signing_threshold(&self) -> Result<u32, Self::Error> {
453 self.signing_threshold.ok_or_else(|| {
454 Self::Error::IllegalState(format!(
455 "Cannot get signing threshold from account {}",
456 self.address_or_scripthash().address()
457 ))
458 })
459 }
460
461 fn get_nr_of_participants(&self) -> Result<u32, Self::Error> {
462 self.nr_of_participants.ok_or_else(|| {
463 Self::Error::IllegalState(format!(
464 "Cannot get signing threshold from account {}",
465 self.address_or_scripthash().address()
466 ))
467 })
468 }
469
470 fn from_verification_script(script: &VerificationScript) -> Result<Self, Self::Error> {
471 let address = ScriptHash::from_script(&script.script());
472
473 let (signing_threshold, nr_of_participants) = if script.is_multi_sig() {
474 (
475 Some(script.get_signing_threshold().unwrap()),
476 Some(script.get_nr_of_accounts().unwrap()),
477 )
478 } else {
479 (None, None)
480 };
481
482 Ok(Self {
483 address_or_scripthash: AddressOrScriptHash::ScriptHash(address),
484 label: Some(address.to_address()),
485 verification_script: Some(script.clone()),
486 signing_threshold: signing_threshold.map(|x| x as u32),
487 nr_of_participants: nr_of_participants.map(|x| x as u32),
488 ..Default::default()
489 })
490 }
491
492 fn from_public_key(public_key: &Secp256r1PublicKey) -> Result<Self, Self::Error> {
493 let script = VerificationScript::from_public_key(public_key);
494 let address = ScriptHash::from_script(&script.script());
495
496 Ok(Self {
497 address_or_scripthash: AddressOrScriptHash::ScriptHash(address),
498 label: Some(address.to_address()),
499 verification_script: Some(script),
500 ..Default::default()
501 })
502 }
503
504 fn set_wallet(&mut self, wallet: Option<Weak<Wallet>>) {
505 self.wallet = wallet;
506 }
507
508 fn get_wallet(&self) -> Option<Arc<Wallet>> {
509 self.wallet.as_ref().and_then(|w| w.upgrade())
510 }
511
512 fn multi_sig_from_public_keys(
513 public_keys: &mut [Secp256r1PublicKey],
514 signing_threshold: u32,
515 ) -> Result<Self, Self::Error> {
516 let script = VerificationScript::from_multi_sig(public_keys, signing_threshold as u8);
517 let addr = ScriptHash::from_script(&script.script());
518
519 Ok(Self {
520 label: Some(addr.to_address()),
521 verification_script: Some(script),
522 signing_threshold: Some(signing_threshold),
523 nr_of_participants: Some(public_keys.len() as u32),
524 address_or_scripthash: AddressOrScriptHash::ScriptHash(addr),
525 ..Default::default()
526 })
527 }
528
529 fn multi_sig_from_addr(
530 address: String,
531 signing_threshold: u8,
532 nr_of_participants: u8,
533 ) -> Result<Self, Self::Error> {
534 Ok(Self {
535 label: Option::from(address.clone()),
536 signing_threshold: Some(signing_threshold as u32),
537 nr_of_participants: Some(nr_of_participants as u32),
538 address_or_scripthash: AddressOrScriptHash::Address(address),
539 ..Default::default()
540 })
541 }
542
543 fn from_address(address: &str) -> Result<Self, Self::Error> {
544 let address = Address::from_str(address)
545 .map_err(|_| Self::Error::IllegalState(format!("Invalid address format: {address}")))?;
546
547 Ok(Self {
548 address_or_scripthash: AddressOrScriptHash::Address(address.clone()),
549 label: Some(address),
550 ..Default::default()
551 })
552 }
553
554 fn from_script_hash(script_hash: &H160) -> Result<Self, Self::Error> {
555 let address = script_hash.to_address();
556 Self::from_address(&address)
557 }
558
559 fn create() -> Result<Self, Self::Error> {
560 let key_pair = KeyPair::new_random();
561 Self::from_key_pair(key_pair, None, None)
562 }
563
564 fn is_multi_sig(&self) -> bool {
565 self.signing_threshold.is_some() && self.nr_of_participants.is_some()
566 }
567}
568
569impl PrehashSigner<Secp256r1Signature> for Account {
570 fn sign_prehash(&self, _prehash: &[u8]) -> Result<Secp256r1Signature, Error> {
571 let key_pair = self.key_pair.as_ref().ok_or_else(|| Error::new())?;
572
573 let signature = key_pair.private_key.sign_prehash(_prehash).map_err(|_| Error::new())?;
574
575 Ok(signature)
576 }
577}
578
579impl Account {
580 pub fn to_nep6_account(&self) -> Result<NEP6Account, ProviderError> {
581 if self.key_pair.is_some() && self.encrypted_private_key.is_none() {
582 return Err(ProviderError::IllegalState(
583 "Account private key is available but not encrypted.".to_string(),
584 ));
585 }
586
587 if self.verification_script.is_none() {
588 return Ok(NEP6Account::new(
589 self.address_or_scripthash.address().clone(),
590 self.label.clone(),
591 self.is_default,
592 self.is_locked,
593 self.encrypted_private_key.clone(),
594 None,
595 None,
596 ));
597 }
598
599 let mut parameters = Vec::new();
600 let script_data = self.verification_script.as_ref().unwrap();
601
602 if script_data.is_multi_sig() {
603 for i in 0..script_data.get_nr_of_accounts().unwrap() {
604 parameters.push(NEP6Parameter {
605 param_name: format!("signature{i}"),
606 param_type: ContractParameterType::Signature,
607 });
608 }
609 } else if script_data.is_single_sig() {
610 parameters.push(NEP6Parameter {
611 param_name: "signature".to_string(),
612 param_type: ContractParameterType::Signature,
613 });
614 }
615
616 let script_encoded = script_data.script().to_base64();
617 let contract = NEP6Contract {
618 script: Some(script_encoded),
619 is_deployed: false, nep6_parameters: parameters,
621 };
622
623 Ok(NEP6Account::new(
624 self.address_or_scripthash.address().clone(),
625 self.label.clone(),
626 self.is_default,
627 self.is_locked,
628 self.encrypted_private_key.clone(),
629 Some(contract),
630 None,
631 ))
632 }
633
634 pub async fn get_nep17_balances<P>(
635 &self,
636 provider: &RpcClient<P>,
637 ) -> Result<HashMap<H160, u64>, ProviderError>
638 where
639 P: JsonRpcProvider,
640 {
641 let response =
642 provider.get_nep17_balances(self.address_or_scripthash().script_hash()).await?;
643 let mut balances = HashMap::new();
644 for balance in response.balances {
645 let asset_hash = balance.asset_hash;
646 let amount = balance.amount.parse::<u64>().map_err(|e| {
647 ProviderError::CustomError(format!("Failed to parse balance amount: {e}"))
648 })?;
649 balances.insert(asset_hash, amount);
650 }
651 Ok(balances)
652 }
653}
654
655#[cfg(test)]
656mod tests {
657 use super::*;
658 use crate::{
659 neo_clients::{BodyRegexMatcher, HttpProvider, MockClient},
660 neo_config::TestConstants,
661 };
662
663 }