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::*;
use age::secrecy::Secret;
use sea_orm::ActiveValue::Set;
use sea_orm::{ConnectionTrait, TransactionTrait};
use sea_orm::{DbErr, ModelTrait};
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,
	},
	/// Add a derivation to an existing SS58 Address
	#[clap(long_about = "Add a derivation to an existing SS58 Address.\n\
		\n\
		Only \"sr25519\" crypto scheme is supported for derivations.\n\
		\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,
	},
	/// 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,
}

#[derive(Clone, Default, Debug, clap::Parser)]
pub enum ListChoice {
	/// List all <Base> SS58 Addresses and their linked derivations in the vault
	#[default]
	All,
	/// List <Base> and Derivation SS58 Addresses linked to the selected one
	For {
		#[clap(flatten)]
		address_or_vault_name: AddressOrVaultNameGroup,
	},
	/// List all <Base> SS58 Addresses in the vault
	Base,
}

#[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')]
	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 => {
				let all_account_tree_node_hierarchies =
					vault_account::fetch_all_base_account_tree_node_hierarchies(db).await?;
				let table =
					display::compute_vault_accounts_table(&all_account_tree_node_hierarchies)?;

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

				let table = display::compute_vault_accounts_table(&base_account_tree_nodes)?;

				println!("available <Base> SS58 Addresses:");
				println!("{table}");
			}
			ListChoice::For {
				address_or_vault_name,
			} => {
				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(&[base_account_tree_node])?;

				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 } => {
			let vault_data_for_import =
				prompt_secret_and_compute_vault_data_to_import(secret_format)?;

			//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)?
				);
				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?;

			println!();
			let _account =
				create_base_account_for_vault_data_to_import(&txn, &vault_data_for_import, None)
					.await?;

			txn.commit().await?;

			println!("Change done");
		}
		Subcommand::Derive {
			address_or_vault_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();

			if base_account.crypto_scheme.is_none() {
				return Err(GcliError::DatabaseError(DbErr::Custom(format!("Crypto scheme is not set for the base account:{base_account} - should never happen"))));
			}

			if let Some(crypto_scheme) = base_account.crypto_scheme {
				if CryptoScheme::from(crypto_scheme) == CryptoScheme::Ed25519 {
					println!(
						"Only \"{}\" crypto scheme is supported for derivations.",
						Into::<&str>::into(CryptoScheme::Sr25519),
					);
					println!();
					println!(
                        "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"
                    );
					return Ok(());
				}
			}

			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}");

			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!();
			let derivation_path = inputs::prompt_vault_derivation_path()?;

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

			let derivation_keypair =
				compute_keypair(CryptoScheme::Sr25519, &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(),
			)
			.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?;

				let account = create_base_account_for_vault_data_to_import(
					&txn,
					&vault_data_to_import,
					Some(&vault_data_from_file.password),
				)
				.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) -> CryptoScheme {
	match secret_format {
		SecretFormat::Seed => CryptoScheme::Sr25519,
		SecretFormat::Substrate => CryptoScheme::Sr25519,
		SecretFormat::Predefined => CryptoScheme::Sr25519,
		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.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,
	prompt_fn: F,
) -> Result<VaultDataToImport, GcliError>
where
	F: Fn() -> (String, P),
	P: Into<KeyPair>,
{
	let (secret, pair) = prompt_fn();
	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,
) -> Result<VaultDataToImport, GcliError> {
	match secret_format {
		SecretFormat::Substrate => {
			create_vault_data_to_import(secret_format, prompt_secret_substrate_and_compute_keypair)
		}
		SecretFormat::Seed => {
			create_vault_data_to_import(secret_format, prompt_seed_and_compute_keypair)
		}
		SecretFormat::G1v1 => {
			create_vault_data_to_import(secret_format, prompt_secret_cesium_and_compute_keypair)
		}
		SecretFormat::Predefined => {
			create_vault_data_to_import(secret_format, 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: Option<&String>,
) -> Result<vault_account::Model, GcliError>
where
	C: ConnectionTrait,
{
	let address_to_import = vault_data.key_pair.address().to_string();
	println!("Trying to import for SS58 address :'{}'", address_to_import);
	println!();

	let vault_account = if let Some(existing_vault_account) =
		vault_account::find_by_id(db_tx, &DbAccountId::from(address_to_import.clone())).await?
	{
		if existing_vault_account.is_base_account() {
			println!("You are trying to add {address_to_import} as a <Base> account while it already exists as a <Base> account.");
			println!();
			println!("Do you want to:");
			println!("1. keep the existing <Base> account and cancel import");
			println!("2. overwrite existing <Base> account with the new encrypted key (children will be re-parented)");
		} else {
			// Existing derivation account
			let account_tree_node_hierarchy =
				vault_account::fetch_base_account_tree_node_hierarchy_unwrapped(
					db_tx,
					&address_to_import,
				)
				.await?;
			let account_tree_node_for_address = vault_account::get_account_tree_node_for_address(
				&account_tree_node_hierarchy,
				&address_to_import,
			);

			let base_parent_hierarchy_account_tree_node =
				vault_account::get_base_parent_hierarchy_account_tree_node(
					&account_tree_node_for_address,
				);

			let parent_hierarchy_table =
				display::compute_vault_accounts_table(&[base_parent_hierarchy_account_tree_node])?;

			println!("You are trying to add {address_to_import} as a <Base> account");
			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}");
			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 the new <Base> account (children will be re-parented)");
		}

		let result = inputs::select_action("Your choice?", vec!["1", "2"])?;
		match result {
			"2" => {
				let encrypted_suri =
					compute_encrypted_suri(password, vault_data.secret_suri.clone())?;

				println!(
					"(Optional) Enter a name for the vault entry (leave empty to remove the name)"
				);
				let 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).into(),
				));
				vault_account.encrypted_suri = Set(Some(encrypted_suri));
				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("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())?;

		println!("(Optional) Enter a name for the vault entry");
		let name = inputs::prompt_vault_name_and_check_availability(db_tx, None).await?;

		let crypto_scheme = map_secret_format_to_crypto_scheme(secret_format);

		let base_account = vault_account::create_base_account(
			db_tx,
			&address_to_import,
			name.as_ref(),
			crypto_scheme,
			encrypted_suri,
		)
		.await?;
		println!("Creating <Base> account {base_account}");

		base_account
	};

	Ok(vault_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,
) -> 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" => {
				println!(
					"(Optional) Enter a name for the vault entry (leave empty to remove the name)"
				);
				let 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 {
		println!("(Optional) Enter a name for the vault entry");
		let name = 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 ask for password if not present and compute the encrypted suri
fn compute_encrypted_suri(
	password: Option<&String>,
	secret_suri: String,
) -> Result<Vec<u8>, GcliError> {
	let password = match password.cloned() {
		Some(password) => password,
		_ => {
			println!("Enter password to protect the key");
			inputs::prompt_password_confirm()?
		}
	};

	Ok(encrypt(secret_suri.as_bytes(), password).map_err(|e| anyhow!(e))?)
}

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?
	{