1use std::{
6 collections::{BTreeMap, HashSet},
7 fmt::{Display, Formatter, Write},
8 fs,
9 fs::File,
10 io::BufReader,
11 path::{Path, PathBuf},
12};
13
14use anyhow::{Context, anyhow, bail, ensure};
15use bip32::DerivationPath;
16use bip39::{Language, Mnemonic, Seed};
17use iota_types::{
18 base_types::IotaAddress,
19 crypto::{
20 EncodeDecodeBase64, IotaKeyPair, PublicKey, Signature, SignatureScheme, enum_dispatch,
21 get_key_pair_from_rng,
22 },
23};
24use rand::{SeedableRng, rngs::StdRng};
25use regex::Regex;
26use serde::{Deserialize, Deserializer, Serialize, Serializer};
27use serde_with::{DisplayFromStr, serde_as};
28use shared_crypto::intent::{Intent, IntentMessage};
29use tracing::{debug, info};
30
31use crate::{
32 key_derive::{derive_key_pair_from_path, generate_new_key},
33 random_names::{random_name, random_names},
34 serde_iota_keypair, serde_public_key,
35};
36
37#[derive(Serialize, Deserialize)]
38#[enum_dispatch(AccountKeystore)]
39pub enum Keystore {
40 File(FileBasedKeystore),
41 InMem(InMemKeystore),
42}
43
44#[enum_dispatch]
45pub trait AccountKeystore: Send + Sync {
46 fn add_key(
47 &mut self,
48 alias: Option<String>,
49 key: impl Into<StoredKey>,
50 ) -> Result<(), anyhow::Error>;
51 fn remove_key(&mut self, address: &IotaAddress) -> Result<(), anyhow::Error>;
52 fn keys(&self) -> Vec<&StoredKey>;
53 fn get_key(&self, address: &IotaAddress) -> Result<&StoredKey, anyhow::Error>;
54
55 fn sign_hashed(&self, address: &IotaAddress, msg: &[u8])
56 -> Result<Signature, signature::Error>;
57
58 fn sign_secure<T>(
59 &self,
60 address: &IotaAddress,
61 msg: &T,
62 intent: Intent,
63 ) -> Result<Signature, signature::Error>
64 where
65 T: Serialize;
66 fn addresses(&self) -> Vec<IotaAddress> {
67 self.keys().into_iter().map(|k| k.address()).collect()
68 }
69 fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)>;
70 fn aliases(&self) -> Vec<&Alias>;
71 fn aliases_mut(&mut self) -> Vec<&mut Alias>;
72 fn alias_names(&self) -> Vec<&str> {
73 self.aliases()
74 .into_iter()
75 .map(|a| a.alias.as_str())
76 .collect()
77 }
78 fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error>;
80 fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error>;
81 fn alias_exists(&self, alias: &str) -> bool {
83 self.alias_names().contains(&alias)
84 }
85
86 fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error>;
87
88 fn update_alias(
89 &mut self,
90 old_alias: &str,
91 new_alias: Option<&str>,
92 ) -> Result<String, anyhow::Error>;
93
94 fn update_alias_value(
96 &mut self,
97 old_alias: &str,
98 new_alias: Option<&str>,
99 ) -> Result<String, anyhow::Error> {
100 if !self.alias_exists(old_alias) {
101 bail!("The provided alias {old_alias} does not exist");
102 }
103
104 let new_alias_name = self.create_alias(new_alias.map(str::to_string))?;
105
106 for a in self.aliases_mut() {
107 if a.alias == old_alias {
108 let pk = &a.public_key_base64;
109 *a = Alias {
110 alias: new_alias_name.clone(),
111 public_key_base64: pk.clone(),
112 };
113 }
114 }
115 Ok(new_alias_name)
116 }
117
118 fn generate_and_add_new_key(
119 &mut self,
120 key_scheme: SignatureScheme,
121 alias: Option<String>,
122 derivation_path: Option<DerivationPath>,
123 word_length: Option<String>,
124 ) -> Result<(IotaAddress, String, SignatureScheme), anyhow::Error> {
125 let (address, kp, scheme, phrase) =
126 generate_new_key(key_scheme, derivation_path, word_length)?;
127 self.add_key(alias, kp)?;
128 Ok((address, phrase, scheme))
129 }
130
131 fn import_from_mnemonic(
132 &mut self,
133 phrase: &str,
134 key_scheme: SignatureScheme,
135 derivation_path: Option<DerivationPath>,
136 alias: Option<String>,
137 ) -> Result<IotaAddress, anyhow::Error> {
138 let mnemonic = Mnemonic::from_phrase(phrase, Language::English)
139 .map_err(|e| anyhow::anyhow!("Invalid mnemonic phrase: {:?}", e))?;
140 let seed = Seed::new(&mnemonic, "");
141 self.import_from_seed(seed.as_bytes(), key_scheme, derivation_path, alias)
142 }
143
144 fn import_from_seed(
145 &mut self,
146 seed: &[u8],
147 key_scheme: SignatureScheme,
148 derivation_path: Option<DerivationPath>,
149 alias: Option<String>,
150 ) -> Result<IotaAddress, anyhow::Error> {
151 match derive_key_pair_from_path(seed, derivation_path, &key_scheme) {
152 Ok((address, kp)) => {
153 self.add_key(alias, kp)?;
154 Ok(address)
155 }
156 Err(e) => Err(anyhow!("error getting keypair {:?}", e)),
157 }
158 }
159
160 fn import_from_external(
161 &mut self,
162 source: &str,
163 public_key: PublicKey,
164 derivation_path: Option<DerivationPath>,
165 alias: Option<String>,
166 ) -> Result<(), anyhow::Error> {
167 self.add_key(
168 alias,
169 StoredKey::External {
170 derivation_path,
171 public_key,
172 source: source.to_string(),
173 },
174 )
175 }
176}
177
178impl Display for Keystore {
179 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
180 let mut writer = String::new();
181 match self {
182 Keystore::File(file) => {
183 writeln!(writer, "Keystore Type: File")?;
184 write!(writer, "Keystore Path : {:?}", file.path)?;
185 write!(f, "{writer}")
186 }
187 Keystore::InMem(_) => {
188 writeln!(writer, "Keystore Type: InMem")?;
189 write!(f, "{writer}")
190 }
191 }
192 }
193}
194
195#[derive(Serialize, Deserialize, Clone, Debug)]
196pub struct Alias {
197 pub alias: String,
198 pub public_key_base64: String,
199}
200
201#[expect(clippy::large_enum_variant)]
202#[serde_as]
203#[derive(Serialize, Deserialize, Debug, Clone)]
204#[serde(
205 tag = "type", content = "value", rename_all = "snake_case"
208)]
209pub enum StoredKey {
210 #[serde(with = "serde_iota_keypair")]
211 KeyPair(IotaKeyPair),
212 External {
213 source: String,
214 #[serde_as(as = "Option<DisplayFromStr>")]
215 #[serde(skip_serializing_if = "Option::is_none")]
216 derivation_path: Option<DerivationPath>,
217 #[serde(rename = "public_key_base64_with_flag", with = "serde_public_key")]
218 public_key: PublicKey,
219 },
220}
221
222impl From<IotaKeyPair> for StoredKey {
223 fn from(keypair: IotaKeyPair) -> Self {
224 StoredKey::KeyPair(keypair)
225 }
226}
227
228impl StoredKey {
229 pub fn address(&self) -> IotaAddress {
230 IotaAddress::from(&self.public())
231 }
232
233 pub fn public(&self) -> PublicKey {
234 match self {
235 StoredKey::KeyPair(keypair) => keypair.public(),
236 StoredKey::External { public_key, .. } => public_key.clone(),
237 }
238 }
239
240 pub fn derivation_path(&self) -> Option<DerivationPath> {
241 match self {
242 StoredKey::KeyPair(_) => None,
243 StoredKey::External {
244 derivation_path, ..
245 } => derivation_path.clone(),
246 }
247 }
248
249 pub fn external_source(&self) -> Option<String> {
250 match self {
251 StoredKey::KeyPair(_) => None,
252 StoredKey::External { source, .. } => Some(source.clone()),
253 }
254 }
255
256 pub fn as_keypair(&self) -> Result<&IotaKeyPair, anyhow::Error> {
257 match self {
258 StoredKey::KeyPair(keypair) => Ok(keypair),
259 StoredKey::External { .. } => bail!("Cannot get key pair for External keys."),
260 }
261 }
262}
263
264#[derive(Default)]
265pub struct FileBasedKeystore {
266 keys: BTreeMap<IotaAddress, StoredKey>,
267 aliases: BTreeMap<IotaAddress, Alias>,
268 path: PathBuf,
269}
270
271impl Serialize for FileBasedKeystore {
272 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
273 where
274 S: Serializer,
275 {
276 serializer.serialize_str(self.path.to_str().unwrap_or(""))
277 }
278}
279
280impl<'de> Deserialize<'de> for FileBasedKeystore {
281 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
282 where
283 D: Deserializer<'de>,
284 {
285 use serde::de::Error;
286 FileBasedKeystore::new(&PathBuf::from(String::deserialize(deserializer)?))
287 .map_err(D::Error::custom)
288 }
289}
290#[derive(Serialize, Deserialize, Debug)]
291pub struct FileBasedKeystoreFile {
292 pub version: u8,
293 pub keys: Vec<AliasedKey>,
294}
295
296#[derive(Serialize, Deserialize, Debug)]
297pub struct AliasedKey {
298 pub alias: String,
299 pub address: IotaAddress,
300 pub key: StoredKey,
301}
302
303impl AccountKeystore for FileBasedKeystore {
304 fn sign_hashed(
305 &self,
306 address: &IotaAddress,
307 msg: &[u8],
308 ) -> Result<Signature, signature::Error> {
309 let stored_key = self.keys.get(address).ok_or_else(|| {
310 signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
311 })?;
312
313 match stored_key {
314 StoredKey::KeyPair(keypair) => Ok(Signature::new_hashed(msg, keypair)),
315 StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
316 "sign_hashed is not supported for external type: {source} [{address}]"
317 ))),
318 }
319 }
320 fn sign_secure<T>(
321 &self,
322 address: &IotaAddress,
323 msg: &T,
324 intent: Intent,
325 ) -> Result<Signature, signature::Error>
326 where
327 T: Serialize,
328 {
329 let stored_key = self.keys.get(address).ok_or_else(|| {
330 signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
331 })?;
332
333 let intent_msg = &IntentMessage::new(intent, msg);
334 match stored_key {
335 StoredKey::KeyPair(keypair) => Ok(Signature::new_secure(intent_msg, keypair)),
336 StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
337 "sign_secure is not supported for external type: {source} [{address}]",
338 ))),
339 }
340 }
341
342 fn add_key(
343 &mut self,
344 alias: Option<String>,
345 key: impl Into<StoredKey>,
346 ) -> Result<(), anyhow::Error> {
347 let key = key.into();
348 let address: IotaAddress = (&key.public()).into();
349 let alias = self.create_alias(alias)?;
350 self.aliases.insert(
351 address,
352 Alias {
353 alias,
354 public_key_base64: key.public().encode_base64(),
355 },
356 );
357 self.keys.insert(address, key);
358 self.save()?;
359 Ok(())
360 }
361
362 fn remove_key(&mut self, address: &IotaAddress) -> Result<(), anyhow::Error> {
363 self.aliases.remove(address);
364 self.keys.remove(address);
365 self.save()?;
366 Ok(())
367 }
368
369 fn aliases(&self) -> Vec<&Alias> {
372 self.aliases.values().collect()
373 }
374
375 fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)> {
376 self.aliases.iter().collect::<Vec<_>>()
377 }
378
379 fn aliases_mut(&mut self) -> Vec<&mut Alias> {
382 self.aliases.values_mut().collect()
383 }
384
385 fn keys(&self) -> Vec<&StoredKey> {
386 self.keys.values().collect()
387 }
388
389 fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error> {
393 match alias {
394 Some(a) if self.alias_exists(&a) => {
395 bail!("Alias {a} already exists. Please choose another alias.")
396 }
397 Some(a) => validate_alias(&a),
398 None => Ok(random_name(
399 &self
400 .alias_names()
401 .into_iter()
402 .map(|x| x.to_string())
403 .collect::<HashSet<_>>(),
404 )),
405 }
406 }
407
408 fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error> {
410 self.addresses_with_alias()
411 .iter()
412 .find(|x| x.1.alias == alias)
413 .ok_or_else(|| anyhow!("Cannot resolve alias {alias} to an address"))
414 .map(|x| x.0)
415 }
416
417 fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error> {
419 match self.aliases.get(address) {
420 Some(alias) => Ok(alias.alias.clone()),
421 None => bail!("Cannot find alias for address {address}"),
422 }
423 }
424
425 fn get_key(&self, address: &IotaAddress) -> Result<&StoredKey, anyhow::Error> {
426 match self.keys.get(address) {
427 Some(key) => Ok(key),
428 None => Err(anyhow!("Cannot find key for address: [{address}]")),
429 }
430 }
431
432 fn update_alias(
435 &mut self,
436 old_alias: &str,
437 new_alias: Option<&str>,
438 ) -> Result<String, anyhow::Error> {
439 let new_alias_name = self.update_alias_value(old_alias, new_alias)?;
440 self.save()?;
441 Ok(new_alias_name)
442 }
443}
444
445impl FileBasedKeystore {
446 pub fn new_from_v1(path: &PathBuf) -> Result<Self, anyhow::Error> {
447 let keys = if path.exists() {
448 let reader =
449 BufReader::new(File::open(path).with_context(|| {
450 format!("Cannot open the keystore file: {}", path.display())
451 })?);
452 let kp_strings: Vec<String> = serde_json::from_reader(reader).with_context(|| {
453 format!("Cannot deserialize the keystore file: {}", path.display(),)
454 })?;
455 kp_strings
456 .iter()
457 .map(|kpstr| {
458 let key = IotaKeyPair::decode(kpstr);
459 key.map(|k| (IotaAddress::from(&k.public()), StoredKey::KeyPair(k)))
460 })
461 .collect::<Result<BTreeMap<_, _>, _>>()
462 .map_err(|e| anyhow!("Invalid keystore file: {}. {}", path.display(), e))?
463 } else {
464 BTreeMap::new()
465 };
466
467 let mut aliases_path = path.clone();
469 aliases_path.set_extension("aliases");
470
471 let aliases = if aliases_path.exists() {
472 let reader = BufReader::new(File::open(&aliases_path).with_context(|| {
473 format!(
474 "Cannot open aliases file in keystore: {}",
475 aliases_path.display()
476 )
477 })?);
478
479 let aliases: Vec<Alias> = serde_json::from_reader(reader).with_context(|| {
480 format!(
481 "Cannot deserialize aliases file in keystore: {}",
482 aliases_path.display(),
483 )
484 })?;
485
486 aliases
487 .into_iter()
488 .map(|alias| {
489 let key = PublicKey::decode_base64(&alias.public_key_base64);
490 key.map(|k| (Into::<IotaAddress>::into(&k), alias))
491 })
492 .collect::<Result<BTreeMap<_, _>, _>>()
493 .map_err(|e| {
494 anyhow!(
495 "Invalid aliases file in keystore: {}. {}",
496 aliases_path.display(),
497 e
498 )
499 })?
500 } else if keys.is_empty() {
501 BTreeMap::new()
502 } else {
503 let names: Vec<String> = random_names(HashSet::new(), keys.len());
504 let aliases = keys
505 .iter()
506 .zip(names)
507 .map(|((iota_address, ikp), alias)| {
508 let public_key_base64 = ikp.public().encode_base64();
509 (
510 *iota_address,
511 Alias {
512 alias,
513 public_key_base64,
514 },
515 )
516 })
517 .collect::<BTreeMap<_, _>>();
518 let aliases_store = serde_json::to_string_pretty(&aliases.values().collect::<Vec<_>>())
519 .with_context(|| {
520 format!(
521 "Cannot serialize aliases to file in keystore: {}",
522 aliases_path.display()
523 )
524 })?;
525 fs::write(aliases_path, aliases_store)?;
526 aliases
527 };
528
529 Ok(Self {
530 keys,
531 aliases,
532 path: path.to_path_buf(),
533 })
534 }
535
536 fn needs_migration(path: &PathBuf) -> Result<bool, anyhow::Error> {
537 let mut aliases_path = path.clone();
538 aliases_path.set_extension("aliases");
539 if aliases_path.exists() {
540 debug!(
542 "An alias file exists at {}, assuming keystore is in v1 format",
543 aliases_path.display()
544 );
545 return Ok(true);
546 }
547 if path.exists() {
550 let reader =
551 BufReader::new(File::open(path).with_context(|| {
552 format!("Cannot open the keystore file: {}", path.display())
553 })?);
554 let is_v1_format = serde_json::from_reader::<_, Vec<String>>(reader).is_ok();
557 debug!(
558 "Keystore file at {} is in v1 format: {}",
559 path.display(),
560 is_v1_format,
561 );
562 Ok(is_v1_format)
563 } else {
564 Ok(false)
566 }
567 }
568
569 fn migrate_v1_to_v2(path: &PathBuf) -> Result<Self, anyhow::Error> {
570 let migrated = Self::new_from_v1(path)?;
571 let mut backup_path = path.clone();
574 backup_path.set_extension(
575 path.extension()
578 .and_then(|ext| Some(ext.to_str()?.to_owned() + ".migrated"))
579 .unwrap_or(String::from("migrated")),
580 );
581 fs::rename(path, &backup_path).with_context(|| {
582 format!(
583 "Failed to rename the old keystore file to {}",
584 backup_path.display()
585 )
586 })?;
587 let mut aliases_path = path.clone();
588 aliases_path.set_extension("aliases");
589 let mut backup_aliases_path = aliases_path.clone();
590 backup_aliases_path.set_extension("aliases.migrated");
591 fs::rename(&aliases_path, &backup_aliases_path).with_context(|| {
592 format!(
593 "Failed to rename the old aliases file to {}",
594 backup_aliases_path.display()
595 )
596 })?;
597
598 info!(
599 "Migrated {} keys in keystore from v1 to v2 format. Old files have been renamed to {} and {}",
600 migrated.keys.len(),
601 backup_path.display(),
602 backup_aliases_path.display()
603 );
604
605 migrated.save()?;
607
608 Ok(migrated)
609 }
610
611 pub fn new(path: &PathBuf) -> Result<Self, anyhow::Error> {
612 if Self::needs_migration(path)? {
613 return Self::migrate_v1_to_v2(path);
614 }
615
616 let (keys, aliases) = if path.exists() {
617 let reader =
618 BufReader::new(File::open(path).with_context(|| {
619 format!("Cannot open the keystore file: {}", path.display())
620 })?);
621
622 let file: FileBasedKeystoreFile = serde_json::from_reader(reader)
623 .with_context(|| {
624 format!("Cannot deserialize the keystore file: {}", path.display(),)
625 })
626 .map_err(|e| anyhow!("Invalid keystore file: {}. {}", path.display(), e))?;
627
628 let aliases = file
629 .keys
630 .iter()
631 .map(|aliased| {
632 let public_key = aliased.key.public();
633 (
634 IotaAddress::from(&public_key),
635 Alias {
636 alias: aliased.alias.clone(),
637 public_key_base64: public_key.encode_base64(),
638 },
639 )
640 })
641 .collect::<BTreeMap<_, _>>();
642
643 let keys = file
644 .keys
645 .into_iter()
646 .map(|aliased| (IotaAddress::from(&aliased.key.public()), aliased.key))
647 .collect::<BTreeMap<_, _>>();
648
649 (keys, aliases)
650 } else {
651 (BTreeMap::new(), BTreeMap::new())
652 };
653
654 Ok(Self {
655 keys,
656 aliases,
657 path: path.to_path_buf(),
658 })
659 }
660
661 pub fn set_path(&mut self, path: &Path) {
662 self.path = path.to_path_buf();
663 }
664
665 pub fn save(&self) -> Result<(), anyhow::Error> {
666 let file = FileBasedKeystoreFile {
667 version: 2,
668 keys: self
669 .keys
670 .iter()
671 .map(|(address, key)| AliasedKey {
672 alias: self
673 .aliases
674 .get(address)
675 .map_or_else(|| self.create_alias(None).unwrap(), |a| a.alias.clone()),
676 address: *address,
677 key: key.clone(),
678 })
679 .collect(),
680 };
681
682 let store = serde_json::to_string_pretty(&file).with_context(|| {
683 format!("Cannot serialize keystore to file: {}", self.path.display())
684 })?;
685 fs::write(&self.path, store)
686 .map_err(|e| anyhow!("Couldn't save keystore to {}: {e}", self.path.display()))?;
687 Ok(())
688 }
689}
690
691#[derive(Default, Serialize, Deserialize)]
692pub struct InMemKeystore {
693 aliases: BTreeMap<IotaAddress, Alias>,
694 keys: BTreeMap<IotaAddress, StoredKey>,
695}
696
697impl AccountKeystore for InMemKeystore {
698 fn sign_hashed(
699 &self,
700 address: &IotaAddress,
701 msg: &[u8],
702 ) -> Result<Signature, signature::Error> {
703 let stored_key = self.keys.get(address).ok_or_else(|| {
704 signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
705 })?;
706
707 match stored_key {
708 StoredKey::KeyPair(keypair) => Ok(Signature::new_hashed(msg, keypair)),
709 StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
710 "sign_hashed is not supported for external type: {source} [{address}]"
711 ))),
712 }
713 }
714 fn sign_secure<T>(
715 &self,
716 address: &IotaAddress,
717 msg: &T,
718 intent: Intent,
719 ) -> Result<Signature, signature::Error>
720 where
721 T: Serialize,
722 {
723 let stored_key = self.keys.get(address).ok_or_else(|| {
724 signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
725 })?;
726
727 let intent_msg = &IntentMessage::new(intent, msg);
728 match stored_key {
729 StoredKey::KeyPair(keypair) => Ok(Signature::new_secure(intent_msg, keypair)),
730 StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
731 "sign_secure is not supported for external type: {source} [{address}]",
732 ))),
733 }
734 }
735
736 fn add_key(
737 &mut self,
738 alias: Option<String>,
739 key: impl Into<StoredKey>,
740 ) -> Result<(), anyhow::Error> {
741 let key = key.into();
742 let address: IotaAddress = (&key.public()).into();
743 let alias = alias.unwrap_or_else(|| {
744 random_name(
745 &self
746 .aliases()
747 .iter()
748 .map(|x| x.alias.clone())
749 .collect::<HashSet<_>>(),
750 )
751 });
752
753 let public_key_base64 = key.public().encode_base64();
754 let alias = Alias {
755 alias,
756 public_key_base64,
757 };
758 self.aliases.insert(address, alias);
759 self.keys.insert(address, key);
760 Ok(())
761 }
762
763 fn remove_key(&mut self, address: &IotaAddress) -> Result<(), anyhow::Error> {
764 self.aliases.remove(address);
765 self.keys.remove(address);
766 Ok(())
767 }
768
769 fn aliases(&self) -> Vec<&Alias> {
771 self.aliases.values().collect()
772 }
773
774 fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)> {
775 self.aliases.iter().collect::<Vec<_>>()
776 }
777
778 fn keys(&self) -> Vec<&StoredKey> {
779 self.keys.values().collect()
780 }
781
782 fn get_key(&self, address: &IotaAddress) -> Result<&StoredKey, anyhow::Error> {
783 match self.keys.get(address) {
784 Some(key) => Ok(key),
785 None => Err(anyhow!("Cannot find key for address: [{address}]")),
786 }
787 }
788
789 fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error> {
791 match self.aliases.get(address) {
792 Some(alias) => Ok(alias.alias.clone()),
793 None => bail!("Cannot find alias for address {address}"),
794 }
795 }
796
797 fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error> {
799 self.addresses_with_alias()
800 .iter()
801 .find(|x| x.1.alias == alias)
802 .ok_or_else(|| anyhow!("Cannot resolve alias {alias} to an address"))
803 .map(|x| x.0)
804 }
805
806 fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error> {
810 match alias {
811 Some(a) if self.alias_exists(&a) => {
812 bail!("Alias {a} already exists. Please choose another alias.")
813 }
814 Some(a) => validate_alias(&a),
815 None => Ok(random_name(
816 &self
817 .alias_names()
818 .into_iter()
819 .map(|x| x.to_string())
820 .collect::<HashSet<_>>(),
821 )),
822 }
823 }
824
825 fn aliases_mut(&mut self) -> Vec<&mut Alias> {
826 self.aliases.values_mut().collect()
827 }
828
829 fn update_alias(
832 &mut self,
833 old_alias: &str,
834 new_alias: Option<&str>,
835 ) -> Result<String, anyhow::Error> {
836 self.update_alias_value(old_alias, new_alias)
837 }
838}
839
840impl InMemKeystore {
841 pub fn new_insecure_for_tests(initial_key_number: usize) -> Self {
842 let mut rng = StdRng::from_seed([0; 32]);
843 let keys = (0..initial_key_number)
844 .map(|_| get_key_pair_from_rng(&mut rng))
845 .map(|(ad, k)| (ad, IotaKeyPair::Ed25519(k).into()))
846 .collect::<BTreeMap<IotaAddress, StoredKey>>();
847
848 let aliases = keys
849 .iter()
850 .zip(random_names(HashSet::new(), keys.len()))
851 .map(|((iota_address, ikp), alias)| {
852 let public_key_base64 = ikp.public().encode_base64();
853 (
854 *iota_address,
855 Alias {
856 alias,
857 public_key_base64,
858 },
859 )
860 })
861 .collect::<BTreeMap<_, _>>();
862
863 Self { aliases, keys }
864 }
865}
866
867fn validate_alias(alias: &str) -> Result<String, anyhow::Error> {
868 let re = Regex::new(r"^[A-Za-z][A-Za-z0-9-_\.]*$")
869 .map_err(|_| anyhow!("Cannot build the regex needed to validate the alias naming"))?;
870 let alias = alias.trim();
871 ensure!(
872 re.is_match(alias),
873 "Invalid alias. A valid alias must start with a letter and can contain only letters, digits, hyphens (-), dots (.), or underscores (_)."
874 );
875 Ok(alias.to_string())
876}
877
878#[cfg(test)]
879mod tests {
880 use crate::keystore::validate_alias;
881
882 #[test]
883 fn validate_alias_test() {
884 assert!(validate_alias("A.B_dash").is_ok());
886 assert!(validate_alias("A.B-C1_dash").is_ok());
887 assert!(validate_alias("abc_123.iota").is_ok());
888 assert!(validate_alias("A.B-C_dash!").is_err());
890 assert!(validate_alias(".B-C_dash!").is_err());
891 assert!(validate_alias("_test").is_err());
892 assert!(validate_alias("123").is_err());
893 assert!(validate_alias("@@123").is_err());
894 assert!(validate_alias("@_Ab").is_err());
895 assert!(validate_alias("_Ab").is_err());
896 assert!(validate_alias("^A").is_err());
897 assert!(validate_alias("-A").is_err());
898 }
899}