From 954d6fac2e729aa87c7b2658aa21a3906c7035a9 Mon Sep 17 00:00:00 2001 From: TymanWasTaken Date: Sun, 2 Jul 2023 19:24:04 -0400 Subject: [PATCH] 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 | 3 +- Cargo.toml | 4 ++ crates/asp/src/keys/mod.rs | 3 +- src/commands/keys/export.rs | 124 +++++++++++++++++++++++------------- src/commands/keys/list.rs | 16 +++-- src/commands/mod.rs | 44 +++++++++++++ 6 files changed, 141 insertions(+), 53 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index 9154760..26184f7 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -6,6 +6,7 @@ "josekit", "PKCS", "Pkey", - "printdoc" + "printdoc", + "writedoc" ] } \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 6a337b6..67b1d17 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/asp/src/keys/mod.rs b/crates/asp/src/keys/mod.rs index cad277c..5f0d963 100644 --- a/crates/asp/src/keys/mod.rs +++ b/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> { diff --git a/src/commands/keys/export.rs b/src/commands/keys/export.rs index 17fe650..742b7fc 100644 --- a/src/commands/keys/export.rs +++ b/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(()) diff --git a/src/commands/keys/list.rs b/src/commands/keys/list.rs index 5891997..52c53a5 100644 --- a/src/commands/keys/list.rs +++ b/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::::try_into(entry.key_type).context("Unable to get key type from database")?, alias = entry.alias - } + }; } Ok(()) diff --git a/src/commands/mod.rs b/src/commands/mod.rs index d185605..c719a8b 100644 --- a/src/commands/mod.rs +++ b/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; +} + +pub(self) enum KeysQueryResult { + None, + One(KeysModel), + Many(Vec), +} + +#[async_trait::async_trait] +impl KeysEntityExt for KeysEntity { + type ResultEnum = KeysQueryResult; + + async fn query_key(db: &DatabaseConnection, query: &str) -> anyhow::Result { + 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), + }) + } +}