From 386f060a41212e678052c6a055084f35e0deca4d Mon Sep 17 00:00:00 2001 From: StNicolay Date: Fri, 4 Aug 2023 00:23:02 +0300 Subject: [PATCH] Implemented account altering --- cryptography/src/account.rs | 6 +- entity/src/account.rs | 37 +++++++++- src/callbacks/alter.rs | 127 ++++++++++++++++++++++++++++++++ src/callbacks/decrypt.rs | 4 +- src/callbacks/delete.rs | 4 +- src/callbacks/mod.rs | 2 +- src/commands/set_master_pass.rs | 4 +- src/macros.rs | 9 ++- src/main.rs | 3 +- src/utils.rs | 2 +- 10 files changed, 185 insertions(+), 13 deletions(-) create mode 100644 src/callbacks/alter.rs diff --git a/cryptography/src/account.rs b/cryptography/src/account.rs index 628a484..33bc0b5 100644 --- a/cryptography/src/account.rs +++ b/cryptography/src/account.rs @@ -5,14 +5,14 @@ use rand::{rngs::OsRng, RngCore}; use sea_orm::ActiveValue::Set; use sha2::Sha256; -struct Cipher { +pub struct Cipher { chacha: ChaCha20Poly1305, } impl Cipher { /// Creates a new cipher from a master password and the salt #[inline] - fn new(password: &[u8], salt: &[u8]) -> Self { + pub fn new(password: &[u8], salt: &[u8]) -> Self { let key = pbkdf2_hmac_array::(password, salt, 480000); Self { @@ -31,7 +31,7 @@ impl Cipher { /// Decrypts the value with the current cipher. The 12 byte nonce is expected to be at the end of the value #[inline] - fn decrypt(&self, value: &[u8]) -> crate::Result> { + pub fn decrypt(&self, value: &[u8]) -> crate::Result> { let (data, nonce) = value.split_at(value.len() - 12); self.chacha.decrypt(nonce.into(), data).map_err(Into::into) } diff --git a/entity/src/account.rs b/entity/src/account.rs index 8f565c7..c395f12 100644 --- a/entity/src/account.rs +++ b/entity/src/account.rs @@ -1,7 +1,7 @@ //! `SeaORM` Entity. Generated by sea-orm-codegen 0.11.3 use futures::Stream; -use sea_orm::{entity::prelude::*, QueryOrder, QuerySelect, Statement}; +use sea_orm::{entity::prelude::*, ActiveValue::Set, QueryOrder, QuerySelect, Statement}; #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)] #[sea_orm(table_name = "account")] @@ -88,6 +88,7 @@ impl Entity { } /// Gets a name by a hex of a SHA256 hash of the name + #[inline] pub async fn get_name_by_hash( user_id: u64, hash: String, @@ -102,4 +103,38 @@ impl Entity { .map(|result| result.try_get_by_index(0)) .transpose() } + + #[inline] + pub async fn get_salt( + user_id: u64, + name: String, + db: &DatabaseConnection, + ) -> crate::Result>> { + Self::find_by_id((user_id, name)) + .select_only() + .column(Column::Salt) + .into_tuple() + .one(db) + .await + } + + #[inline] + pub async fn update_name( + user_id: u64, + original_name: String, + new_name: String, + db: &DatabaseConnection, + ) -> crate::Result<()> { + Self::update_many() + .set(ActiveModel { + name: Set(new_name), + ..Default::default() + }) + .filter(Column::UserId.eq(user_id)) + .filter(Column::Name.eq(original_name)) + .exec(db) + .await?; + + Ok(()) + } } diff --git a/src/callbacks/alter.rs b/src/callbacks/alter.rs new file mode 100644 index 0000000..1fe4272 --- /dev/null +++ b/src/callbacks/alter.rs @@ -0,0 +1,127 @@ +use super::AlterableField::{self, *}; +use crate::{change_state, prelude::*}; +use account::ActiveModel; +use futures::TryFutureExt; +use sea_orm::ActiveValue::Set; +use tokio::{task::spawn_blocking, try_join}; + +#[inline] +async fn update_account( + user_id: u64, + db: &DatabaseConnection, + name: String, + field: AlterableField, + field_value: String, + master_pass: String, +) -> crate::Result<()> { + if let Name = field { + Account::update_name(user_id, name, field_value, db).await?; + + return Ok(()); + } + + let salt = Account::get_salt(user_id, name.clone(), db).await?.unwrap(); + + let field_value = spawn_blocking(move || { + let cipher = Cipher::new(master_pass.as_bytes(), &salt); + cipher.encrypt(field_value.as_bytes()).unwrap() + }) + .await?; + + let mut model = ActiveModel { + user_id: Set(user_id), + name: Set(name), + ..Default::default() + }; + + match field { + Login => model.enc_login = Set(field_value), + Pass => model.enc_password = Set(field_value), + _ => unreachable!(), + } + + model.update(db).await?; + + Ok(()) +} + +#[inline] +#[allow(clippy::too_many_arguments)] +async fn get_master_pass( + bot: Throttle, + msg: Message, + db: DatabaseConnection, + dialogue: MainDialogue, + mut ids: MessageIds, + name: String, + field: AlterableField, + field_value: String, + master_pass: String, +) -> crate::Result<()> { + dialogue.exit().await?; + let user_id = msg.from().ok_or(NoUserInfo)?.id.0; + + update_account(user_id, &db, name, field, field_value, master_pass).await?; + + ids.alter_message( + &bot, + "Success. Choose the account to view", + menu_markup("get", user_id, &db).await?, + None, + ) + .await?; + + Ok(()) +} + +handler!(get_field(name:String, field:AlterableField, field_value:String), "Send the master password", State::GetMasterPass, get_master_pass); + +pub async fn alter( + bot: Throttle, + q: CallbackQuery, + db: DatabaseConnection, + dialogue: MainDialogue, + (hash, field): (super::NameHash, AlterableField), +) -> crate::Result<()> { + let user_id = q.from.id.0; + let mut ids: MessageIds = q.message.as_ref().unwrap().into(); + + let name = match name_from_hash(&db, user_id, &hash).await? { + Some(name) => name, + None => { + try_join!( + ids.alter_message( + &bot, + "Account wasn't found. Choose another one", + menu_markup("get", user_id, &db).await?, + None, + ), + bot.answer_callback_query(q.id).send().err_into() + )?; + + return Ok(()); + } + }; + + let text = match field { + Name => { + change_state!(dialogue, ids, (name, field), State::GetNewName, get_field); + "Send new account name" + } + Login => { + change_state!(dialogue, ids, (name, field), State::GetLogin, get_field); + "Send new account login" + } + Pass => { + change_state!(dialogue, ids, (name, field), State::GetPassword, get_field); + "Send new account password" + } + }; + + try_join!( + ids.alter_message(&bot, text, None, None), + bot.answer_callback_query(q.id).send().err_into() + )?; + + Ok(()) +} diff --git a/src/callbacks/decrypt.rs b/src/callbacks/decrypt.rs index 3635230..5ae0cdf 100644 --- a/src/callbacks/decrypt.rs +++ b/src/callbacks/decrypt.rs @@ -66,5 +66,7 @@ pub async fn decrypt( .await?; bot.answer_callback_query(q.id).await?; - change_state!(dialogue, msg, (name), State::GetMasterPass, get_master_pass) + change_state!(dialogue, msg, (name), State::GetMasterPass, get_master_pass); + + Ok(()) } diff --git a/src/callbacks/delete.rs b/src/callbacks/delete.rs index f37ce02..9b1d160 100644 --- a/src/callbacks/delete.rs +++ b/src/callbacks/delete.rs @@ -61,5 +61,7 @@ pub async fn delete( (name), State::GetMasterPass, get_master_pass - ) + ); + + Ok(()) } diff --git a/src/callbacks/mod.rs b/src/callbacks/mod.rs index 7df3de5..7a39d81 100644 --- a/src/callbacks/mod.rs +++ b/src/callbacks/mod.rs @@ -1,6 +1,6 @@ //! This module consists of endpoints to handle callbacks -crate::export_handlers!(decrypt, delete, delete_message, get, get_menu); +crate::export_handlers!(decrypt, delete, delete_message, get, get_menu, alter); use crate::errors::InvalidCommand; use base64::{engine::general_purpose::STANDARD_NO_PAD as B64_ENGINE, Engine as _}; diff --git a/src/commands/set_master_pass.rs b/src/commands/set_master_pass.rs index 70114a0..3797041 100644 --- a/src/commands/set_master_pass.rs +++ b/src/commands/set_master_pass.rs @@ -46,5 +46,7 @@ pub async fn set_master_pass( (), State::GetNewMasterPass, get_master_pass - ) + ); + + Ok(()) } diff --git a/src/macros.rs b/src/macros.rs index ccc8534..5e8778e 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -7,7 +7,6 @@ macro_rules! change_state { $previous, ))) .await?; - Ok(()) }}; } @@ -22,7 +21,9 @@ macro_rules! first_handler { ) -> $crate::Result<()> { let previous = bot.send_message(msg.chat.id, $message).await?; - $crate::change_state!(dialogue, &previous, (), $next_state, $next_func) + $crate::change_state!(dialogue, &previous, (), $next_state, $next_func); + + Ok(()) } }; } @@ -42,7 +43,9 @@ macro_rules! handler { ) -> $crate::Result<()> { ids.alter_message(&bot, $message, None, None).await?; - $crate::change_state!(dialogue, ids, ($($param),*), $next_state, $next_func) + $crate::change_state!(dialogue, ids, ($($param),*), $next_state, $next_func); + + Ok(()) } }; } diff --git a/src/main.rs b/src/main.rs index 4e31dbe..f073bb5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -62,7 +62,8 @@ fn get_dispatcher( .branch(case![CallbackCommand::DeleteMessage].endpoint(callbacks::delete_message)) .branch(case![CallbackCommand::Get(hash)].endpoint(callbacks::get)) .branch(case![CallbackCommand::Decrypt(hash)].endpoint(callbacks::decrypt)) - .branch(case![CallbackCommand::DeleteAccount(hash)].endpoint(callbacks::delete)); + .branch(case![CallbackCommand::DeleteAccount(hash)].endpoint(callbacks::delete)) + .branch(case![CallbackCommand::Alter(hash, field)].endpoint(callbacks::alter)); let handler = dptree::entry() .enter_dialogue::, State>() diff --git a/src/utils.rs b/src/utils.rs index 95280d8..33b1093 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -9,7 +9,7 @@ pub async fn delete_message(bot: Throttle, msg: Message) { /// Returns true if the field is valid #[inline] pub fn validate_field(field: &str) -> bool { - if field.is_empty() { + if !(1..255).contains(&field.len()) { return false; } field