From 5f756888f5ec0041b8d2a2fad7e75225ec76148f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <bkchr@users.noreply.github.com>
Date: Wed, 23 Sep 2020 18:30:12 +0200
Subject: [PATCH] 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>
---
 client/cli/src/commands/utils.rs  | 12 +++++++---
 client/cli/src/commands/vanity.rs | 38 +++++++++++++++++++++++--------
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/client/cli/src/commands/utils.rs b/client/cli/src/commands/utils.rs
index a3298b222ad2b..8da71ec95540e 100644
--- a/client/cli/src/commands/utils.rs
+++ b/client/cli/src/commands/utils.rs
@@ -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.
 #[macro_export]
 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 {
 			$crate::CryptoScheme::Ecdsa => {
 				$method::<sp_core::ecdsa::Pair, $($generics),*>($($params),*)
diff --git a/client/cli/src/commands/vanity.rs b/client/cli/src/commands/vanity.rs
index e6f923f73c45c..33b9025c13fbc 100644
--- a/client/cli/src/commands/vanity.rs
+++ b/client/cli/src/commands/vanity.rs
@@ -22,10 +22,11 @@ use crate::{
 	error, utils, with_crypto_scheme,
 	CryptoSchemeFlag, NetworkSchemeFlag, OutputTypeFlag,
 };
-use sp_core::crypto::Ss58Codec;
+use sp_core::crypto::{Ss58Codec, Ss58AddressFormat};
 use structopt::StructOpt;
 use rand::{rngs::OsRng, RngCore};
 use sp_runtime::traits::IdentifyAccount;
+use utils::print_from_uri;
 
 /// The `vanity` command
 #[derive(Debug, StructOpt)]
@@ -54,23 +55,29 @@ pub struct VanityCmd {
 impl VanityCmd {
 	/// Run the command
 	pub fn run(&self) -> error::Result<()> {
-		let formated_seed = with_crypto_scheme!(self.crypto_scheme.scheme, generate_key(&self.pattern))?;
-		use utils::print_from_uri;
+		let formated_seed = with_crypto_scheme!(
+			self.crypto_scheme.scheme,
+			generate_key(&self.pattern, self.network_scheme.network.clone().unwrap_or_default()),
+		)?;
+
 		with_crypto_scheme!(
 			self.crypto_scheme.scheme,
 			print_from_uri(
 				&formated_seed,
 				None,
 				self.network_scheme.network.clone(),
-				self.output_scheme.output_type.clone()
-			)
+				self.output_scheme.output_type.clone(),
+			),
 		);
 		Ok(())
 	}
 }
 
 /// 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
 		Pair: sp_core::Pair,
 		Pair::Public: IdentifyAccount,
@@ -91,7 +98,7 @@ fn generate_key<Pair>(desired: &str) -> Result<String, &'static str>
 		}
 
 		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);
 		if score > best || desired.len() < 2 {
 			best = score;
@@ -171,13 +178,26 @@ mod tests {
 
 	#[test]
 	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!(
 			sr25519::Pair::from_seed_slice(&hex::decode(&seed[2..]).unwrap())
 				.unwrap()
 				.public()
 				.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]