More lenient key querying, smaller release binaries, don't panic on stdout close, formatting

The key querying first tries exact match on a fingerprint, and then falls back to contains() on the alias

.vscode/settings.json

@@ -6,6 +6,7 @@
         "josekit",
         "PKCS",
         "Pkey",
-        "printdoc"
+        "printdoc",
+        "writedoc"
     ]
 }

Cargo.toml

@@ -29,3 +29,7 @@ async-trait = "0.1.68"
 
 [profile.release]
 strip = true
+opt-level = "z"
+lto = true
+codegen-units = 1
+panic = "abort"

crates/asp/src/keys/mod.rs

@@ -105,7 +105,8 @@ impl AspKey {
 					})
 				}
 			}
-		})().or(Err(AspKeyError::GenerationError))
+		})()
+		.or(Err(AspKeyError::GenerationError))
 	}
 
 	pub fn create_signer(&self) -> anyhow::Result<Box<dyn JwsSigner>> {

src/commands/keys/export.rs

@@ -1,12 +1,18 @@
+use anstyle::{AnsiColor, Color as AnstyleColor, Reset, Style as Anstyle};
 use anyhow::{anyhow, Context};
 use argon2::{password_hash::SaltString, Argon2, PasswordHasher};
 use asp::keys::AspKey;
 use clap::{Parser, ValueEnum};
 use data_encoding::BASE64_NOPAD;
 use dialoguer::{theme::ColorfulTheme, Password};
-use sea_orm::EntityTrait;
+use indoc::writedoc;
 
-use crate::{commands::AspmSubcommand, entities::keys};
+use std::io::Write;
+
+use crate::{
+	commands::{AspmSubcommand, KeysEntityExt, KeysQueryResult},
+	entities::keys::{Entity as KeysEntity, Model as KeysModel},
+};
 
 #[derive(ValueEnum, Debug, Clone)]
 pub enum KeyExportFormat {
@@ -25,58 +31,86 @@ pub struct KeysExportCommand {
 	/// The format to export the key into. All of these formats can then be correctly imported back into this tool, but only some are encrypted, so keep that in mind when handling the exported keys.
 	#[clap(value_enum, ignore_case = true)]
 	format: KeyExportFormat,
-	/// The fingerprint of the key to export. This can be obtained with the `keys list` command.
-	fingerprint: String,
+	/// The key to export. This can either be a fingerprint obtained with the `keys list` command, or an alias to search for.
+	key: String,
 }
 
 #[async_trait::async_trait]
 impl AspmSubcommand for KeysExportCommand {
 	async fn execute(&self, state: crate::AspmState) -> Result<(), anyhow::Error> {
 		// Fetch key from db
-		let entry = keys::Entity::find_by_id(&self.fingerprint)
-			.one(&state.db)
+		let entry = KeysEntity::query_key(&state.db, &self.key)
 			.await
-			.context("Unable to fetch key from database")?;
-
-		if let Some(entry) = entry {
-			let key_password = std::env::var("KEY_PASSWORD").or_else(|_| {
-				Password::with_theme(&ColorfulTheme::default())
-					.with_prompt("Please enter a password to decrypt the key with")
-					.interact()
-					.context("Unable to prompt on stderr")
-			})?;
-
-			let argon_salt = SaltString::from_b64(
-				&BASE64_NOPAD.encode(self.fingerprint.to_uppercase().as_bytes()),
-			)
-			.context("Unable to decode argon2 salt")?;
-			let argon2 = Argon2::default();
-			let hash = argon2
-				.hash_password(key_password.as_bytes(), &argon_salt)
-				.or(Err(anyhow!("Unable to derive encryption key")))?;
-			let aes_key = hash.hash.context("Unable to derive encryption key")?;
-			let aes_key = &aes_key.as_bytes()[0..32];
-
-			if let Ok(decrypted) = AspKey::from_encrypted(aes_key, &entry.encrypted) {
-				let export = match self.format {
-					KeyExportFormat::Encrypted => decrypted
-						.export_encrypted(aes_key)
-						.context("Unable to convert key into encrypted format")?,
-					KeyExportFormat::PKCS8 => decrypted
-						.into_pkcs8()
-						.context("Unable to convert key into PKCS#8 format")?,
-					KeyExportFormat::Jwk => decrypted.jwk.to_string(),
-				};
-
-				println!("{export}")
-			} else {
-				eprintln!("There was an error decrypting the key, please make sure the password you entered was correct");
+			.context("Unable to query keys from database")?;
+		let key: KeysModel = match entry {
+			KeysQueryResult::None => {
+				eprintln!(
+					"{style}No keys matching the given query were found{reset}",
+					style = Anstyle::new()
+						.fg_color(Some(AnstyleColor::Ansi(AnsiColor::BrightRed)))
+						.render(),
+					reset = Reset.render()
+				);
+				std::process::exit(1);
 			}
-		} else {
+			KeysQueryResult::One(key) => key,
+			KeysQueryResult::Many(mut keys) => {
+				eprintln!(
+					"{style}More than one keys matching the given query were found{reset}",
+					style = Anstyle::new()
+						.fg_color(Some(AnstyleColor::Ansi(AnsiColor::BrightYellow)))
+						.render(),
+					reset = Reset::default().render()
+				);
+				keys.remove(0)
+			}
+		};
+
+		let key_password = std::env::var("KEY_PASSWORD").or_else(|_| {
+			Password::with_theme(&ColorfulTheme::default())
+				.with_prompt("Please enter a password to decrypt the key with")
+				.interact()
+				.context("Unable to prompt on stderr")
+		})?;
+
+		let argon_salt =
+			SaltString::from_b64(&BASE64_NOPAD.encode(key.fingerprint.to_uppercase().as_bytes()))
+				.context("Unable to decode argon2 salt")?;
+		let argon2 = Argon2::default();
+		let hash = argon2
+			.hash_password(key_password.as_bytes(), &argon_salt)
+			.or(Err(anyhow!("Unable to derive encryption key")))?;
+		let aes_key = hash.hash.context("Unable to derive encryption key")?;
+		let aes_key = &aes_key.as_bytes()[0..32];
+
+		if let Ok(decrypted) = AspKey::from_encrypted(aes_key, &key.encrypted) {
+			let export = match self.format {
+				KeyExportFormat::Encrypted => decrypted
+					.export_encrypted(aes_key)
+					.context("Unable to convert key into encrypted format")?,
+				KeyExportFormat::PKCS8 => decrypted
+					.into_pkcs8()
+					.context("Unable to convert key into PKCS#8 format")?,
+				KeyExportFormat::Jwk => decrypted.jwk.to_string(),
+			};
+
 			eprintln!(
-				"The specified key fingerprint {fingerprint} does not exist",
-				fingerprint = self.fingerprint
-			)
+				"{style}Exported key \"{alias}\" with fingerprint {fingerprint}:{reset}",
+				style = Anstyle::new()
+							.bold()
+							.underline()
+							.fg_color(Some(AnstyleColor::Ansi(AnsiColor::BrightMagenta)))
+							.render(),
+				alias = key.alias,
+				fingerprint = decrypted.fingerprint,
+				reset = Reset.render()
+			);
+			let _ = writedoc!(
+				std::io::stdout(),
+				"{export}"
+			);
+		} else {
+			eprintln!("There was an error decrypting the key, please make sure the password you entered was correct");
 		}
 
 		Ok(())

src/commands/keys/list.rs

@@ -2,9 +2,11 @@ use anstyle::{AnsiColor, Reset, Style as Anstyle};
 use anyhow::Context;
 use asp::keys::AspKeyType;
 use clap::Parser;
-use indoc::printdoc;
+use indoc::{writedoc};
 use sea_orm::EntityTrait;
 
+use std::io::Write;
+
 use crate::{commands::AspmSubcommand, entities::keys};
 
 /// A command to list all saved keys, along with their fingerprints and types
@@ -38,21 +40,23 @@ impl AspmSubcommand for KeysListCommand {
 			.render();
 
 		// Print output
-		println!(
-			"{header_style}Saved keys ({n} total):{reset}\n",
+		let _ = writedoc!(
+			std::io::stdout(),
+			"{header_style}Saved keys ({n} total):{reset}\n\n",
 			n = entries.len(),
 		);
 		for entry in entries.iter() {
-			printdoc! {
+			let _ = writedoc! {
+				std::io::stdout(),
 				"
 					{alias_style}{alias}:{reset}
 						{key_style}Fingerprint{reset}  {value_style}{fingerprint}{reset}
-						{key_style}Key Type{reset}     {value_style}{key_type:?}{reset}
+						{key_style}Key Type{reset}     {value_style}{key_type:?}{reset}\n
 				",
 				fingerprint = entry.fingerprint,
 				key_type = TryInto::<AspKeyType>::try_into(entry.key_type).context("Unable to get key type from database")?,
 				alias = entry.alias
-			}
+			};
 		}
 
 		Ok(())

src/commands/mod.rs

@@ -1,10 +1,54 @@
 pub mod keys;
 
 use clap::Parser;
+use sea_orm::{ColumnTrait, Condition, DatabaseConnection, EntityTrait, QueryFilter, QueryOrder};
 
+use crate::entities::keys::{Column as KeysColumn, Entity as KeysEntity, Model as KeysModel};
 use crate::AspmState;
 
 #[async_trait::async_trait]
 pub trait AspmSubcommand: Parser {
 	async fn execute(&self, state: AspmState) -> Result<(), anyhow::Error>;
 }
+
+#[async_trait::async_trait]
+pub(self) trait KeysEntityExt {
+	type ResultEnum;
+
+	/// Queries the database for a specific key entity, first checking by fingerprint and then checking
+	async fn query_key(db: &DatabaseConnection, query: &str) -> anyhow::Result<Self::ResultEnum>;
+}
+
+pub(self) enum KeysQueryResult {
+	None,
+	One(KeysModel),
+	Many(Vec<KeysModel>),
+}
+
+#[async_trait::async_trait]
+impl KeysEntityExt for KeysEntity {
+	type ResultEnum = KeysQueryResult;
+
+	async fn query_key(db: &DatabaseConnection, query: &str) -> anyhow::Result<Self::ResultEnum> {
+		let mut models = KeysEntity::find()
+			.filter(
+				Condition::any()
+					.add(KeysColumn::Fingerprint.eq(query))
+					.add(KeysColumn::Alias.contains(query)),
+			)
+			.order_by(
+				// This order_by will assign a higher order priority if the fingerprint is EXACTLY the query, and everything else gets lower
+				// This effectively means that if there was an exact fingerprint match, it is ordered first, and any LIKE `%alias%` matches come after
+				KeysColumn::Fingerprint,
+				sea_orm::Order::Field(sea_orm::Values(vec![query.into()])),
+			)
+			.all(db)
+			.await?;
+
+		Ok(match models.len() {
+			0 => Self::ResultEnum::None,
+			1 => Self::ResultEnum::One(models.remove(0)),
+			_ => Self::ResultEnum::Many(models),
+		})
+	}
+}