Fix vanity command by taking the network into account (#7192)

* Fix vanity command by taking the network into account

* Delete empty line

* Apply suggestions from code review

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Change test

* Stupid me

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
This commit is contained in:
Bastian Köcher
2020-09-23 18:30:12 +02:00
committed by GitHub
parent 4e5f9fb72b
commit a00f0bf81c
2 changed files with 38 additions and 12 deletions
+9 -3
View File
@@ -215,10 +215,16 @@ pub fn read_message(msg: Option<&String>, should_decode: bool) -> Result<Vec<u8>
/// Allows for calling $method with appropriate crypto impl. /// Allows for calling $method with appropriate crypto impl.
#[macro_export] #[macro_export]
macro_rules! with_crypto_scheme { macro_rules! with_crypto_scheme {
($scheme:expr, $method:ident($($params:expr),*)) => { (
with_crypto_scheme!($scheme, $method<>($($params),*)) $scheme:expr,
$method:ident ( $($params:expr),* $(,)?) $(,)?
) => {
$crate::with_crypto_scheme!($scheme, $method<>($($params),*))
}; };
($scheme:expr, $method:ident<$($generics:ty),*>($($params:expr),*)) => { (
$scheme:expr,
$method:ident<$($generics:ty),*>( $( $params:expr ),* $(,)?) $(,)?
) => {
match $scheme { match $scheme {
$crate::CryptoScheme::Ecdsa => { $crate::CryptoScheme::Ecdsa => {
$method::<sp_core::ecdsa::Pair, $($generics),*>($($params),*) $method::<sp_core::ecdsa::Pair, $($generics),*>($($params),*)
+29 -9
View File
@@ -22,10 +22,11 @@ use crate::{
error, utils, with_crypto_scheme, error, utils, with_crypto_scheme,
CryptoSchemeFlag, NetworkSchemeFlag, OutputTypeFlag, CryptoSchemeFlag, NetworkSchemeFlag, OutputTypeFlag,
}; };
use sp_core::crypto::Ss58Codec; use sp_core::crypto::{Ss58Codec, Ss58AddressFormat};
use structopt::StructOpt; use structopt::StructOpt;
use rand::{rngs::OsRng, RngCore}; use rand::{rngs::OsRng, RngCore};
use sp_runtime::traits::IdentifyAccount; use sp_runtime::traits::IdentifyAccount;
use utils::print_from_uri;
/// The `vanity` command /// The `vanity` command
#[derive(Debug, StructOpt)] #[derive(Debug, StructOpt)]
@@ -54,23 +55,29 @@ pub struct VanityCmd {
impl VanityCmd { impl VanityCmd {
/// Run the command /// Run the command
pub fn run(&self) -> error::Result<()> { pub fn run(&self) -> error::Result<()> {
let formated_seed = with_crypto_scheme!(self.crypto_scheme.scheme, generate_key(&self.pattern))?; let formated_seed = with_crypto_scheme!(
use utils::print_from_uri; self.crypto_scheme.scheme,
generate_key(&self.pattern, self.network_scheme.network.clone().unwrap_or_default()),
)?;
with_crypto_scheme!( with_crypto_scheme!(
self.crypto_scheme.scheme, self.crypto_scheme.scheme,
print_from_uri( print_from_uri(
&formated_seed, &formated_seed,
None, None,
self.network_scheme.network.clone(), self.network_scheme.network.clone(),
self.output_scheme.output_type.clone() self.output_scheme.output_type.clone(),
) ),
); );
Ok(()) Ok(())
} }
} }
/// genertae a key based on given pattern /// genertae a key based on given pattern
fn generate_key<Pair>(desired: &str) -> Result<String, &'static str> fn generate_key<Pair>(
desired: &str,
network_override: Ss58AddressFormat,
) -> Result<String, &'static str>
where where
Pair: sp_core::Pair, Pair: sp_core::Pair,
Pair::Public: IdentifyAccount, Pair::Public: IdentifyAccount,
@@ -91,7 +98,7 @@ fn generate_key<Pair>(desired: &str) -> Result<String, &'static str>
} }
let p = Pair::from_seed(&seed); let p = Pair::from_seed(&seed);
let ss58 = p.public().into_account().to_ss58check(); let ss58 = p.public().into_account().to_ss58check_with_version(network_override);
let score = calculate_score(&desired, &ss58); let score = calculate_score(&desired, &ss58);
if score > best || desired.len() < 2 { if score > best || desired.len() < 2 {
best = score; best = score;
@@ -171,13 +178,26 @@ mod tests {
#[test] #[test]
fn test_generation_with_single_char() { fn test_generation_with_single_char() {
let seed = generate_key::<sr25519::Pair>("j").unwrap(); let seed = generate_key::<sr25519::Pair>("ab", Default::default()).unwrap();
assert!( assert!(
sr25519::Pair::from_seed_slice(&hex::decode(&seed[2..]).unwrap()) sr25519::Pair::from_seed_slice(&hex::decode(&seed[2..]).unwrap())
.unwrap() .unwrap()
.public() .public()
.to_ss58check() .to_ss58check()
.contains("j")); .contains("ab")
);
}
#[test]
fn generate_key_respects_network_override() {
let seed = generate_key::<sr25519::Pair>("ab", Ss58AddressFormat::PolkadotAccount).unwrap();
assert!(
sr25519::Pair::from_seed_slice(&hex::decode(&seed[2..]).unwrap())
.unwrap()
.public()
.to_ss58check_with_version(Ss58AddressFormat::PolkadotAccount)
.contains("ab")
);
} }
#[test] #[test]