vault.rs

mod display;

use crate::commands::cesium::compute_g1v1_public_key;
use crate::entities::vault_account;
use crate::entities::vault_account::{AccountTreeNode, ActiveModel, DbAccountId};
use crate::inputs::{trim_and_reduce_empty_as_none, validate_derivation_path, validate_vault_name};
use crate::keys::seed_from_cesium;
use crate::*;
use age::secrecy::Secret;
use sea_orm::ActiveValue::Set;
use sea_orm::ModelTrait;
use sea_orm::{ConnectionTrait, TransactionTrait};
use sp_core::crypto::AddressUri;
use std::cell::RefCell;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::rc::Rc;

/// vault subcommands
#[derive(Clone, Debug, clap::Parser)]
pub enum Subcommand {
	/// List available SS58 Addresses in the vault
	#[clap(subcommand)]
	List(ListChoice),
	/// Use specific SS58 Address (changes the config Address)
	Use {
		#[clap(flatten)]
		address_or_vault_name: AddressOrVaultNameGroup,
	},
	/// Generate a mnemonic
	Generate,
	/// Import key from (substrate uri) or other format with interactive prompt
	#[clap(
		long_about = "Import key from (substrate uri) or other format with interactive prompt.\n\
		\n\
		This will create a <Base> account in the vault for the provided/computed Substrate URI \n\
		and associated SS58 Address.\n\
		\n\
		If using default format (or specifically \"substrate\") a derivation path is supported\n\
		in the substrate uri value."
	)]
	Import {
		/// Secret key format (substrate, seed, g1v1)
		#[clap(short = 'S', long, required = false, default_value = SecretFormat::Substrate)]
		secret_format: SecretFormat,

		/// Crypto scheme to use (sr25519, ed25519)
		#[clap(short = 'c', long, required = false, default_value = CryptoScheme::Ed25519)]
		crypto_scheme: CryptoScheme,

		/// Substrate URI to import (non-interactive mode)
		#[clap(short = 'u', long, required = false)]
		uri: Option<String>,

		/// G1v1 ID (non-interactive mode for g1v1 format)
		#[clap(long, required = false)]
		g1v1_id: Option<String>,

		/// G1v1 password (non-interactive mode for g1v1 format)
		#[clap(long, required = false)]
		g1v1_password: Option<String>,

		/// Password for encrypting the key (non-interactive mode)
		#[clap(short = 'p', long, required = false, conflicts_with_all=["no_password"])]
		password: Option<String>,

		/// Use empty password for encrypting the key (non-interactive mode)
		#[clap(long, required = false)]
		no_password: bool,

		/// Name for the wallet entry (non-interactive mode) - "" empty string will be considered as None
		#[clap(short = 'n', long, required = false)]
		name: Option<String>,
	},
	/// Add a derivation to an existing SS58 Address
	#[clap(long_about = "Add a derivation to an existing SS58 Address.\n\
		\n\
		Both \"sr25519\" and \"ed25519\" crypto schemes are supported
		\n\
		Use command `vault list base` to see available <Base> account and their crypto scheme\n\
		And then use command 'vault list for' to find all accounts linked to that <Base> account")]
	#[clap(alias = "deriv")]
	#[clap(alias = "derivation")]
	Derive {
		#[clap(flatten)]
		address_or_vault_name: AddressOrVaultNameGroup,

		/// Derivation path (non-interactive mode)
		#[clap(short = 'd', long, required = false)]
		derivation_path: Option<String>,

		/// Password to decrypt the <Base> account key (non-interactive mode)
		#[clap(short = 'p', long, required = false, requires = "derivation_path", conflicts_with_all=["no_password"])]
		password: Option<String>,

		/// Use empty password to decrypt the <Base> account key (non-interactive mode)
		#[clap(long, required = false, requires = "derivation_path")]
		no_password: bool,

		/// Name for the wallet entry (non-interactive mode) - "" empty string will be considered as None
		#[clap(short = 'n', long, required = false, requires = "derivation_path")]
		name: Option<String>,
	},
	/// Give a meaningful name to an SS58 Address in the vault
	Rename {
		/// SS58 Address
		address: AccountId,
	},
	/// Remove an SS58 Address from the vault together with its linked derivations
	#[clap(long_about = "Remove an SS58 Address from the vault\n\
		\n\
		If a <Base> Address is given it will also remove the saved key")]
	Remove {
		#[clap(flatten)]
		address_or_vault_name: AddressOrVaultNameGroup,
	},
	/// Inspect a vault entry, retrieving its Substrate URI (will provide more data in a future version)
	Inspect {
		#[clap(flatten)]
		address_or_vault_name: AddressOrVaultNameGroup,
	},
	/// (deprecated) List available key files (needs to be migrated with command `vault migrate` in order to use them)
	#[deprecated(
		note = "Should be removed in a future version when db persistence of vault is present for a while"
	)]
	ListFiles,
	/// (deprecated) Migrate old key files into db (will have to provide password for each key)
	#[deprecated(
		note = "Should be removed in a future version when db persistence of vault is present for a while"
	)]
	Migrate,
	/// Show where vault db (or old keys) is stored
	Where,
}

/// List subcommands
#[derive(Clone, Debug, clap::Subcommand)]
pub enum ListChoice {
	/// List all accounts
	#[clap(alias = "a")]
	All {
		/// Show G1v1 public key for ed25519 keys
		#[clap(long)]
		show_g1v1: bool,
		/// Show wallet type (g1v1 or mnemonic)
		#[clap(long)]
		show_type: bool,
	},
	/// List only base accounts
	#[clap(alias = "b")]
	Base {
		/// Show G1v1 public key for ed25519 keys
		#[clap(long)]
		show_g1v1: bool,
		/// Show wallet type (g1v1 or mnemonic)
		#[clap(long)]
		show_type: bool,
	},
	/// List accounts for a specific address
	#[clap(alias = "f")]
	For {
		#[clap(flatten)]
		address_or_vault_name: AddressOrVaultNameGroup,
		/// Show G1v1 public key for ed25519 keys
		#[clap(long)]
		show_g1v1: bool,
		/// Show wallet type (g1v1 or mnemonic)
		#[clap(long)]
		show_type: bool,
	},
}

impl Default for ListChoice {
	fn default() -> Self {
		ListChoice::All {
			show_g1v1: false,
			show_type: false,
		}
	}
}

#[derive(Debug, clap::Args, Clone)]
#[group(required = true, multiple = false)]
pub struct AddressOrVaultNameGroup {
	/// SS58 Address
	#[clap(short)]
	address: Option<AccountId>,
	/// Name of an SS58 Address in the vault
	#[clap(short = 'v')]
	vault_name: Option<String>,
}

pub struct VaultDataToImport {
	secret_format: SecretFormat,
	secret_suri: String,
	key_pair: KeyPair,
}

// encrypt input with passphrase
pub fn encrypt(input: &[u8], passphrase: String) -> Result<Vec<u8>, age::EncryptError> {
	let encryptor = age::Encryptor::with_user_passphrase(Secret::new(passphrase));
	let mut encrypted = vec![];
	let mut writer = encryptor.wrap_output(age::armor::ArmoredWriter::wrap_output(
		&mut encrypted,
		age::armor::Format::AsciiArmor,
	)?)?;
	writer.write_all(input)?;
	writer.finish().and_then(|armor| armor.finish())?;
	Ok(encrypted)
}

// decrypt cypher with passphrase
pub fn decrypt(input: &[u8], passphrase: String) -> Result<Vec<u8>, age::DecryptError> {
	let age::Decryptor::Passphrase(decryptor) =
		age::Decryptor::new(age::armor::ArmoredReader::new(input))?
	else {
		unimplemented!()
	};
	let mut decrypted = vec![];
	let mut reader = decryptor.decrypt(&Secret::new(passphrase.to_owned()), None)?;
	reader.read_to_end(&mut decrypted)?;
	Ok(decrypted)
}

/// handle vault commands
pub async fn handle_command(data: Data, command: Subcommand) -> Result<(), GcliError> {
	let db = data.connect_db();

	// match subcommand
	match command {
		Subcommand::List(choice) => match choice {
			ListChoice::All {
				show_g1v1,
				show_type,
			} => {
				let all_account_tree_node_hierarchies =
					vault_account::fetch_all_base_account_tree_node_hierarchies(db).await?;

				let table = display::compute_vault_accounts_table_with_g1v1(
					&all_account_tree_node_hierarchies,
					show_g1v1,
					show_type,
				)?;

				println!("available SS58 Addresses:");
				println!("{table}");
			}
			ListChoice::Base {
				show_g1v1,
				show_type,
			} => {
				let base_account_tree_nodes =
					vault_account::fetch_only_base_account_tree_nodes(db).await?;

				let table = display::compute_vault_accounts_table_with_g1v1(
					&base_account_tree_nodes,
					show_g1v1,
					show_type,
				)?;

				println!("available <Base> SS58 Addresses:");
				println!("{table}");
			}
			ListChoice::For {
				address_or_vault_name,
				show_g1v1,
				show_type,
			} => {
				let account_tree_node =
					retrieve_account_tree_node(db, address_or_vault_name).await?;

				let base_account_tree_node =
					vault_account::get_base_account_tree_node(&account_tree_node);

				let table = display::compute_vault_accounts_table_with_g1v1(
					&[base_account_tree_node],
					show_g1v1,
					show_type,
				)?;

				println!(
					"available SS58 Addresses linked to {}:",
					account_tree_node.borrow().account
				);
				println!("{table}");
			}
		},
		Subcommand::ListFiles => {
			let vault_key_addresses = fetch_vault_key_addresses(&data).await?;

			let table = display::compute_vault_key_files_table(&vault_key_addresses)?;

			println!("available key files (needs to be migrated with command `vault migrate` in order to use them):");
			println!("{table}");
		}
		Subcommand::Use {
			address_or_vault_name,
		} => {
			let account = retrieve_vault_account(db, address_or_vault_name).await?;

			println!("Using: {}", account);

			let updated_cfg = conf::Config {
				address: Some(account.address.0),
				..data.cfg
			};

			//This updated configuration will be picked up with next GCli execution
			conf::save(&updated_cfg);
		}
		Subcommand::Generate => {
			// TODO allow custom word count
			let mnemonic = bip39::Mnemonic::generate(12).unwrap();
			println!("{mnemonic}");
		}
		Subcommand::Import {
			secret_format,
			crypto_scheme,
			uri,
			g1v1_id,
			g1v1_password,
			password,
			no_password,
			name,
		} => {
			let vault_data_for_import = if let Some(uri_str) = uri {
				// Non-interactive mode with provided URI
				if secret_format != SecretFormat::Substrate {
					return Err(GcliError::Input(format!(
						"URI can only be provided directly with secret_format=substrate, got: {:?}",
						secret_format
					)));
				}

				// Create keypair from provided URI
				let key_pair = compute_keypair(crypto_scheme, &uri_str)?;

				VaultDataToImport {
					secret_format,
					secret_suri: uri_str,
					key_pair,
				}
			} else if let (Some(id), Some(pwd)) = (&g1v1_id, &g1v1_password) {
				// Non-interactive mode with provided G1v1 ID and password
				if secret_format != SecretFormat::G1v1 {
					return Err(GcliError::Input(format!(
						"G1v1 ID and password can only be provided directly with secret_format=g1v1, got: {:?}", 
						secret_format
					)));
				}

				// Create keypair from provided G1v1 ID and password
				let seed = seed_from_cesium(id, pwd);
				let secret_suri = format!("0x{}", hex::encode(seed));

				// G1v1 always uses Ed25519
				let key_pair = compute_keypair(CryptoScheme::Ed25519, &secret_suri)?;

				VaultDataToImport {
					secret_format,
					secret_suri,
					key_pair,
				}
			} else {
				// Interactive mode
				prompt_secret_and_compute_vault_data_to_import(secret_format, crypto_scheme)?
			};

			//Extra check for SecretFormat::G1v1 (old cesium) - showing the G1v1 cesium public key for confirmation
			if secret_format == SecretFormat::G1v1 {
				println!(
					"The G1v1 public key for the provided secret is: '{}'",
					compute_g1v1_public_key(&vault_data_for_import.key_pair)?
				);

				// Skip confirmation in non-interactive mode
				let is_non_interactive_g1v1 = g1v1_id.is_some() && g1v1_password.is_some();
				if !is_non_interactive_g1v1 {
					let confirmed = inputs::confirm_action("Is it the correct one (if not, you should try again to input G1v1 id/password) ?".to_string())?;
					if !confirmed {
						return Ok(());
					}
				}
			}

			let txn = db.begin().await?;

			// Handle password in non-interactive mode
			let provided_password = if no_password {
				Some(String::new()) // Empty password
			} else {
				password
			};

			let _account = create_base_account_for_vault_data_to_import(
				&txn,
				&vault_data_for_import,
				provided_password.as_ref(),
				Some(crypto_scheme),
				name,
			)
			.await?;

			txn.commit().await?;

			println!("Change done");
		}
		Subcommand::Derive {
			address_or_vault_name,
			derivation_path,
			password,
			no_password,
			name,
		} => {
			let account_tree_node_to_derive =
				retrieve_account_tree_node(db, address_or_vault_name).await?;

			let account_to_derive = account_tree_node_to_derive.borrow().account.clone();

			let base_account_tree_node =
				vault_account::get_base_account_tree_node(&account_tree_node_to_derive);

			let base_account = &base_account_tree_node.borrow().account.clone();

			println!("Adding derivation to: {account_to_derive}");

			let base_parent_hierarchy_account_tree_node_to_derive =
				vault_account::get_base_parent_hierarchy_account_tree_node(
					&account_tree_node_to_derive,
				);

			let parent_hierarchy_table_account_to_derive =
				display::compute_vault_accounts_table(&[
					base_parent_hierarchy_account_tree_node_to_derive,
				])?;

			println!();
			println!("Its parent hierarchy is this:");
			println!("{parent_hierarchy_table_account_to_derive}");
			println!();

			println!("The linked <Base> account is {base_account}");

			// Handle password from non-interactive mode or ask for it
			let password = if no_password {
				String::new()
			} else if let Some(password) = password {
				password
			} else {
				println!("Enter password to decrypt the <Base> account key");
				inputs::prompt_password()?
			};

			let account_to_derive_secret_suri = vault_account::compute_suri_account_tree_node(
				&account_tree_node_to_derive,
				password,
			)?;

			println!();

			// Handle derivation_path from non-interactive mode or ask for it
			let derivation_path = if let Some(derivation_path) = derivation_path {
				validate_derivation_path(derivation_path.clone())?;
				derivation_path
			} else {
				inputs::prompt_vault_derivation_path()?
			};

			let derivation_secret_suri =
				format!("{account_to_derive_secret_suri}{derivation_path}");

			let crypto_scheme = base_account
				.crypto_scheme
				.map(CryptoScheme::from)
				.unwrap_or(CryptoScheme::Ed25519); // Fallback to Ed25519 if not defined (should never happen)

			let derivation_keypair = compute_keypair(crypto_scheme, &derivation_secret_suri)?;

			let derivation_address: String = derivation_keypair.address().to_string();

			let txn = db.begin().await?;

			println!();
			let _derivation = create_derivation_account(
				&txn,
				&derivation_address,
				&derivation_path,
				&account_to_derive.address.to_string(),
				name,
			)
			.await?;

			txn.commit().await?;
			println!("Change done");
		}
		Subcommand::Rename { address } => {
			let account =
				vault_account::find_by_id(db, &DbAccountId::from(address.clone())).await?;

			if account.is_none() {
				println!("No vault entry found for address:'{address}'");
				println!("You might want to import it first with 'vault import'");
				return Ok(());
			}

			let account = account.unwrap();

			println!(
				"Current name for address:'{address}' is {:?}",
				&account.name
			);

			println!("Enter new name for address (leave empty to remove the name)");
			let name =
				inputs::prompt_vault_name_and_check_availability(db, account.name.as_ref()).await?;

			let _account = vault_account::update_account_name(db, account, name.as_ref()).await?;

			println!("Rename done");
		}
		Subcommand::Remove {
			address_or_vault_name,
		} => {
			let account_tree_node_to_delete =
				retrieve_account_tree_node(db, address_or_vault_name).await?;

			let txn = db.begin().await?;

			let account_to_delete = account_tree_node_to_delete.borrow().account.clone();
			let address_to_delete = account_tree_node_to_delete.borrow().account.address.clone();

			//If account to delete has children; also delete all linked derivations
			if !account_tree_node_to_delete.borrow().children.is_empty() {
				let table =
					display::compute_vault_accounts_table(&[account_tree_node_to_delete.clone()])?;

				println!("All addresses linked to: {account_to_delete}");
				println!("{table}");

				println!(
					"This {} account has {} addresses in total",
					account_to_delete.account_type(),
					vault_account::count_accounts_in_account_tree_node_and_children(
						&account_tree_node_to_delete
					)
				);

				let confirmation_message = if account_to_delete.is_base_account() {
					"Are you sure you want to delete it along with the saved key?"
				} else {
					"Are you sure you want to delete it?"
				};

				let confirmed = inputs::confirm_action(confirmation_message.to_string())?;

				if !confirmed {
					return Ok(());
				}

				for account_to_delete in
					vault_account::extract_accounts_depth_first_from_account_tree_node(
						&account_tree_node_to_delete,
					)? {
					let delete_result = account_to_delete.delete(&txn).await?;
					println!("Deleting {} address", delete_result.rows_affected);
				}
			} else {
				let delete_result = account_to_delete.delete(&txn).await?;
				println!("Deleting {} address", delete_result.rows_affected);
			}

			txn.commit().await?;

			println!("Done removing address:'{address_to_delete}'");
		}
		Subcommand::Inspect {
			address_or_vault_name,
		} => {
			let account_tree_node_to_derive =
				retrieve_account_tree_node(db, address_or_vault_name).await?;

			println!("Enter password to decrypt the <Base> account key");
			let password = inputs::prompt_password()?;

			let account_to_derive_secret_suri = vault_account::compute_suri_account_tree_node(
				&account_tree_node_to_derive,
				password,
			)?;

			println!("Substrate URI: '{account_to_derive_secret_suri}'")
		}
		Subcommand::Migrate => {
			println!("Migrating existing key files to db");

			let vault_key_addresses = fetch_vault_key_addresses(&data).await?;

			let table = display::compute_vault_key_files_table(&vault_key_addresses)?;
			println!("available key files to possibly migrate:");
			println!("{table}");

			for address in vault_key_addresses {
				//Check if we already have a vault_derivation for that address
				let existing_account =
					vault_account::find_by_id(db, &DbAccountId::from_str(&address)?).await?;

				if existing_account.is_some() {
					//Already migrated
					continue;
				}

				println!();
				println!("Trying to migrate key {address}");
				let vault_data_from_file = match try_fetch_vault_data_from_file(&data, &address) {
					Ok(Some(vault_data)) => vault_data,
					Ok(None) => {
						println!("No vault entry file found for address {address}");
						continue;
					}
					Err(e) => {
						println!("Error while fetching vault data for address {address}: {e}");
						println!("Continuing to next one");
						continue;
					}
				};

				let vault_data_to_import = VaultDataToImport {
					secret_format: vault_data_from_file.secret_format,
					secret_suri: vault_data_from_file.secret,
					key_pair: vault_data_from_file.key_pair,
				};

				let txn = db.begin().await?;

				// Old key files were in Sr25519 format (and had the Address associated to that scheme)
				let account = create_base_account_for_vault_data_to_import(
					&txn,
					&vault_data_to_import,
					Some(&vault_data_from_file.password),
					Some(CryptoScheme::Sr25519),
					None,
				)
				.await;

				match account {
					Ok(_account) => {
						txn.commit().await?;
						println!("Change done");
					}
					Err(error) => {
						println!("Error occurred: {error}");
						println!("Continuing to next key");
					}
				}
			}

			println!("Migration done");
		}
		Subcommand::Where => {
			println!("{}", data.project_dir.data_dir().to_str().unwrap());
		}
	};

	Ok(())
}

/// Method used to separate vault `name` part from optional `derivation` part in computed names that can be provided by users in the different `vault` commands using `AddressOrVaultNameGroup`
fn parse_vault_name_and_derivation_path_from_user_input(
	user_input_name: String,
) -> Result<(String, Option<String>), GcliError> {
	if user_input_name.contains("/") {
		user_input_name.find("/").map_or(
			Err(GcliError::Input("Invalid format".to_string())),
			|idx| {
				let (prefix, derivation_path) = user_input_name.split_at(idx);
				Ok((prefix.to_string(), Some(derivation_path.to_string())))
			},
		)
	} else {
		Ok((user_input_name, None))
	}
}

/// Method that can be used to parse a Substrate URI (which can also be only a derivation path)
///
/// Does some internal verification (done by sp_core::address_uri::AddressUri)
///
/// It extracts the (optional) `phrase` and the (optional) recomposed full `derivation path`
///
/// It also checks if a derivation `password` was provided and returns an error if one was found
pub fn parse_prefix_and_derivation_path_from_suri(
	raw_string: String,
) -> Result<(Option<String>, Option<String>), GcliError> {
	let address_uri =
		AddressUri::parse(&raw_string).map_err(|e| GcliError::Input(e.to_string()))?;

	if let Some(pass) = address_uri.pass {
		return Err(GcliError::Input(format!(
			"Having a password in the derivation path is not supported (password:'{}')",
			pass
		)));
	}

	let full_path = if address_uri.paths.is_empty() {
		None
	} else {
		Some("/".to_owned() + &address_uri.paths.into_iter().collect::<Vec<_>>().join("/"))
	};

	Ok((address_uri.phrase.map(|s| s.to_string()), full_path))
}

fn map_secret_format_to_crypto_scheme(
	secret_format: SecretFormat,
	override_crypto_scheme: Option<CryptoScheme>,
) -> CryptoScheme {
	// If a crypto_scheme is explicitly specified, use it except for G1v1 which must always use Ed25519
	if let Some(scheme) = override_crypto_scheme {
		if secret_format == SecretFormat::G1v1 {
			// G1v1 must always use Ed25519
			CryptoScheme::Ed25519
		} else {
			scheme
		}
	} else {
		// Default behavior if no crypto_scheme is specified
		match secret_format {
			// All formats use Ed25519 by default
			SecretFormat::Seed => CryptoScheme::Ed25519,
			SecretFormat::Substrate => CryptoScheme::Ed25519,
			SecretFormat::Predefined => CryptoScheme::Ed25519,
			SecretFormat::G1v1 => CryptoScheme::Ed25519,
		}
	}
}

/// This method will scan files in the data directory and return the addresses of the vault keys found
#[deprecated(
	note = "Should be removed in a future version when db persistence of vault is present for a while"
)]
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();
		// If potential_address is a valid AccountId
		if AccountId::from_str(filename).is_ok() {
			vault_key_addresses.push(filename.to_string());
		}
	});

	Ok(vault_key_addresses)
}

pub async fn retrieve_vault_account_for_name<C>(
	db: &C,
	name_input: &String,
) -> Result<vault_account::Model, GcliError>
where
	C: ConnectionTrait,
{
	let account_tree_node = retrieve_account_tree_node_for_name(db, name_input).await?;

	//Need this extra step to avoid borrowing issues
	let account = account_tree_node.borrow().account.clone();
	Ok(account)
}

pub async fn retrieve_account_tree_node<C>(
	db: &C,
	address_or_vault_name: AddressOrVaultNameGroup,
) -> Result<Rc<RefCell<AccountTreeNode>>, GcliError>
where
	C: ConnectionTrait,
{
	let account_tree_node = if let Some(name_input) = &address_or_vault_name.vault_name {
		retrieve_account_tree_node_for_name(db, name_input).await?
	} else if let Some(address) = &address_or_vault_name.address {
		let base_account_tree_node =
			vault_account::fetch_base_account_tree_node_hierarchy_unwrapped(
				db,
				&address.to_string(),
			)
			.await?;

		let account_tree_node_for_address = vault_account::get_account_tree_node_for_address(
			&base_account_tree_node,
			&address.to_string(),
		);

		Rc::clone(&account_tree_node_for_address)
	} 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(account_tree_node)
}

pub async fn retrieve_account_tree_node_for_name<C>(
	db: &C,
	name_input: &String,
) -> Result<Rc<RefCell<AccountTreeNode>>, GcliError>
where
	C: ConnectionTrait,
{
	let (name, derivation_path_opt) =
		parse_vault_name_and_derivation_path_from_user_input(name_input.to_string())?;

	let account_for_name = vault_account::find_by_name(db, &name).await?;

	let account_for_name = account_for_name.ok_or(GcliError::Input(format!(
		"No account found with name:'{name}'"
	)))?;

	let base_account_tree_node = vault_account::fetch_base_account_tree_node_hierarchy_unwrapped(
		db,
		&account_for_name.address.to_string(),
	)
	.await?;

	let account_tree_node_for_name = vault_account::get_account_tree_node_for_address(
		&base_account_tree_node,
		&account_for_name.address.to_string(),
	);

	Ok(match derivation_path_opt {
		None => Rc::clone(&account_tree_node_for_name),
		Some(path) => {
			let account_tree_node_for_name_input =
				vault_account::compute_name_map_for_account_tree_node(&account_tree_node_for_name)?
					.get(name_input)
					.cloned()
					.ok_or(GcliError::Input(format!(
						"No account found with name:'{name}' and path:'{path}'"
					)))?;

			Rc::clone(&account_tree_node_for_name_input)
		}
	})
}

pub async fn retrieve_vault_account<C>(
	db: &C,
	address_or_vault_name: AddressOrVaultNameGroup,
) -> Result<vault_account::Model, GcliError>
where
	C: ConnectionTrait,
{
	let account_tree_node = retrieve_account_tree_node(db, address_or_vault_name).await?;

	//Need this extra step to avoid borrowing issues
	let account = account_tree_node.borrow().account.clone();
	Ok(account)
}

fn create_vault_data_to_import<F, P>(
	secret_format: SecretFormat,
	crypto_scheme: CryptoScheme,
	prompt_fn: F,
) -> Result<VaultDataToImport, GcliError>
where
	F: Fn(CryptoScheme) -> (String, P),
	P: Into<KeyPair>,
{
	let (secret, pair) = prompt_fn(crypto_scheme);
	let key_pair = pair.into();
	Ok(VaultDataToImport {
		secret_format,
		secret_suri: secret,
		key_pair,
	})
}

fn prompt_secret_and_compute_vault_data_to_import(
	secret_format: SecretFormat,
	crypto_scheme: CryptoScheme,
) -> Result<VaultDataToImport, GcliError> {
	match secret_format {
		SecretFormat::Substrate => create_vault_data_to_import(
			secret_format,
			crypto_scheme,
			prompt_secret_substrate_and_compute_keypair,
		),
		SecretFormat::Seed => create_vault_data_to_import(
			secret_format,
			crypto_scheme,
			prompt_seed_and_compute_keypair,
		),
		SecretFormat::G1v1 => {
			// G1v1 always uses Ed25519, ignore crypto_scheme
			create_vault_data_to_import(
				secret_format,
				CryptoScheme::Ed25519,
				prompt_secret_cesium_and_compute_keypair,
			)
		}
		SecretFormat::Predefined => create_vault_data_to_import(
			secret_format,
			crypto_scheme,
			prompt_predefined_and_compute_keypair,
		),
	}
}

/// Creates a `base` vault account for vault_data provided and returns it
///
/// Does extra checks and asks for user input in case the address is already present in the vault.
///
/// Can request password and (optional) name to the user at the proper time
///
/// Typically used for `vault import|migrate` commands
pub async fn create_base_account_for_vault_data_to_import<C>(
	db_tx: &C,
	vault_data: &VaultDataToImport,
	password_opt: Option<&String>,
	crypto_scheme: Option<CryptoScheme>,
	name_opt: Option<String>,
) -> Result<vault_account::Model, GcliError>
where
	C: ConnectionTrait,
{
	let address = vault_data.key_pair.address().to_string();

	// Check if the account already exists
	let existing_vault_account =
		vault_account::find_by_id(db_tx, &DbAccountId(vault_data.key_pair.address())).await?;

	let password = match password_opt {
		Some(password) => password.clone(),
		None => inputs::prompt_password_query("Enter password to encrypt the key: ")?,
	};

	let encrypted_suri = compute_encrypted_suri(password.clone(), vault_data.secret_suri.clone())?;

	if let Some(existing_vault_account) = existing_vault_account {
		// Existing account
		match inputs::confirm_action(format!(
			"Account {} already exists. Do you want to update it?",
			existing_vault_account
		))? {
			true => {
				let name = if let Some(name) = name_opt {
					validate_vault_name(&name)?;
					trim_and_reduce_empty_as_none(name)
				} else {
					println!("(Optional) Enter a name for the vault entry (leave empty to remove the name)");
					inputs::prompt_vault_name_and_check_availability(
						db_tx,
						existing_vault_account.name.as_ref(),
					)
					.await?
				};

				// Since links are made based on address / parent(address) we can just edit the existing entry and it should be fine
				let mut vault_account: ActiveModel = existing_vault_account.into();
				vault_account.path = Set(None);
				vault_account.parent = Set(None);
				vault_account.crypto_scheme = Set(Some(
					map_secret_format_to_crypto_scheme(vault_data.secret_format, crypto_scheme)
						.into(),
				));
				vault_account.encrypted_suri = Set(Some(encrypted_suri));
				vault_account.name = Set(name);
				vault_account.secret_format = Set(Some(vault_data.secret_format.into()));
				let updated_vault_account =
					vault_account::update_account(db_tx, vault_account).await?;

				println!("Updating vault account {updated_vault_account}");
				Ok(updated_vault_account)
			}
			_ => Err(GcliError::Input("import canceled".into())),
		}
	} else {
		//New entry
		let secret_format = vault_data.secret_format;

		let encrypted_suri = compute_encrypted_suri(password, vault_data.secret_suri.clone())?;

		let name = if let Some(name) = name_opt {
			validate_vault_name(&name)?;
			trim_and_reduce_empty_as_none(name)
		} else {
			println!("(Optional) Enter a name for the vault entry");
			inputs::prompt_vault_name_and_check_availability(db_tx, None).await?
		};

		let crypto_scheme = map_secret_format_to_crypto_scheme(secret_format, crypto_scheme);

		let account = vault_account::create_base_account(
			db_tx,
			&address,
			name.as_ref(),
			crypto_scheme,
			encrypted_suri,
			secret_format,
		)
		.await?;

		println!("Creating <Base> account {account}");
		Ok(account)
	}
}

/// Creates a `derivation` vault account for data provided and returns it
///
/// Does extra checks and asks for user input in case the address is already present in the vault.
///
/// Can request (optional) name to the user at the proper time
///
/// Typically used for `vault derive` command
pub async fn create_derivation_account<C>(
	db_tx: &C,
	derivation_address: &String,
	derivation_path: &String,
	parent_address: &String,
	name_opt: Option<String>,
) -> Result<vault_account::Model, GcliError>
where
	C: ConnectionTrait,
{
	println!("Trying to create derivation with address '{derivation_address}'");
	println!();
	let vault_account = if let Some(existing_vault_account) =
		vault_account::find_by_id(db_tx, &DbAccountId::from(derivation_address.clone())).await?
	{
		// Existing account
		println!("You are trying to derive '{derivation_path}' from parent '{parent_address}'");

		if existing_vault_account.is_base_account() {
			println!(
				"but it is already present as a direct <Base> account '{}'",
				existing_vault_account.address
			);

			println!("Do you want to:");
			println!("1. keep the existing <Base> account and cancel import");
			println!("2. delete the existing <Base> account and associated key and replace it with the new derivation account (children will be re-parented)");
		} else {
			//Existing derivation
			let existing_account_tree_node_hierarchy =
				vault_account::fetch_base_account_tree_node_hierarchy_unwrapped(
					db_tx,
					derivation_address,
				)
				.await?;

			let existing_account_tree_node_for_address =
				vault_account::get_account_tree_node_for_address(
					&existing_account_tree_node_hierarchy,
					derivation_address,
				);

			let base_parent_hierarchy_existing_account_tree_node =
				vault_account::get_base_parent_hierarchy_account_tree_node(
					&existing_account_tree_node_for_address,
				);

			let parent_hierarchy_table_existing_account =
				display::compute_vault_accounts_table(&[
					base_parent_hierarchy_existing_account_tree_node,
				])?;

			println!(
				"but it is already present as `{}` derivation of '{}' account.",
				existing_vault_account.path.clone().unwrap(),
				existing_vault_account.parent.clone().unwrap()
			);

			println!();
			println!("Its parent hierarchy is this:");
			println!("{parent_hierarchy_table_existing_account}");
			println!();
			println!("Do you want to:");
			println!("1. keep the existing derivation and cancel import");
			println!("2. delete the derivation account and replace it with this new derivation (children will be re-parented)");
		}

		let result = inputs::select_action("Your choice?", vec!["1", "2"])?;
		match result {
			"2" => {
				let name = if let Some(name) = name_opt {
					validate_vault_name(&name)?;
					trim_and_reduce_empty_as_none(name)
				} else {
					println!("(Optional) Enter a name for the vault entry (leave empty to remove the name)");
					inputs::prompt_vault_name_and_check_availability(
						db_tx,
						existing_vault_account.name.as_ref(),
					)
					.await?
				};

				// Since links are made based on address / parent(address) we can just edit the existing entry and it should be fine
				let mut vault_account: ActiveModel = existing_vault_account.into();
				vault_account.path = Set(Some(derivation_path.clone()));
				vault_account.parent = Set(Some(DbAccountId::from(parent_address.clone())));
				vault_account.crypto_scheme = Set(None);
				vault_account.encrypted_suri = Set(None);
				vault_account.name = Set(name.clone());
				let updated_vault_account =
					vault_account::update_account(db_tx, vault_account).await?;

				println!("Updating vault account {updated_vault_account}");
				updated_vault_account
			}
			_ => {
				return Err(GcliError::Input("derive canceled".into()));
			}
		}
	} else {
		let name = if let Some(name) = name_opt {
			validate_vault_name(&name)?;
			trim_and_reduce_empty_as_none(name)
		} else {
			println!("(Optional) Enter a name for the vault entry");
			inputs::prompt_vault_name_and_check_availability(db_tx, None).await?
		};

		let derivation = vault_account::create_derivation_account(
			db_tx,
			derivation_address,
			name.as_ref(),
			derivation_path,
			parent_address,
		)
		.await?;
		println!("Creating derivation account {derivation}");

		derivation
	};

	Ok(vault_account)
}

/// Function will compute the encrypted suri
fn compute_encrypted_suri(password: String, secret_suri: String) -> Result<Vec<u8>, GcliError> {
	encrypt(secret_suri.as_bytes(), password).map_err(|e| GcliError::Input(e.to_string()))
}

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_substrate_vault_key_file(data: &Data, address: &str) -> Result<Option<PathBuf>, GcliError> {
	let path = get_vault_key_path(data, address);
	if path.exists() {
		return Ok(Some(path));
	}

	Ok(None)
}

/// try to get secret in keystore, prompt for the password and compute the keypair
pub async fn try_fetch_key_pair(
	data: &Data,
	address: AccountId,
) -> Result<Option<KeyPair>, GcliError> {
	if let Some(account_tree_node_hierarchy) =
		vault_account::fetch_base_account_tree_node_hierarchy(
			data.connect_db(),
			&address.to_string(),
		)
		.await?
	{
		let account_tree_node = vault_account::get_account_tree_node_for_address(
			&account_tree_node_hierarchy,
			&address.to_string(),
		);

		println!("(Vault: {})", account_tree_node.borrow().account);

		let password = inputs::prompt_password()?;
		let secret_suri =
			vault_account::compute_suri_account_tree_node(&account_tree_node, password)?;

		let base_account_tree_node = vault_account::get_base_account_tree_node(&account_tree_node);

		let base_account = &base_account_tree_node.borrow().account.clone();

		let key_pair = compute_keypair(base_account.crypto_scheme.unwrap().into(), &secret_suri)?;

		//To be safe
		if address != key_pair.address() {
			return Err(GcliError::Input(format!(
				"Computed address {} does not match the expected address {}",
				key_pair.address(),
				address
			)));
		}

		Ok(Some(key_pair))
	} else {
		Ok(None)
	}
}

pub fn compute_keypair(
	crypto_scheme: CryptoScheme,
	secret_suri: &str,
) -> Result<KeyPair, GcliError> {
	let key_pair = match crypto_scheme {
		CryptoScheme::Sr25519 => pair_from_sr25519_str(secret_suri)?.into(),
		CryptoScheme::Ed25519 => pair_from_ed25519_str(secret_suri)?.into(),
	};
	Ok(key_pair)
}

pub struct VaultDataFromFile {
	secret_format: SecretFormat,
	secret: String,
	#[allow(dead_code)]
	path: PathBuf,
	password: String,
	key_pair: KeyPair,
}

/// try to get secret in keystore, prompt for the password and compute the keypair
#[deprecated(
	note = "Should be removed in a future version when db persistence of vault is present for a while"
)]
pub fn try_fetch_vault_data_from_file(
	data: &Data,
	address: &str,
) -> Result<Option<VaultDataFromFile>, GcliError> {
	if let Some(path) = find_substrate_vault_key_file(data, address)? {
		println!("Enter password to unlock account {address}");
		let password = inputs::prompt_password()?;
		let mut file = std::fs::OpenOptions::new().read(true).open(path.clone())?;
		let mut cypher = vec![];
		file.read_to_end(&mut cypher)?;
		let secret_vec =
			decrypt(&cypher, password.clone()).map_err(|e| GcliError::Input(e.to_string()))?;
		let secret = String::from_utf8(secret_vec).map_err(|e| anyhow!(e))?;

		let key_pair = pair_from_sr25519_str(&secret)?.into();

		Ok(Some(VaultDataFromFile {
			secret_format: SecretFormat::Substrate,
			secret,
			path,
			password,
			key_pair,
		}))
	} else {
		Ok(None)
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use rstest::rstest;

	/// test that armored encryption/decryption work as intended
	#[test]
	fn test_encrypt_decrypt() {
		let plaintext = b"Hello world!";
		let passphrase = "this is not a good passphrase".to_string();
		let encrypted = encrypt(plaintext, passphrase.clone()).unwrap();
		let decrypted = decrypt(&encrypted, passphrase).unwrap();
		assert_eq!(decrypted, plaintext);
	}

	#[rstest]
	#[case(
		String::from("bottom drive obey lake curtain smoke basket hold race lonely fit walk//0"),
		Some(String::from(
			"bottom drive obey lake curtain smoke basket hold race lonely fit walk"
		)),
		Some(String::from("//0"))
	)]
	#[case(
		String::from("0xfac7959dbfe72f052e5a0c3c8d6530f202b02fd8f9f5ca3580ec8deb7797479e//0"),
		Some(String::from("0xfac7959dbfe72f052e5a0c3c8d6530f202b02fd8f9f5ca3580ec8deb7797479e")),
		Some(String::from("//0"))
	)]
	#[case(
		String::from(
			"bottom drive obey lake curtain smoke basket hold race lonely fit walk//Alice"
		),
		Some(String::from(
			"bottom drive obey lake curtain smoke basket hold race lonely fit walk"
		)),
		Some(String::from("//Alice"))
	)]
	#[case(
		String::from(
			"bottom drive obey lake curtain smoke basket hold race lonely fit walk//Alice//Bob/soft1/soft2"
		),
		Some(String::from("bottom drive obey lake curtain smoke basket hold race lonely fit walk")),
		Some(String::from("//Alice//Bob/soft1/soft2"))
	)]
	#[case(
		String::from("bottom drive obey lake curtain smoke basket hold race lonely fit walk"),
		Some(String::from(
			"bottom drive obey lake curtain smoke basket hold race lonely fit walk"
		)),
		None
	)]
	#[case(
		String::from("0xfac7959dbfe72f052e5a0c3c8d6530f202b02fd8f9f5ca3580ec8deb7797479e"),
		Some(String::from("0xfac7959dbfe72f052e5a0c3c8d6530f202b02fd8f9f5ca3580ec8deb7797479e")),
		None
	)]
	#[case(
		String::from("fac7959dbfe72f052e5a0c3c8d6530f202b02fd8f9f5ca3580ec8deb7797479e"),
		Some(String::from("fac7959dbfe72f052e5a0c3c8d6530f202b02fd8f9f5ca3580ec8deb7797479e")),
		None
	)]
	#[case(
		String::from("someVaultName//Alice"),
		Some(String::from("someVaultName")),
		Some(String::from("//Alice"))
	)]
	#[case(
		String::from("someVaultName"),
		Some(String::from("someVaultName")),
		None
	)]
	fn test_parse_prefix_and_derivation_path_from_suri(
		#[case] raw_string: String,
		#[case] expected_prefix: Option<String>,
		#[case] expected_derivation_path: Option<String>,
	) {
		let (root_secret, derivation_path) =
			parse_prefix_and_derivation_path_from_suri(raw_string).unwrap();
		assert_eq!(expected_prefix, root_secret);
		assert_eq!(expected_derivation_path, derivation_path);
	}

	#[rstest]
	#[case(
		String::from("//Alice//Bob/soft1/soft2"),
		None,
		Some(String::from("//Alice//Bob/soft1/soft2"))
	)]
	#[case(String::from(""), None, None)]
	#[case(String::from("//0"), None, Some(String::from("//0")))]
	fn test_parse_prefix_and_derivation_path_from_suri_works_with_empty_prefix_phrase(
		#[case] raw_string: String,
		#[case] expected_prefix: Option<String>,
		#[case] expected_derivation_path: Option<String>,
	) {
		let (root_secret, derivation_path) =
			parse_prefix_and_derivation_path_from_suri(raw_string).unwrap();
		assert_eq!(expected_prefix, root_secret);
		assert_eq!(expected_derivation_path, derivation_path);
	}

	#[rstest]
	#[case(
        String::from(
            "bottom drive obey lake curtain smoke basket hold race lonely fit walk//Alice//Bob/soft1/soft2///password"
        ),
    )]
	#[case(String::from(
		"bottom drive obey lake curtain smoke basket hold race lonely fit walk///password"
	))]
	#[case(String::from(
		"bottom drive obey lake curtain smoke basket hold race lonely fit walk///"
	))]
	#[case(
        String::from(
            "bottom drive obey lake curtain smoke basket hold race lonely fit walk///password//NotDerivations//Still/password/part"
        ),
    )]
	fn test_parse_prefix_and_derivation_path_from_suri_does_not_allow_password(
		#[case] raw_string: String,
	) {
		let result = parse_prefix_and_derivation_path_from_suri(raw_string);
		match result.unwrap_err() {
			GcliError::Input(err) => {
				println!("Error message: {}", err);
				assert!(
					err.starts_with("Having a password in the derivation path is not supported")
				);
			}
			other => panic!("Should have been an Input error; got: {:?}", other),
		}
	}
}