From 6f83bcccf94b88181d86358a922e61e3d3a2dad8 Mon Sep 17 00:00:00 2001 From: Awiteb Date: Sat, 11 May 2024 13:56:32 +0300 Subject: [PATCH] feat: Support TOTP --- src/cli/add_command.rs | 6 ++++- src/cli/edit_command.rs | 21 +++++++++++------ src/cli/get_command.rs | 32 +++++++++++++++++++++----- src/cli/list_command.rs | 31 +++++++++++++++++-------- src/vault/bitwarden.rs | 4 +++- src/vault/cipher.rs | 50 +++++++++++++++++++++++++++++++++++++++++ src/vault/mod.rs | 20 +++++++++++++++++ 7 files changed, 141 insertions(+), 23 deletions(-) diff --git a/src/cli/add_command.rs b/src/cli/add_command.rs index d712b1d..8c3fa8d 100644 --- a/src/cli/add_command.rs +++ b/src/cli/add_command.rs @@ -31,10 +31,14 @@ use crate::{ pub struct Add { #[command(flatten)] vault_info: Vault, - /// The password, if there is no value for it you will prompt it + /// The password, if there is no value you will prompt it #[arg(short, long)] #[allow(clippy::option_option)] password: Option>, + /// The TOTP secret, if there is no value you will prompt it + #[arg(short, long)] + #[allow(clippy::option_option)] + totp_secret: Option>, /// Add a custom field to the vault #[arg(name = "KEY=VALUE", short = 'c', long = "custom")] #[arg(value_parser = clap_parsers::kv_parser)] diff --git a/src/cli/edit_command.rs b/src/cli/edit_command.rs index 21e67a2..81e3c88 100644 --- a/src/cli/edit_command.rs +++ b/src/cli/edit_command.rs @@ -28,31 +28,35 @@ pub struct Edit { #[arg(short, long)] /// The new vault name - name: Option, + name: Option, #[arg(short, long)] /// The new vault username - username: Option, + username: Option, #[arg(short, long)] /// The new password, if there is no value for it you will prompt it #[allow(clippy::option_option)] - password: Option>, + password: Option>, #[arg(short, long)] /// The new vault service - service: Option, + service: Option, #[arg(short = 'o', long)] /// The new vault note - note: Option, + note: Option, + /// The TOTP secret, if there is no value you will prompt it + #[arg(short, long)] + #[allow(clippy::option_option)] + totp_secret: Option>, /// The custom field, make its value empty to delete it /// /// If the custom field not exist will created it, if it's will update it #[arg(name = "KEY=VALUE", short = 'c', long = "custom")] #[arg(value_parser = clap_parsers::kv_parser)] - pub custom_fields: Vec<(String, String)>, + custom_fields: Vec<(String, String)>, /// Force edit, will not return error if there is a problem with the args. /// /// For example, duplication in the custom fields and try to editing nothing #[arg(short, long)] - force: bool, + force: bool, } impl LprsCommand for Edit { @@ -75,6 +79,9 @@ impl LprsCommand for Edit { if self.password.is_some() { vault.password = utils::user_secret(self.password, "New vault password:")?; } + if self.totp_secret.is_some() { + vault.totp_secret = utils::user_secret(self.totp_secret, "TOTP Secret:")?; + } if let Some(new_username) = self.username { vault.username = Some(new_username); } diff --git a/src/cli/get_command.rs b/src/cli/get_command.rs index 73f72eb..6c6a39e 100644 --- a/src/cli/get_command.rs +++ b/src/cli/get_command.rs @@ -20,7 +20,7 @@ use clap::Args; use crate::{ utils, - vault::{Vault, Vaults}, + vault::{cipher, Vault, Vaults}, LprsCommand, LprsError, LprsResult, @@ -34,6 +34,8 @@ enum VaultGetField { Password, Service, Note, + TotpSecret, + TotpCode, Custom(String), } @@ -48,6 +50,8 @@ impl FromStr for VaultGetField { "password" => Self::Password, "service" => Self::Service, "note" => Self::Note, + "totp_secret" => Self::TotpSecret, + "totp_code" => Self::TotpCode, _ => Self::Custom(input.to_owned()), }) } @@ -63,6 +67,8 @@ impl VaultGetField { Self::Password => vault.password.as_deref(), Self::Service => vault.service.as_deref(), Self::Note => vault.note.as_deref(), + Self::TotpSecret => vault.totp_secret.as_deref(), + Self::TotpCode => None, Self::Custom(custom_field) => vault.custom_fields.get(custom_field).map(|x| x.as_str()), } } @@ -76,6 +82,8 @@ impl VaultGetField { Self::Password => "password", Self::Service => "service", Self::Note => "note", + Self::TotpSecret => "totp_secret", + Self::TotpCode => "totp_code", Self::Custom(field) => field.as_str(), } } @@ -90,8 +98,10 @@ pub struct Get { location: String, /// A Specific field to get. /// - /// Can be [name,username,password,service,note,"string"] where the string - /// means a custom field + /// Can be [name, username, password, service, note, totp_secret, totp_code, + /// "string"] + /// + /// where the string means a custom field #[arg(value_parser = VaultGetField::from_str)] field: Option, } @@ -106,13 +116,25 @@ impl LprsCommand for Get { print!("{index}"); return Ok(()); } + if field == VaultGetField::TotpCode { + if let Some(ref totp_secret) = vault.totp_secret { + let totp_code = cipher::totp_now(totp_secret, &vault.totp_hash)?.0; + print!("{totp_code}"); + return Ok(()); + } else { + return Err(LprsError::Other( + "There is no TOTP secret to get TOTP code".to_owned(), + )); + } + } if let Some(value) = field.get_from_vault(vault) { print!("{value}") } else { return Err(LprsError::Other(format!( - "There is no value for `{}`", - field.as_str() + "There is no value for `{}` at \"{}\" vault", + field.as_str(), + vault.name ))); } } else { diff --git a/src/cli/list_command.rs b/src/cli/list_command.rs index 2ae356b..a9db358 100644 --- a/src/cli/list_command.rs +++ b/src/cli/list_command.rs @@ -17,7 +17,13 @@ use clap::Args; use inquire::{InquireError, Select}; -use crate::{vault::Vaults, LprsCommand, LprsError, LprsResult}; +use crate::{ + vault::{cipher, Vaults}, + LprsCommand, + LprsError, + LprsResult, + RESERVED_FIELD_PREFIX, +}; #[derive(Debug, Args)] #[command(author, version, about, long_about = None)] @@ -35,7 +41,7 @@ pub struct List { } impl LprsCommand for List { - fn run(self, vault_manager: Vaults) -> LprsResult<()> { + fn run(self, mut vault_manager: Vaults) -> LprsResult<()> { if vault_manager.vaults.is_empty() { return Err(LprsError::Other( "Looks like there is no vaults to list".to_owned(), @@ -95,13 +101,20 @@ impl LprsCommand for List { log::debug!("The user selected the vault at index: {vault_idx}"); - println!( - "{}", - vault_manager - .vaults - .get(vault_idx - 1) - .expect("The index is correct") - ); + let vault = vault_manager + .vaults + .get_mut(vault_idx - 1) + .expect("The index is correct"); + + if let Some(ref totp_secret) = vault.totp_secret { + let (code, remaining) = cipher::totp_now(totp_secret, &vault.totp_hash)?; + vault.custom_fields.insert( + format!("{RESERVED_FIELD_PREFIX}TOTP Code"), + format!("{code} ({remaining}s remaining)"), + ); + } + + println!("{vault}",); } Ok(()) diff --git a/src/vault/bitwarden.rs b/src/vault/bitwarden.rs index 8e66576..ed22bdf 100644 --- a/src/vault/bitwarden.rs +++ b/src/vault/bitwarden.rs @@ -20,7 +20,7 @@ use serde::{Deserialize, Serialize}; -use super::{Vault, Vaults}; +use super::{cipher::TotpHash, Vault, Vaults}; #[derive(Clone, Debug, Deserialize, Serialize)] pub struct BitWardenLoginData { @@ -84,6 +84,8 @@ impl From for Vault { .into_iter() .map(|nv| (nv.name, nv.value)) .collect(), + None::, + TotpHash::default(), ) } } diff --git a/src/vault/cipher.rs b/src/vault/cipher.rs index a627a21..7422015 100644 --- a/src/vault/cipher.rs +++ b/src/vault/cipher.rs @@ -17,13 +17,63 @@ use std::time::{SystemTime, UNIX_EPOCH}; use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit}; +use base32::Alphabet as Base32Alphabet; +use clap::ValueEnum; use rand::{rngs::StdRng, Rng, SeedableRng}; +use serde::{Deserialize, Serialize}; use crate::{LprsError, LprsResult}; type Aes256CbcEnc = cbc::Encryptor; type Aes256CbcDec = cbc::Decryptor; +#[derive(Default, Clone, Debug, ValueEnum, Eq, PartialEq, Deserialize, Serialize)] +/// The TOTP hash functions +pub enum TotpHash { + /// Sha1 hash function + #[default] + Sha1, + /// Sha256 hash function + Sha256, + /// Sha512 hash function + Sha512, +} + + +/// Create the TOTP code of the current time +/// +/// ## Errors +/// - If the given `secret_base32` are vaild base32 +pub fn totp_now(secret_base32: &str, hash_function: &TotpHash) -> LprsResult<(String, u8)> { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("SystemTime before UNIX EPOCH!") + .as_secs(); + let remaining = 30 - (now % 30) as u8; + let secret = base32::decode(Base32Alphabet::RFC4648 { padding: true }, secret_base32) + .ok_or_else(|| LprsError::Base32("Can't decode the TOTP secret".to_owned()))?; + Ok(match hash_function { + TotpHash::Sha1 => { + ( + totp_lite::totp_custom::(30, 6, &secret, now), + remaining, + ) + } + TotpHash::Sha256 => { + ( + totp_lite::totp_custom::(30, 6, &secret, now), + remaining, + ) + } + TotpHash::Sha512 => { + ( + totp_lite::totp_custom::(30, 6, &secret, now), + remaining, + ) + } + }) +} + /// Encrypt the given data by the given key using AES-256 CBC /// /// Note: The IV will be add it to the end of the ciphertext (Last 16 bytes) diff --git a/src/vault/mod.rs b/src/vault/mod.rs index 8e62884..be76a67 100644 --- a/src/vault/mod.rs +++ b/src/vault/mod.rs @@ -29,6 +29,8 @@ mod bitwarden; pub use bitwarden::*; +use self::cipher::TotpHash; + #[derive(Clone, Debug, ValueEnum, Eq, PartialEq)] /// The vaults format pub enum Format { @@ -61,6 +63,12 @@ pub struct Vault { /// The vault custom fields #[arg(skip)] pub custom_fields: BTreeMap, + /// The TOTP secret + #[arg(skip)] + pub totp_secret: Option, + /// The TOTP hash function + #[arg(long, value_name = "HASH_FUNCTION", value_enum, default_value_t)] + pub totp_hash: TotpHash, } /// The vaults manager @@ -76,6 +84,7 @@ pub struct Vaults { impl Vault { /// Create new [`Vault`] instance + #[allow(clippy::too_many_arguments)] pub fn new( name: impl Into, username: Option>, @@ -83,6 +92,8 @@ impl Vault { service: Option>, note: Option>, custom_fields: BTreeMap, + totp_secret: Option>, + totp_hash: TotpHash, ) -> Self { Self { name: name.into(), @@ -91,6 +102,8 @@ impl Vault { service: service.map(Into::into), note: note.map(Into::into), custom_fields, + totp_secret: totp_secret.map(Into::into), + totp_hash, } } @@ -163,6 +176,8 @@ impl Vaults { .iter() .map(|(key, value)| (encrypt(key), encrypt(value))) .collect(), + v.totp_secret.as_ref().map(|t| encrypt(t)), + v.totp_hash.clone(), )) }) .collect::>>()?, @@ -199,6 +214,8 @@ impl Vaults { .into_iter() .map(|(key, value)| LprsResult::Ok((decrypt(&key)?, decrypt(&value)?))) .collect::>()?, + v.totp_secret.as_ref().and_then(|t| decrypt(t).ok()), + v.totp_hash, )) }) .collect() @@ -266,6 +283,9 @@ impl fmt::Display for Vault { if let Some(ref note) = self.note { write!(f, "\nNote:\n{note}")?; } + if let Some(ref totp_secret) = self.totp_secret { + write!(f, "\nTOTP Secret: {totp_secret}")?; + } for (key, value) in &self.custom_fields { write!( f,