diff --git a/src/commands/aspe/create.rs b/src/commands/aspe/create.rs index 6a46ff3..d147d4e 100644 --- a/src/commands/aspe/create.rs +++ b/src/commands/aspe/create.rs @@ -23,7 +23,7 @@ use crate::{ entities::{prelude::*, profiles}, }; -/// +/// Upload an existing profile to an ASPE server #[derive(Parser, Debug)] pub struct AspeCreateCommand { /// The fingerprint or alias of the profile to upload @@ -130,7 +130,10 @@ impl AspmSubcommand for AspeCreateCommand { .context("Unable to encode the profile as a JWT and sign it")?; match server.post_request(encoded_request).await { - Ok(_) => Ok(()), + Ok(_) => { + println!("Successfully uploaded profile!"); + Ok(()) + }, Err(AspeRequestFailure::BadRequest) => bail!("The ASPE server rejected the request due to invalid data"), Err(AspeRequestFailure::TooLarge) => bail!("The ASPE server rejected the request as being too large"), Err(AspeRequestFailure::RateLimited) => bail!("The ASPE server rejected the request due to a ratelimit"), diff --git a/src/commands/keys/delete.rs b/src/commands/keys/delete.rs index b130a07..e730ae4 100644 --- a/src/commands/keys/delete.rs +++ b/src/commands/keys/delete.rs @@ -1,16 +1,16 @@ -use anstyle::{AnsiColor, Color as AnstyleColor, Reset, Style as Anstyle}; +use anstyle::{AnsiColor, Reset, Style as Anstyle}; use anyhow::Context; use asp::keys::AspKeyType; use clap::Parser; use dialoguer::{console::Term, theme::ColorfulTheme, Confirm}; use indoc::writedoc; -use sea_orm::ModelTrait; +use sea_orm::ModelTrait as _; use std::io::Write; use crate::{ - commands::{AspmSubcommand, KeysEntityExt, KeysQueryResult}, - entities::prelude::*, + commands::AspmSubcommand, + utils, }; /// Deletes a saved key, after asking for confirmation. @@ -27,32 +27,7 @@ pub struct KeysDeleteCommand { impl AspmSubcommand for KeysDeleteCommand { async fn execute(self, state: crate::AspmState) -> Result<(), anyhow::Error> { // Fetch key from db - let entry = Keys::query_key(&state.db, &self.key) - .await - .context("Unable to query keys from database")?; - let key = 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); - } - 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.render() - ); - keys.remove(0) - } - }; + let key = utils::keys::query_key(&state.db, self.key).await?; if !self.no_confirm { // Construct styles diff --git a/src/commands/keys/export.rs b/src/commands/keys/export.rs index c465ced..6894e3a 100644 --- a/src/commands/keys/export.rs +++ b/src/commands/keys/export.rs @@ -3,8 +3,7 @@ use aes_gcm::{ Aes256Gcm, Key, }; use anstyle::{AnsiColor, Color as AnstyleColor, Reset, Style as Anstyle}; -use anyhow::{anyhow, Context}; -use argon2::{password_hash::SaltString, Argon2, PasswordHasher}; +use anyhow::Context as _; use asp::keys::AspKey; use clap::{Parser, ValueEnum}; use data_encoding::{BASE64, BASE64_NOPAD}; @@ -12,10 +11,7 @@ use dialoguer::{theme::ColorfulTheme, Password}; use josekit::jwk::Jwk; use serde::{Deserialize, Serialize}; -use crate::{ - commands::{AspmSubcommand, KeysEntityExt, KeysQueryResult}, - entities::prelude::*, -}; +use crate::{commands::AspmSubcommand, utils}; #[derive(Serialize, Deserialize)] #[serde(tag = "alg", rename = "scrypt")] @@ -63,32 +59,7 @@ pub struct KeysExportCommand { impl AspmSubcommand for KeysExportCommand { async fn execute(self, state: crate::AspmState) -> Result<(), anyhow::Error> { // Fetch key from db - let entry = Keys::query_key(&state.db, &self.key) - .await - .context("Unable to query keys from database")?; - let key = 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); - } - 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.render() - ); - keys.remove(0) - } - }; + let key = utils::keys::query_key(&state.db, self.key).await?; let key_password = std::env::var("KEY_PASSWORD").or_else(|_| { Password::with_theme(&ColorfulTheme::default()) @@ -97,18 +68,10 @@ impl AspmSubcommand for KeysExportCommand { .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]; + let derived_key = utils::keys::derive_encryption_key(key.fingerprint, &key_password)?; - if let Ok(decrypted) = Aes256Gcm::new(Key::::from_slice(aes_key)) - .decrypt((&aes_key[0..12]).into(), key.cipher_text.as_slice()) + if let Ok(decrypted) = Aes256Gcm::new(Key::::from_slice(&derived_key)) + .decrypt((&derived_key[0..12]).into(), key.cipher_text.as_slice()) { let decrypted = AspKey::from_jwk(Jwk::from_bytes(decrypted)?)?; diff --git a/src/commands/mod.rs b/src/commands/mod.rs index cfb67ee..2f9f128 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -3,10 +3,8 @@ pub mod profiles; pub mod aspe; use clap::Parser; -use sea_orm::{ColumnTrait, Condition, DatabaseConnection, EntityTrait, QueryFilter, QueryOrder}; use tokio::runtime::Runtime; -use crate::entities::keys::{Column as KeysColumn, Entity as KeysEntity, Model as KeysModel}; use crate::AspmState; #[async_trait::async_trait] @@ -18,45 +16,3 @@ pub trait AspmSubcommand: Parser + Sync + Send { runtime.block_on(self.execute(state)) } } - -#[async_trait::async_trait] -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; -} - -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), - }) - } -} diff --git a/src/main.rs b/src/main.rs index 7c3e945..77e2a3d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ mod commands; #[allow(warnings)] // This is autogenerated, no use in showing warnings mod entities; +mod utils; use anstyle::{AnsiColor, Color as AnstyleColor, Style as Anstyle}; use anyhow::Context; diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..c0fc618 --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,112 @@ +pub mod keys { + use anstyle::{AnsiColor, Color as AnstyleColor, Reset, Style as Anstyle}; + use anyhow::{anyhow, Context as _}; + use argon2::{password_hash::SaltString, Argon2, PasswordHasher}; + use data_encoding::BASE64_NOPAD; + use sea_orm::{ + ColumnTrait as _, Condition, DatabaseConnection, EntityTrait as _, QueryFilter as _, + QueryOrder as _, + }; + + use crate::entities::{ + keys::{Column as KeysColumn, Entity as KeysEntity, Model as KeysModel}, + prelude::Keys, + }; + + pub fn derive_encryption_key, P: AsRef<[u8]>>( + fingerprint: S, + password: P, + ) -> anyhow::Result> { + let argon_salt = + SaltString::from_b64(&BASE64_NOPAD.encode(fingerprint.as_ref().to_uppercase().as_bytes())) + .context("Unable to derive argon2 salt")?; + let argon2 = Argon2::default(); + let hash = argon2 + .hash_password(password.as_ref(), &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]; + + Ok(aes_key.to_vec()) + } + + #[async_trait::async_trait] + 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 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), + }) + } + } + + pub async fn query_key>( + db: &DatabaseConnection, + query: S, + ) -> anyhow::Result { + let entry = Keys::query_key(db, query.as_ref()) + .await + .context("Unable to query keys from database")?; + 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); + } + KeysQueryResult::One(key) => Ok(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.render() + ); + Ok(keys.remove(0)) + } + } + } +}