keys.rs

use crate::*;
use clap::builder::OsStr;
use sr25519::Pair as Sr25519Pair;
use std::str::FromStr;

pub const SUBSTRATE_MNEMONIC: &str =
	"bottom drive obey lake curtain smoke basket hold race lonely fit walk";

/// secret format
#[derive(Clone, Copy, Debug, Eq, PartialEq, Default)]
pub enum SecretFormat {
	/// Raw 32B seed
	Seed,
	/// Substrate secret key or BIP39 mnemonic (optionally followed by derivation path)
	#[default]
	Substrate,
	/// Predefined (Alice, Bob, ...)
	Predefined,
	/// Cesium (scrypt + nacl)
	Cesium,
}
impl FromStr for SecretFormat {
	type Err = std::io::Error;

	fn from_str(s: &str) -> std::io::Result<Self> {
		match s {
			"seed" => Ok(SecretFormat::Seed),
			"substrate" => Ok(SecretFormat::Substrate),
			"predefined" => Ok(SecretFormat::Predefined),
			"cesium" => Ok(SecretFormat::Cesium),
			_ => Err(std::io::Error::from(std::io::ErrorKind::InvalidInput)),
		}
	}
}
impl From<SecretFormat> for &'static str {
	fn from(val: SecretFormat) -> &'static str {
		match val {
			SecretFormat::Seed => "seed",
			SecretFormat::Substrate => "substrate",
			SecretFormat::Predefined => "predefined",
			SecretFormat::Cesium => "cesium",
		}
	}
}
impl From<SecretFormat> for OsStr {
	fn from(val: SecretFormat) -> OsStr {
		OsStr::from(Into::<&str>::into(val))
	}
}

/// wrapper type for keys + signature
pub enum KeyPair {
	Sr25519(Sr25519Pair),
	Nacl(nacl::sign::Keypair),
}
impl KeyPair {
	pub fn address(&self) -> AccountId {
		match self {
			KeyPair::Sr25519(keypair) => keypair.public().into(),
			KeyPair::Nacl(keypair) => keypair.pkey.into(),
		}
	}
}
// can not derive clone because nacl does not implement it
impl Clone for KeyPair {
	fn clone(&self) -> Self {
		match self {
			KeyPair::Sr25519(keypair) => KeyPair::Sr25519(keypair.clone()),
			KeyPair::Nacl(keypair) => KeyPair::Nacl(nacl::sign::Keypair {
				skey: keypair.skey,
				pkey: keypair.pkey,
			}),
		}
	}
}
impl From<Sr25519Pair> for KeyPair {
	fn from(pair: Sr25519Pair) -> KeyPair {
		KeyPair::Sr25519(pair)
	}
}
impl From<nacl::sign::Keypair> for KeyPair {
	fn from(pair: nacl::sign::Keypair) -> KeyPair {
		KeyPair::Nacl(pair)
	}
}
pub enum Signature {
	Sr25519(sr25519::Signature),
	Nacl(Vec<u8>),
}

/// get keypair in any possible way
/// at this point, if secret is predefined, it's not replaced yet
pub fn get_keypair(
	secret_format: SecretFormat,
	secret: Option<&str>,
) -> Result<KeyPair, GcliError> {
	match (secret_format, secret) {
		(SecretFormat::Cesium, None) => Ok(prompt_secret(SecretFormat::Cesium)),
		(SecretFormat::Predefined, Some(deriv)) => pair_from_predefined(deriv).map(|v| v.into()),
		(_, Some(secret)) => Ok(pair_from_secret(secret_format, secret)?.into()),
		_ => Err(GcliError::Logic(
			"can not get keypair from available options".to_string(),
		)),
	}
}

/// get keypair from given secret
/// if secret is predefined, secret should contain the predefined value
pub fn pair_from_secret(
	secret_format: SecretFormat,
	secret: &str,
) -> Result<Sr25519Pair, GcliError> {
	match secret_format {
		SecretFormat::Substrate => pair_from_str(secret),
		SecretFormat::Predefined => pair_from_str(secret), /* if predefined, secret arg is replaced in config */
		SecretFormat::Seed => pair_from_seed(secret),
		SecretFormat::Cesium => Err(GcliError::Logic(
			"cesium format incompatible with single secret".to_string(),
		)),
	}
}

/// get keypair from given string secret
pub fn pair_from_str(secret: &str) -> Result<Sr25519Pair, GcliError> {
	Sr25519Pair::from_string(secret, None)
		.map_err(|_| GcliError::Input("Invalid secret".to_string()))
}

/// get keypair from given seed
pub fn pair_from_seed(secret: &str) -> Result<Sr25519Pair, GcliError> {
	let mut seed = [0; 32];
	hex::decode_to_slice(secret, &mut seed)
		.map_err(|_| GcliError::Input("Invalid secret".to_string()))?;
	let pair = Sr25519Pair::from_seed(&seed);
	Ok(pair)
}

/// get mnemonic from predefined derivation path
pub fn predefined_mnemonic(deriv: &str) -> String {
	format!("{SUBSTRATE_MNEMONIC}//{deriv}")
}

/// get keypair from predefined secret
pub fn pair_from_predefined(deriv: &str) -> Result<Sr25519Pair, GcliError> {
	pair_from_str(&predefined_mnemonic(deriv))
}

/// get keypair from Cesium id/pwd
pub fn pair_from_cesium(id: String, pwd: String) -> nacl::sign::Keypair {
	let params = scrypt::Params::new(12u8, 16u32, 1u32, 32).unwrap();
	let seed = &mut [0u8; 32];
	scrypt::scrypt(&pwd.into_bytes(), &id.into_bytes(), &params, seed).unwrap();
	nacl::sign::generate_keypair(seed)
}

/// ask user to input a secret
pub fn prompt_secret_substrate() -> Sr25519Pair {
	loop {
		let mnemonic = &rpassword::prompt_password("Mnemonic: ").unwrap();
		match pair_from_str(mnemonic) {
			Ok(pair) => return pair,
			Err(_) => println!("Invalid secret"),
		}
	}
}

/// ask user pass (Cesium format)
pub fn prompt_secret_cesium() -> nacl::sign::Keypair {
	let id = rpassword::prompt_password("Cesium id: ").unwrap();
	let pwd = rpassword::prompt_password("Cesium password: ").unwrap();
	pair_from_cesium(id, pwd)
}

/// ask user to input a seed
pub fn prompt_seed() -> Sr25519Pair {
	loop {
		let seed = &rpassword::prompt_password("Seed: ").unwrap();
		match pair_from_seed(seed) {
			Ok(pair) => return pair,
			Err(_) => println!("Invalid seed"),
		}
	}
}

/// ask user pass (Cesium format)
pub fn prompt_predefined() -> Sr25519Pair {
	let deriv = rpassword::prompt_password("Enter derivation path: ").unwrap();
	pair_from_predefined(&deriv).expect("invalid secret")
}

/// ask user secret in relevant format
pub fn prompt_secret(secret_format: SecretFormat) -> KeyPair {
	match secret_format {
		SecretFormat::Substrate => prompt_secret_substrate().into(),
		SecretFormat::Cesium => prompt_secret_cesium().into(),
		SecretFormat::Seed => prompt_seed().into(),
		SecretFormat::Predefined => prompt_predefined().into(),
	}
}