Nicolas80 · 18e061cd · 6d1be328 · 9f712d7b · 97abd338 · f740fc65
--- a/src/commands/vault.rs

+ 327

− 33
+++ b/src/commands/vault.rs

+ 327

− 33
+use crate::entities::vault_user;
 use crate::*;
 use age::secrecy::Secret;
+use comfy_table::{Cell, Table};
+use sea_orm::ActiveValue::Set;
+use sea_orm::PaginatorTrait;
+use sea_orm::{ActiveModelTrait, EntityTrait};
+use sea_orm::{ColumnTrait, QueryFilter};
+use std::collections::BTreeMap;
 use std::io::{Read, Write};
+use std::path::PathBuf;
 /// define universal dividends subcommands
 #[derive(Clone, Default, Debug, clap::Parser)]
 pub enum Subcommand {
-	#[default]
 	/// List available keys
+	#[default]
 	List,
-	/// Show where vault stores secret
+	/// Use specific vault key (changes the config address)
-	Where,
+	Use {
+		#[clap(flatten)]
+		address_or_vault_name: AddressOrVaultNameGroup,
+	},
 	/// Generate a mnemonic
 	Generate,
-	/// Import mnemonic with interactive prompt
+	/// Import key from (substrate)mnemonic or other format with interactive prompt
-	Import,
+	Import {
+		/// Secret key format (substrate, seed, cesium)
+		#[clap(short = 'S', long, required = false, default_value = SecretFormat::Substrate)]
+		secret_format: SecretFormat,
+	},
+	/// Give a meaningful vault name to an account address
+	Rename {
+		/// Account Address
+		address: AccountId,
+		/// Vault name for that account
+		name: String,
+	},
+	/// Remove a vault key and potentially associated vault name
+	Remove {
+		#[clap(flatten)]
+		address_or_vault_name: AddressOrVaultNameGroup,
+	},
+	/// Show where vault stores secret
+	Where,
+}
+struct VaultItem {
+	secret_format: SecretFormat,
+	secret: keys::Secret,
+	address: AccountId,
 }
 // encrypt input with passphrase
 @@ -43,65 +78,324 @@ fn decrypt(input: &[u8], passphrase: String) -> Result<Vec<u8>, age::DecryptErro
 @@ -43,65 +78,324 @@ fn decrypt(input: &[u8], passphrase: String) -> Result<Vec<u8>, age::DecryptErro
 }
 /// handle ud commands
-pub fn handle_command(data: Data, command: Subcommand) -> Result<(), GcliError> {
+pub async fn handle_command(data: Data, command: Subcommand) -> Result<(), GcliError> {
 	// match subcommand
 	match command {
 		Subcommand::List => {
-			if let Ok(entries) = std::fs::read_dir(data.project_dir.data_dir()) {
+			let vault_key_addresses = fetch_vault_key_addresses(&data).await?;
-				println!("available keys:");
-				entries.for_each(|e| println!("{}", e.unwrap().file_name().to_str().unwrap()));
+			let table = compute_vault_key_table(&data, &vault_key_addresses).await?;
-			} else {
-				println!("could not read project dir");
+			println!("available keys:");
-			}
+			println!("{table}");
 		}
-		Subcommand::Where => {
+		Subcommand::Use {
-			println!("{}", data.project_dir.data_dir().to_str().unwrap());
+			address_or_vault_name,
+		} => {
+			let address = retrieve_address_string(&data, address_or_vault_name).await?;
+			let secret_and_path = find_vault_key(&data, address.to_string().as_str())?;
+			if secret_and_path.is_none() {
+				println!("No vault entry found for address {address}");
+				println!("You might want to import it first with 'vault import'");
+				return Ok(());
+			}
+			//FIXME not sure if this is ok (but since it's a CLI; this data instance won't be used afterwards)
+			let mut data = data;
+			data.cfg.address = Some(AccountId::from_str(&address).expect("invalid address"));
+			println!("Using key {address}");
+			conf::save_config(&data);
 		}
 		Subcommand::Generate => {
 			// TODO allow custom word count
 			let mnemonic = bip39::Mnemonic::generate(12).unwrap();
 			println!("{mnemonic}");
 		}
-		Subcommand::Import => {
+		Subcommand::Import { secret_format } => {
-			let mnemonic = rpassword::prompt_password("Mnemonic: ")?;
+			let vault_item = prompt_secret_and_compute_vault_item(secret_format)?;
 			println!("Enter password to protect the key");
 			let password = rpassword::prompt_password("Password: ")?;
-			let address = store_mnemonic(&data, &mnemonic, password)?;
+			let address = store_vault_item(&data, vault_item, password)?;
 			println!("Stored secret for {address}");
 		}
+		Subcommand::Rename { address, name } => {
+			let secret_format_and_path = find_vault_key(&data, address.to_string().as_str())?;
+			if secret_format_and_path.is_none() {
+				println!("No vault entry found for address {address}");
+				println!("You might want to import it first with 'vault import'");
+				return Ok(());
+			}
+			let the_vault_user = vault_user::Entity::find_by_id(address.to_string())
+				.one(data.connection.as_ref().unwrap())
+				.await?;
+			match the_vault_user {
+				Some(vault_user) => {
+					let old_name = vault_user.name.clone();
+					let mut vault_user: vault_user::ActiveModel = vault_user.into();
+					vault_user.name = Set(name.clone());
+					let _vault_user = vault_user.update(data.connection.as_ref().unwrap()).await?;
+					println!("Renamed {address} from {old_name} to {name}");
+				}
+				None => {
+					let vault_user = vault_user::ActiveModel {
+						address: Set(address.to_string()),
+						name: Set(name.clone()),
+					};
+					let _vault_user = vault_user.insert(data.connection.as_ref().unwrap()).await?;
+					println!("Renamed {address} to {name}");
+				}
+			}
+			let count_users = vault_user::Entity::find()
+				.count(data.connection.as_ref().unwrap())
+				.await?;
+			println!("Found {count_users} vault users");
+			let vault_key_addresses = fetch_vault_key_addresses(&data).await?;
+			let table = compute_vault_key_table(&data, &vault_key_addresses).await?;
+			println!("available keys:");
+			println!("{table}");
+		}
+		Subcommand::Remove {
+			address_or_vault_name,
+		} => {
+			let address = retrieve_address_string(&data, address_or_vault_name).await?;
+			let secret_and_path = find_vault_key(&data, address.to_string().as_str())?;
+			if secret_and_path.is_none() {
+				println!("No vault entry found for address {address}");
+				println!("You might want to import it first with 'vault import'");
+				return Ok(());
+			}
+			let delete_result = vault_user::Entity::delete_by_id(address.to_string())
+				.exec(data.connection.as_ref().unwrap())
+				.await?;
+			dbg!(delete_result.rows_affected);
+			let (_, vault_path) = secret_and_path.unwrap();
+			std::fs::remove_file(vault_path)?;
+			println!("Removed key {address} and potentially associated vault name");
+		}
+		Subcommand::Where => {
+			println!("{}", data.project_dir.data_dir().to_str().unwrap());
+		}
 	};
 	Ok(())
 }
-/// store mnemonic protected with password
+/// This method will scan files in the data directory and return the addresses of the vault keys found
-pub fn store_mnemonic(
+async fn fetch_vault_key_addresses(data: &Data) -> Result<Vec<String>, GcliError> {
+	let mut entries = std::fs::read_dir(data.project_dir.data_dir())?
+		.map(|res| res.map(|e| e.path()))
+		.collect::<Result<Vec<_>, std::io::Error>>()?;
+	// To have consistent ordering
+	entries.sort();
+	let mut vault_key_addresses: Vec<String> = vec![];
+	entries.iter().for_each(|dir_path| {
+		let filename = dir_path.file_name().unwrap().to_str().unwrap();
+		// To only keep the address part of the filename for names like "<ss58 address>-<secret_format>"
+		let potential_address = filename.split("-").next().unwrap();
+		// If potential_address is a valid AccountId
+		if AccountId::from_str(potential_address).is_ok() {
+			vault_key_addresses.push(potential_address.to_string());
+		}
+	});
+	Ok(vault_key_addresses)
+}
+async fn compute_vault_key_table(
+	data: &Data,
+	vault_key_addresses: &[String],
+) -> Result<Table, GcliError> {
+	let mut table = Table::new();
+	table.load_preset(comfy_table::presets::UTF8_BORDERS_ONLY);
+	table.set_header(vec!["Key", "Vault Name"]);
+	let empty_name = "".to_string();
+	let address_to_name_map = get_address_to_vault_name_map(data).await?;
+	vault_key_addresses.iter().for_each(|address| {
+		let vault_name = address_to_name_map.get(address).unwrap_or(&empty_name);
+		table.add_row(vec![Cell::new(address), Cell::new(vault_name)]);
+	});
+	Ok(table)
+}
+async fn get_address_to_vault_name_map(data: &Data) -> Result<BTreeMap<String, String>, GcliError> {
+	let all_vault_users = vault_user::Entity::find()
+		.all(data.connection.as_ref().unwrap())
+		.await?;
+	let mut address_to_name_map: BTreeMap<String, String> = BTreeMap::new();
+	for vault_user in all_vault_users {
+		address_to_name_map.insert(vault_user.address.clone(), vault_user.name.clone());
+	}
+	Ok(address_to_name_map)
+}
+pub async fn retrieve_address_string<T: AddressOrVaultName>(
+	data: &Data,
+	address_or_vault_name: T,
+) -> Result<String, GcliError> {
+	let address = if let Some(name) = address_or_vault_name.name() {
+		let the_vault_user = vault_user::Entity::find()
+			.filter(vault_user::Column::Name.eq(name))
+			.one(data.connection.as_ref().unwrap())
+			.await?;
+		if let Some(vault_user) = the_vault_user {
+			vault_user.address
+		} else {
+			return Err(GcliError::Input(format!(
+				"No vault user found with name {name}"
+			)));
+		}
+	} else if let Some(address) = address_or_vault_name.address() {
+		address.to_string()
+	} else {
+		//Should never happen since clap enforces exactly one of the 2 options
+		return Err(GcliError::Input("No address or name provided".to_string()));
+	};
+	Ok(address)
+}
+fn create_vault_item<F, P>(
+	secret_format: SecretFormat,
+	prompt_fn: F,
+) -> Result<VaultItem, GcliError>
+where
+	F: Fn() -> (keys::Secret, P),
+	P: Into<KeyPair>,
+{
+	let (secret, pair) = prompt_fn();
+	Ok(VaultItem {
+		secret_format,
+		secret,
+		address: pair.into().address(),
+	})
+}
+fn prompt_secret_and_compute_vault_item(
+	secret_format: SecretFormat,
+) -> Result<VaultItem, GcliError> {
+	match secret_format {
+		SecretFormat::Substrate => {
+			create_vault_item(secret_format, prompt_secret_substrate_and_compute_keypair)
+		}
+		SecretFormat::Seed => create_vault_item(secret_format, prompt_seed_and_compute_keypair),
+		SecretFormat::Cesium => {
+			create_vault_item(secret_format, prompt_secret_cesium_and_compute_keypair)
+		}
+		SecretFormat::Predefined => {
+			create_vault_item(secret_format, prompt_predefined_and_compute_keypair)
+		}
+	}
+}
+/// store VaultItem protected with password
+fn store_vault_item(
 	data: &Data,
-	mnemonic: &str,
+	vault_item: VaultItem,
 	password: String,
 ) -> Result<AccountId, GcliError> {
-	// check validity by deriving keypair
+	// write encrypted secret in file identified by address pubkey and secret_format
-	let keypair = pair_from_str(mnemonic)?;
+	let path = get_vault_key_path(
-	let address = keypair.public();
+		data,
-	// write encrypted mnemonic in file identified by pubkey
+		&compute_vault_filename(&vault_item.secret_format, &vault_item.address.to_string()),
-	let path = data.project_dir.data_dir().join(address.to_string());
+	);
 	let mut file = std::fs::File::create(path)?;
-	file.write_all(&encrypt(mnemonic.as_bytes(), password).map_err(|e| anyhow!(e))?[..])?;
-	Ok(keypair.public().into())
+	match vault_item.secret {
+		keys::Secret::SimpleSecret(secret) => {
+			file.write_all(&encrypt(secret.as_bytes(), password).map_err(|e| anyhow!(e))?[..])?;
+		}
+		keys::Secret::DualSecret(id, pwd) => {
+			//Making a simple separation of the 2 parts with a newline character
+			let secret = format!("{id}\n{pwd}");
+			file.write_all(&encrypt(secret.as_bytes(), password).map_err(|e| anyhow!(e))?[..])?;
+		}
+	}
+	Ok(vault_item.address)
 }
-/// try get secret in keystore
+fn compute_vault_filename(secret_format: &SecretFormat, address: &str) -> String {
-pub fn try_fetch_secret(data: &Data, address: AccountId) -> Result<Option<String>, GcliError> {
+	let format_str: &'static str = From::from(*secret_format);
-	let path = data.project_dir.data_dir().join(address.to_string());
+	format!("{}-{}", address, format_str)
+}
+fn get_vault_key_path(data: &Data, vault_filename: &str) -> PathBuf {
+	data.project_dir.data_dir().join(vault_filename)
+}
+/// look for different possible paths for vault keys and return both format and path
+fn find_vault_key(
+	data: &Data,
+	address: &str,
+) -> Result<Option<(SecretFormat, PathBuf)>, GcliError> {
+	let formats = [
+		SecretFormat::Substrate,
+		SecretFormat::Seed,
+		SecretFormat::Cesium,
+		SecretFormat::Predefined,
+	];
+	for &secret_format in &formats {
+		let path = get_vault_key_path(data, &compute_vault_filename(&secret_format, address));
+		if path.exists() {
+			return Ok(Some((secret_format, path)));
+		}
+	}
+	//Also checking for old file name without secret_format which would be for (default) substrate key
+	let path = get_vault_key_path(data, address);
 	if path.exists() {
+		return Ok(Some((SecretFormat::Substrate, path)));
+	}
+	Ok(None)
+}
+/// try to get secret in keystore, prompt for the password and compute the keypair
+pub fn try_fetch_key_pair(data: &Data, address: AccountId) -> Result<Option<KeyPair>, GcliError> {
+	if let Some((secret_format, path)) = find_vault_key(data, address.to_string().as_str())? {
 		println!("Enter password to unlock account {address}");
 		let password = rpassword::prompt_password("Password: ")?;
 		let mut file = std::fs::OpenOptions::new().read(true).open(path)?;
 		let mut cypher = vec![];
 		file.read_to_end(&mut cypher)?;
-		let secret = decrypt(&cypher, password).map_err(|e| GcliError::Input(e.to_string()))?;
+		let secret_vec = decrypt(&cypher, password).map_err(|e| GcliError::Input(e.to_string()))?;
-		let secretstr = String::from_utf8(secret).map_err(|e| anyhow!(e))?;
+		let secret = String::from_utf8(secret_vec).map_err(|e| anyhow!(e))?;
-		Ok(Some(secretstr))
+		//Still need to handle different secret formats
+		match secret_format {
+			SecretFormat::Substrate => Ok(Some(pair_from_str(&secret)?.into())),
+			SecretFormat::Seed => Ok(Some(pair_from_seed(&secret)?.into())),
+			SecretFormat::Cesium => {
+				let mut lines = secret.lines();
+				//Un-wrapping the 2 secrets from each line
+				let id = lines.next().unwrap();
+				let pwd = lines.next().unwrap();
+				Ok(Some(
+					pair_from_cesium(id.to_string(), pwd.to_string()).into(),
+				))
+			}
+			SecretFormat::Predefined => Ok(Some(pair_from_predefined(&secret)?.into())),
+		}
 	} else {
 		Ok(None)
 	}