vault.rs

use crate::entities::vault_account::CryptoType;
use crate::entities::{vault_account, vault_derivation};
use crate::*;
use age::secrecy::Secret;
use comfy_table::{Cell, Table};
use sea_orm::ActiveValue::Set;
use sea_orm::{ActiveModelTrait, EntityTrait, ModelTrait};
use sea_orm::{ColumnTrait, QueryFilter};
use sea_orm::{ConnectionTrait, TransactionTrait};
use std::io::{Read, Write};
use std::path::PathBuf;

/// define universal dividends 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)mnemonic or other format with interactive prompt
	#[clap(
		long_about = "Import key from (substrate)mnemonic or other format with interactive prompt\n\
		\n\
		If a (substrate)mnemonic is provided with a derivation path, it will ensure the base <Account>\n\
		and associated SS58 Address exists before creating the derivation; but please use command \n\
		`vault derivation|derive|deriv` to add a derivation to an existing <Account> instead."
	)]
	Import {
		/// Secret key format (substrate, seed, cesium)
		#[clap(short = 'S', long, required = false, default_value = SecretFormat::Substrate)]
		secret_format: SecretFormat,
	},
	/// Add a derivation to an existing <Account>
	#[clap(long_about = "Add a derivation to an existing <Account>\n\
		\n\
		Only \"substrate\" and \"seed\" format are supported for derivations\n\
		Use command `vault list account` to see available <Account> and their format")]
	#[clap(alias = "deriv")]
	#[clap(alias = "derive")]
	Derivation {
		#[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
	#[clap(long_about = "Remove an SS58 Address from the vault\n\
		\n\
		If an <Account> Address is given it will also remove all linked derivations")]
	Remove {
		#[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)
	ListFiles,
	/// (deprecated)Migrate old key files into db (will have to provide password for each key)
	Migrate,
	/// Show where vault db (or old keys) is stored
	Where,
}

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

pub struct VaultDataToImport {
	secret_format: SecretFormat,
	secret: keys::Secret,
	address: AccountId,
	key_pair: KeyPair,
}

// encrypt input with passphrase
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
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 ud commands
pub async fn handle_command(data: Data, command: Subcommand) -> Result<(), GcliError> {
	// match subcommand
	match command {
		Subcommand::List(choice) => match choice {
			ListChoice::All => {
				let derivations = vault_derivation::list_all_derivations_in_order(
					data.connection.as_ref().unwrap(),
				)
				.await?;

				let table = compute_vault_derivations_table(
					data.connection.as_ref().unwrap(),
					&derivations,
				)
				.await?;

				println!("available SS58 Addresses:");
				println!("{table}");
			}
			ListChoice::Account => {
				let derivations = vault_derivation::list_all_root_derivations_in_order(
					data.connection.as_ref().unwrap(),
				)
				.await?;

				let table = compute_vault_derivations_table(
					data.connection.as_ref().unwrap(),
					&derivations,
				)
				.await?;

				println!("available <Account> SS58 Addresses:");
				println!("{table}");
			}
			ListChoice::For {
				address_or_vault_name,
			} => {
				let selected_derivation =
					retrieve_vault_derivation(&data, address_or_vault_name).await?;

				let linked_derivations = vault_derivation::fetch_all_linked_derivations_in_order(
					data.connection.as_ref().unwrap(),
					&selected_derivation.root_address,
				)
				.await?;

				let table = compute_vault_derivations_table(
					data.connection.as_ref().unwrap(),
					&linked_derivations,
				)
				.await?;

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

			let table = compute_vault_key_files_table(&vault_key_addresses).await?;

			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 derivation = retrieve_vault_derivation(&data, address_or_vault_name).await?;

			//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(&derivation.address).expect("invalid address"));

			println!("Using key {}", derivation);
			conf::save_config(&data);
		}
		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)?;

			println!(
				"Trying to import for address :'{}'",
				vault_data_for_import.address
			);

			if let Some(derivation) =
				vault_derivation::Entity::find_by_id(vault_data_for_import.address.to_string())
					.one(data.connection.as_ref().unwrap())
					.await?
			{
				println!(
					"Vault entry already exists for address:'{}'",
					vault_data_for_import.address
				);

				let linked_derivations = vault_derivation::fetch_all_linked_derivations_in_order(
					data.connection.as_ref().unwrap(),
					&derivation.root_address.clone(),
				)
				.await?;
				println!("Here are all the linked derivations already present in the vault:");

				let table = compute_vault_derivations_table(
					data.connection.as_ref().unwrap(),
					&linked_derivations,
				)
				.await?;
				println!("{table}");

				return Ok(());
			}

			println!("Enter password to protect the key");
			let password = inputs::prompt_password_confirm()?;

			println!("(Optional) Enter a name for the vault entry");
			let name = inputs::prompt_vault_name()?;

			let txn = data.connection.as_ref().unwrap().begin().await?;

			let _derivation = create_derivation_for_vault_data_to_import(
				&txn,
				&vault_data_for_import,
				&password,
				name.as_ref(),
			)
			.await?;

			txn.commit().await?;

			println!("Import done");
		}
		Subcommand::Derivation {
			address_or_vault_name,
		} => {
			let root_derivation = retrieve_vault_derivation(&data, address_or_vault_name).await?;

			if root_derivation.path.is_some() {
				println!("Can only add derivation on a ROOT key.");
				println!(
					"The selected key with address:'{}' already has a ROOT with address:'{}'",
					root_derivation.address, root_derivation.root_address
				);
				println!("You can check for available ROOT keys with command 'vault list root'");
				return Ok(());
			}

			let vault_account = vault_account::Entity::find_by_id(&root_derivation.address)
				.one(data.connection.as_ref().unwrap())
				.await?
				.ok_or(GcliError::Input(format!(
					"Could not find (root) vault account for address:'{}'",
					root_derivation.address
				)))?;

			if vault_account.crypto_type == CryptoType::Ed25519Seed {
				println!(
					"Only \"{}\" and \"{}\" format are supported for derivations",
					Into::<&str>::into(SecretFormat::Substrate),
					Into::<&str>::into(SecretFormat::Seed)
				);
				println!("Use command `vault list root` to see available <Account> and their format");
				return Ok(());
			}

			println!("Adding derivation to root key: {root_derivation}");

			println!("Enter password to decrypt the root key");

			let root_secret_suri = retrieve_suri_from_vault_account(&vault_account)?;

			let derivation_path = inputs::prompt_vault_derivation_path()?;

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

			let derivation_keypair =
				compute_keypair(vault_account.crypto_type, &derivation_secret_suri)?;

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

			let check_derivation = vault_derivation::Entity::find_by_id(&derivation_address)
				.one(data.connection.as_ref().unwrap())
				.await?;

			if check_derivation.is_some() {
				println!("Derivation already exists for address:'{derivation_address}'");
				return Ok(());
			}

			println!("(Optional) Enter a name for the new derivation");
			let name = inputs::prompt_vault_name()?;

			let derivation = vault_derivation::ActiveModel {
				address: Set(derivation_address),
				name: Set(name),
				path: Set(Some(derivation_path)),
				root_address: Set(root_derivation.root_address.clone()),
			};
			let derivation = derivation.insert(data.connection.as_ref().unwrap()).await?;
			println!("Created derivation {}", derivation);
		}
		Subcommand::Rename { address } => {
			let derivation = vault_derivation::Entity::find_by_id(address.to_string())
				.one(data.connection.as_ref().unwrap())
				.await?;

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

			let derivation = derivation.unwrap();

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

			println!("Enter new vault name for the key (leave empty to remove the name)");
			let name = inputs::prompt_vault_name()?;

			let old_name = derivation.name.clone();
			let mut derivation: vault_derivation::ActiveModel = derivation.into();
			derivation.name = Set(name.clone());
			let _derivation = derivation.update(data.connection.as_ref().unwrap()).await?;
			println!(
				"Renamed address:'{address}' from {:?} to {:?}",
				old_name, name
			);
		}
		Subcommand::Remove {
			address_or_vault_name,
		} => {
			let derivation = retrieve_vault_derivation(&data, address_or_vault_name).await?;
			let address_to_delete = derivation.address.clone();

			let txn = data.connection.as_ref().unwrap().begin().await?;

			//If deleting a root derivation; also delete the vault account and all linked derivations
			if derivation.path.is_none() {
				let all_derivations_to_delete =
					vault_derivation::fetch_all_linked_derivations_in_order(
						&txn,
						&address_to_delete,
					)
					.await?;

				let table =
					compute_vault_derivations_table(&txn, &all_derivations_to_delete).await?;

				println!("All derivations linked to the root derivation:");
				println!("{table}");

				println!(
					"This root derivation has {} derivations in total",
					all_derivations_to_delete.len()
				);
				let confirmed = inputs::confirm_action(
                    "Are you sure you want to delete it along with the vault account (and saved key)?".to_string()
                )?;

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

				for derivation_to_delete in all_derivations_to_delete {
					let delete_result = derivation_to_delete.delete(&txn).await?;
					println!("Deleted {} derivation", delete_result.rows_affected);
				}

				let delete_result = vault_account::Entity::delete_by_id(&address_to_delete)
					.exec(&txn)
					.await?;
				println!("Deleted {} vault account", delete_result.rows_affected);
			} else {
				let delete_result = derivation.delete(&txn).await?;
				println!("Deleted {} derivation", delete_result.rows_affected);
			}

			txn.commit().await?;

			println!("Done removing address:'{address_to_delete}'");
		}
		Subcommand::Migrate => {
			println!("Migrating existing key files to db");

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

			let table = compute_vault_key_files_table(&vault_key_addresses).await?;
			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 derivation = vault_derivation::Entity::find_by_id(&address)
					.one(data.connection.as_ref().unwrap())
					.await?;

				if derivation.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: SecretFormat::Substrate,
					secret: keys::Secret::SimpleSecret(vault_data_from_file.secret),
					address: AccountId::from_str(&address).expect("invalid address"),
					key_pair: vault_data_from_file.key_pair,
				};

				let txn = data.connection.as_ref().unwrap().begin().await?;

				let derivation = create_derivation_for_vault_data_to_import(
					&txn,
					&vault_data_to_import,
					&vault_data_from_file.password,
					None,
				)
				.await?;

				txn.commit().await?;
				println!("Import done: {}", derivation);
			}

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

	Ok(())
}

fn parse_prefix_and_derivation_path_from_string(
	raw_string: String,
) -> Result<(String, Option<String>), GcliError> {
	if raw_string.contains("/") {
		raw_string
			.find("/")
			.map_or(Err(GcliError::Input("Invalid format".to_string())), |idx| {
				let (prefix, derivation_path) = raw_string.split_at(idx);
				Ok((prefix.to_string(), Some(derivation_path.to_string())))
			})
	} else {
		Ok((raw_string, None))
	}
}

fn map_secret_format_to_crypto_type(secret_format: SecretFormat) -> CryptoType {
	match secret_format {
		SecretFormat::Seed => vault_account::CryptoType::Sr25519Seed,
		SecretFormat::Substrate => vault_account::CryptoType::Sr25519Mnemonic,
		SecretFormat::Predefined => vault_account::CryptoType::Sr25519Mnemonic,
		SecretFormat::Cesium => vault_account::CryptoType::Ed25519Seed,
	}
}

/// This method will scan files in the data directory and return the addresses of the vault keys found
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_files_table(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 file"]);

	vault_key_addresses.iter().for_each(|address| {
		table.add_row(vec![Cell::new(address)]);
	});

	Ok(table)
}

async fn compute_vault_derivations_table<C>(
	db: &C,
	derivations_ordered: &[vault_derivation::Model],
) -> Result<Table, GcliError>
where
	C: ConnectionTrait,
{
	let mut table = Table::new();
	table.load_preset(comfy_table::presets::UTF8_BORDERS_ONLY);
	table.set_header(vec!["SS58 Address", "Format", "Account/Path", "Name"]);
	let empty_string = "".to_string();
	let root_path = "<Account>".to_string();

	let mut current_root_address = "".to_string();
	let mut current_root_name: Option<String> = None;
	let mut current_vault_format: Option<&str> = None;

	for derivation in derivations_ordered {
		if derivation.root_address != current_root_address {
			// First entry should be a root derivation
			if derivation.path.is_some() {
				return Err(GcliError::Input(
					"Order of derivations parameter is wrong".to_string(),
				));
			}
			current_root_address = derivation.root_address.clone();
			current_root_name = derivation.name.clone();

			let vault_account = vault_account::Entity::find_by_id(current_root_address.clone())
				.one(db)
				.await?
				.ok_or(GcliError::Input(format!(
					"No vault <Account> found with address:'{current_root_address}'"
				)))?;

			current_vault_format = match vault_account.crypto_type {
				CryptoType::Sr25519Mnemonic => Some(SecretFormat::Substrate.into()),
				CryptoType::Sr25519Seed => Some(SecretFormat::Seed.into()),
				CryptoType::Ed25519Seed => Some(SecretFormat::Cesium.into()),
			};
		} else {
			// Validate that the RootAddress is the same as current_root_address
			if derivation.root_address != current_root_address {
				return Err(GcliError::Input(
					"Order of derivations parameter is wrong".to_string(),
				));
			}
		}

		let address = if derivation.path.is_none() {
			derivation.address.clone()
		} else {
			"  ".to_string() + &derivation.address
		};

		let (path, format) = if derivation.path.is_none() {
			(root_path.clone(), current_vault_format.unwrap())
		} else {
			(derivation.path.clone().unwrap(), empty_string.as_str())
		};

		let name = if derivation.name.is_none() {
			if derivation.path.is_none() {
				empty_string.clone()
			} else if let Some(current_root_name) = &current_root_name {
				format!(
					"<{}{}>",
					current_root_name,
					derivation.path.clone().unwrap()
				)
			} else {
				empty_string.clone()
			}
		} else {
			derivation.name.clone().unwrap()
		};

		table.add_row(vec![
			Cell::new(&address),
			Cell::new(format),
			Cell::new(&path),
			Cell::new(&name),
		]);
	}

	Ok(table)
}

pub async fn retrieve_address_string<T: AddressOrVaultName>(
	data: &Data,
	address_or_vault_name: T,
) -> Result<String, GcliError> {
	if let Some(address) = address_or_vault_name.address() {
		return Ok(address.to_string());
	}

	let derivation = retrieve_vault_derivation(data, address_or_vault_name).await?;

	Ok(derivation.address)
}

pub async fn retrieve_vault_derivation<T: AddressOrVaultName>(
	data: &Data,
	address_or_vault_name: T,
) -> Result<vault_derivation::Model, GcliError> {
	let derivation = if let Some(name) = address_or_vault_name.name() {
		let (name, derivation_path_opt) =
			parse_prefix_and_derivation_path_from_string(name.to_string())?;

		let derivation = vault_derivation::Entity::find()
			.filter(vault_derivation::Column::Name.eq(Some(name.clone())))
			.one(data.connection.as_ref().unwrap())
			.await?;

		let derivation = derivation.ok_or(GcliError::Input(format!(
			"No vault SS58 Address found with name:'{name}'"
		)))?;

		match derivation_path_opt {
			None => derivation,
			Some(path) => {
				let sub_derivation = vault_derivation::Entity::find()
					.filter(
						vault_derivation::Column::RootAddress.eq(derivation.root_address.clone()),
					)
					.filter(vault_derivation::Column::Path.eq(Some(path.clone())))
					.one(data.connection.as_ref().unwrap())
					.await?;

				sub_derivation.ok_or(GcliError::Input(format!(
					"No vault derivation found with <Account> name:'{name}' and path:'{path}'"
				)))?
			}
		}
	} else if let Some(address) = address_or_vault_name.address() {
		let derivation = vault_derivation::Entity::find_by_id(address.to_string())
			.one(data.connection.as_ref().unwrap())
			.await?;

		derivation.ok_or(GcliError::Input(format!(
			"No vault entry found with Address:'{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(derivation)
}

fn create_vault_data_to_import<F, P>(
	secret_format: SecretFormat,
	prompt_fn: F,
) -> Result<VaultDataToImport, GcliError>
where
	F: Fn() -> (keys::Secret, P),
	P: Into<KeyPair>,
{
	let (secret, pair) = prompt_fn();
	let key_pair = pair.into();
	Ok(VaultDataToImport {
		secret_format,
		secret,
		address: key_pair.address(),
		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::Cesium => {
			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 derivation and if necessary root vault account and root derivation
///
/// Does it all using "db" parameter that should better be a transaction since multiple operations can be done
pub async fn create_derivation_for_vault_data_to_import<C>(
	db: &C,
	vault_data: &VaultDataToImport,
	password: &str,
	name: Option<&String>,
) -> Result<vault_derivation::Model, GcliError>
where
	C: ConnectionTrait,
{
	//To be safe
	if vault_derivation::Entity::find_by_id(vault_data.address.to_string())
		.one(db)
		.await?
		.is_some()
	{
		return Err(GcliError::Input(format!(
			"Vault entry already exists for address {}",
			vault_data.address
		)));
	}

	let secret_suri: String = match &vault_data.secret_format {
		SecretFormat::Cesium => {
			if let KeyPair::Nacl(keypair) = &vault_data.key_pair {
				// In case of cesium key, we will store the seed suri instead of id/password (so it supports derivations)
				let seed: [u8; 32] = keypair.skey[0..32]
					.try_into()
					.expect("slice with incorrect length");
				format!("0x{}", hex::encode(seed))
			} else {
				return Err(GcliError::Input("Expected KeyPair::Nacl".to_string()));
			}
		}
		SecretFormat::Seed => {
			if let keys::Secret::SimpleSecret(seed_str) = &vault_data.secret {
				format!("0x{seed_str}")
			} else {
				return Err(GcliError::Input("Expected SimpleSecret".to_string()));
			}
		}
		SecretFormat::Substrate | SecretFormat::Predefined => {
			if let keys::Secret::SimpleSecret(secret_suri) = &vault_data.secret {
				secret_suri.clone()
			} else {
				return Err(GcliError::Input("Expected SimpleSecret".to_string()));
			}
		}
	};

	let secret_format = vault_data.secret_format;

	let (root_secret_suri, derivation_path_opt, root_address, derivation_address) =
		compute_root_and_derivation_data(&secret_format, secret_suri)?;

	// Making sure the computed address is the same as the address to import
	let address_to_import = vault_data.address.to_string();
	if let Some(derivation_address) = &derivation_address {
		if *derivation_address != address_to_import {
			return Err(GcliError::Input(format!(
				"Derivation address {} does not match the expected address {}",
				derivation_address, address_to_import
			)));
		}
	} else if root_address != address_to_import {
		return Err(GcliError::Input(format!(
			"Derivation address {} does not match the expected address {}",
			root_address, address_to_import
		)));
	}

	let encrypted_private_key =
		encrypt(root_secret_suri.as_bytes(), password.to_string()).map_err(|e| anyhow!(e))?;

	let _root_account = vault_account::create_vault_account(
		db,
		&root_address,
		map_secret_format_to_crypto_type(secret_format),
		encrypted_private_key,
	)
	.await?;

	let derivation = if let Some(derivation_path) = derivation_path_opt {
		// Root derivation
		let _root_derivation =
			vault_derivation::create_root_vault_derivation(db, &root_address, None).await?;

		// Compute derivation !
		let derivation = vault_derivation::ActiveModel {
			address: Set(derivation_address.unwrap().clone()),
			name: Set(name.cloned()),
			path: Set(Some(derivation_path)),
			root_address: Set(root_address.clone()),
		};
		let derivation = derivation.insert(db).await?;
		println!("Created derivation {}", derivation);
		derivation
	} else {
		let derivation =
			vault_derivation::create_root_vault_derivation(db, &root_address, name).await?;
		println!("Created derivation {}", derivation);
		derivation
	};

	Ok(derivation)
}

fn compute_root_and_derivation_data(
	secret_format: &SecretFormat,
	secret_suri: String,
) -> Result<(String, Option<String>, String, Option<String>), GcliError> {
	let (root_secret_suri, derivation_path_opt) =
		parse_prefix_and_derivation_path_from_string(secret_suri)?;

	let (root_address, derivation_address_opt) = match &secret_format {
		SecretFormat::Cesium => match &derivation_path_opt {
			None => {
				let root_suri = &root_secret_suri;
				let root_pair = pair_from_ed25519_str(root_suri)?;
				let root_address: AccountId = root_pair.public().into();

				(root_address.to_string(), None)
			}
			Some(derivation_path) => {
				let root_suri = &root_secret_suri;
				let root_pair = pair_from_ed25519_str(root_suri)?;
				let root_address: AccountId = root_pair.public().into();

				let derivation_suri = root_suri.clone() + derivation_path;
				let derivation_pair = pair_from_ed25519_str(&derivation_suri)?;
				let derivation_address: AccountId = derivation_pair.public().into();

				(
					root_address.to_string(),
					Some(derivation_address.to_string()),
				)
			}
		},
		SecretFormat::Substrate | SecretFormat::Seed | SecretFormat::Predefined => {
			match &derivation_path_opt {
				None => {
					let root_suri = &root_secret_suri;
					let root_pair = pair_from_sr25519_str(root_suri)?;
					let root_address: AccountId = root_pair.public().into();

					(root_address.to_string(), None)
				}
				Some(derivation_path) => {
					let root_suri = &root_secret_suri;
					let root_pair = pair_from_sr25519_str(root_suri)?;
					let root_address: AccountId = root_pair.public().into();

					let derivation_suri = root_suri.clone() + derivation_path;
					let derivation_pair = pair_from_sr25519_str(&derivation_suri)?;
					let derivation_address: AccountId = derivation_pair.public().into();

					(
						root_address.to_string(),
						Some(derivation_address.to_string()),
					)
				}
			}
		}
	};
	Ok((
		root_secret_suri,
		derivation_path_opt,
		root_address,
		derivation_address_opt,
	))
}

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(derivation) = vault_derivation::Entity::find_by_id(address.to_string())
		.one(data.connection.as_ref().unwrap())
		.await?
	{
		if let Some(vault_account) =
			vault_account::Entity::find_by_id(derivation.root_address.clone())
				.one(data.connection.as_ref().unwrap())
				.await?
		{
			let root_secret_suri = retrieve_suri_from_vault_account(&vault_account)?;

			let secret_suri = if let Some(derivation_path) = derivation.path {
				format!("{root_secret_suri}{derivation_path}")
			} else {
				root_secret_suri
			};

			let key_pair = compute_keypair(vault_account.crypto_type, &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)
		}
	} else {
		Ok(None)
	}
}

pub fn retrieve_suri_from_vault_account(
	vault_account: &vault_account::Model,
) -> Result<String, GcliError> {
	let password = inputs::prompt_password()?;

	let cypher = &vault_account.encrypted_private_key;
	let secret_vec =
		decrypt(cypher, password.clone()).map_err(|e| GcliError::Input(e.to_string()))?;
	let secret_suri = String::from_utf8(secret_vec).map_err(|e| anyhow!(e))?;

	Ok(secret_suri)
}

pub fn compute_keypair(crypto_type: CryptoType, secret_suri: &str) -> Result<KeyPair, GcliError> {
	let key_pair = match crypto_type {
		CryptoType::Sr25519Mnemonic | CryptoType::Sr25519Seed => {
			pair_from_sr25519_str(secret_suri)?.into()
		}
		CryptoType::Ed25519Seed => pair_from_ed25519_str(secret_suri)?.into(),
	};
	Ok(key_pair)
}

pub struct VaultDataFromFile {
	address: String,
	secret_format: SecretFormat,
	secret: String,
	path: PathBuf,
	password: String,
	key_pair: KeyPair,
}

/// try to get secret in keystore, prompt for the password and compute the keypair
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 = rpassword::prompt_password("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 {
			address: address.to_string(),