From 22c754a2564a92918dc35934a2157a457e3159b6 Mon Sep 17 00:00:00 2001 From: StNicolay Date: Wed, 26 Jul 2023 00:52:12 +0300 Subject: [PATCH] added get endpoint, updated state_management --- entity/src/account.rs | 18 +++++++- src/callbacks/delete_message.rs | 1 - src/callbacks/get.rs | 35 ++++++++++++++++ src/callbacks/get_menu.rs | 2 +- src/callbacks/mod.rs | 37 +++++++++-------- src/commands/add_account.rs | 9 +++- src/commands/delete.rs | 9 +++- src/commands/delete_all.rs | 6 ++- src/commands/export.rs | 9 +++- src/commands/get_account.rs | 9 +++- src/commands/import.rs | 9 +++- src/commands/menu.rs | 2 +- src/commands/set_master_pass.rs | 9 ++-- src/errors.rs | 4 +- src/macros.rs | 71 ++++++++++++++++++++++---------- src/main.rs | 3 +- src/markups.rs | 59 ++++++++++++++++++++++++-- src/prelude.rs | 2 +- src/state/generic.rs | 11 +++-- src/state/get_existing_name.rs | 34 ++++----------- src/state/get_login.rs | 10 ++--- src/state/get_master_pass.rs | 19 +++------ src/state/get_new_master_pass.rs | 13 ++---- src/state/get_new_name.rs | 23 ++++------- src/state/get_password.rs | 14 +++---- src/state/get_user.rs | 5 +-- src/state/handler.rs | 4 +- 27 files changed, 278 insertions(+), 149 deletions(-) create mode 100644 src/callbacks/get.rs diff --git a/entity/src/account.rs b/entity/src/account.rs index f8e1ab2..8f565c7 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}; +use sea_orm::{entity::prelude::*, QueryOrder, QuerySelect, Statement}; #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)] #[sea_orm(table_name = "account")] @@ -86,4 +86,20 @@ impl Entity { .await?; Ok(()) } + + /// Gets a name by a hex of a SHA256 hash of the name + pub async fn get_name_by_hash( + user_id: u64, + hash: String, + db: &DatabaseConnection, + ) -> crate::Result> { + db.query_one(Statement::from_sql_and_values( + sea_orm::DatabaseBackend::MySql, + "SELECT `name` FROM `account` WHERE SHA2(`name`, 256) = ? AND `user_id` = ?;", + [hash.into(), user_id.into()], + )) + .await? + .map(|result| result.try_get_by_index(0)) + .transpose() + } } diff --git a/src/callbacks/delete_message.rs b/src/callbacks/delete_message.rs index 04a337d..cbfd235 100644 --- a/src/callbacks/delete_message.rs +++ b/src/callbacks/delete_message.rs @@ -1,5 +1,4 @@ use crate::prelude::*; -use teloxide::{dispatching::DpHandlerDescription, dptree::Handler}; /// Deletes the message from the callback pub async fn delete_message(bot: Throttle, q: CallbackQuery) -> crate::Result<()> { diff --git a/src/callbacks/get.rs b/src/callbacks/get.rs new file mode 100644 index 0000000..8c6948a --- /dev/null +++ b/src/callbacks/get.rs @@ -0,0 +1,35 @@ +use crate::prelude::*; +use teloxide::types::ParseMode; + +pub async fn get( + bot: Throttle, + q: CallbackQuery, + db: DatabaseConnection, + hash: super::NameHash, +) -> crate::Result<()> { + let user_id = q.from.id.0; + let msg = q.message.as_ref().unwrap(); + + let hash = hex::encode(hash); + let name = match Account::get_name_by_hash(user_id, hash, &db).await? { + Some(name) => name, + None => { + bot.edit_message_text( + msg.chat.id, + msg.id, + "Account wan't found. Select another one", + ) + .reply_markup(menu_markup(user_id, &db).await?) + .await?; + return Ok(()); + } + }; + + let text = format!("Name:\n`{name}`\nLogin:\n***\nPassword:\n***"); + bot.send_message(msg.chat.id, text) + .reply_markup(account_markup(&name, true)) + .parse_mode(ParseMode::MarkdownV2) + .await?; + + Ok(()) +} diff --git a/src/callbacks/get_menu.rs b/src/callbacks/get_menu.rs index 32cbb94..ea0eadd 100644 --- a/src/callbacks/get_menu.rs +++ b/src/callbacks/get_menu.rs @@ -20,7 +20,7 @@ pub async fn get_menu( .await?; } - let markup = spawn_blocking(|| menu_markup(names)).await?; + let markup = spawn_blocking(|| menu_markup_sync(names)).await?; bot.edit_message_text(msg.chat.id, msg.id, "Choose your account") .reply_markup(markup) .await?; diff --git a/src/callbacks/mod.rs b/src/callbacks/mod.rs index 0b694e7..8e00e18 100644 --- a/src/callbacks/mod.rs +++ b/src/callbacks/mod.rs @@ -1,26 +1,28 @@ //! This module consists of endpoints to handle callbacks mod delete_message; +mod get; mod get_menu; pub use delete_message::delete_message; +pub use get::get; pub use get_menu::get_menu; use crate::errors::InvalidCommand; use base64::{engine::general_purpose::STANDARD_NO_PAD as B64_ENGINE, Engine as _}; use std::str::FromStr; -use teloxide::types::{CallbackQuery, Message}; +use teloxide::types::CallbackQuery; -type NameHash = [u8; 32]; +type NameHash = Vec; -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Debug)] pub enum AlterableField { Name, Login, Pass, } -#[derive(Clone, Copy)] +#[derive(Clone, Debug)] pub enum CallbackCommand { DeleteMessage, Get(NameHash), @@ -34,7 +36,7 @@ pub enum CallbackCommand { impl CallbackCommand { pub fn from_query(q: CallbackQuery) -> Option { q.message.as_ref()?; - q.data.and_then(|text| text.parse().ok()) + q.data.and_then(|text| dbg!(text.parse()).ok()) } } @@ -45,6 +47,8 @@ impl FromStr for CallbackCommand { use AlterableField::*; use CallbackCommand::*; + println!("{s}"); + match s { "delete_message" => return Ok(DeleteMessage), "get_menu" => return Ok(GetMenu), @@ -56,20 +60,19 @@ impl FromStr for CallbackCommand { (Some(command), Some(name), None) => (command, name), _ => return Err(InvalidCommand::InvalidParams), }; - - let mut name_arr = [0; 32]; - if B64_ENGINE.decode_slice(name, &mut name_arr)? != 32 { - return Err(InvalidCommand::OutputTooShort); - }; + let name_hash = B64_ENGINE.decode(name)?; + if name_hash.len() != 32 { + return Err(InvalidCommand::InvalidOutputLength); + } match command { - "get" => Ok(Get(name_arr)), - "decrypt" => Ok(Decrypt(name_arr)), - "hide" => Ok(Hide(name_arr)), - "an" => Ok(Alter(name_arr, Name)), - "al" => Ok(Alter(name_arr, Login)), - "ap" => Ok(Alter(name_arr, Pass)), - "delete" => Ok(DeleteAccount(name_arr)), + "get" => Ok(Get(name_hash)), + "decrypt" => Ok(Decrypt(name_hash)), + "hide" => Ok(Hide(name_hash)), + "an" => Ok(Alter(name_hash, Name)), + "al" => Ok(Alter(name_hash, Login)), + "ap" => Ok(Alter(name_hash, Pass)), + "delete" => Ok(DeleteAccount(name_hash)), _ => Err(InvalidCommand::InvalidParams), } } diff --git a/src/commands/add_account.rs b/src/commands/add_account.rs index 1634b21..abe5bbf 100644 --- a/src/commands/add_account.rs +++ b/src/commands/add_account.rs @@ -8,11 +8,13 @@ async fn get_master_pass( msg: Message, db: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, name: String, login: String, password: String, master_pass: String, ) -> crate::Result<()> { + ids.delete(&bot).await; let user_id = msg.from().ok_or(NoUserInfo)?.id.0; dialogue.exit().await?; let account = spawn_blocking(move || { @@ -38,4 +40,9 @@ handler!(get_login(name: String, login: String), get_password ); handler!(get_account_name(name: String), "Send login", State::GetLogin, get_login); -handler!(pub add_account(), "Send account name", State::GetNewName, get_account_name); +first_handler!( + add_account, + "Send account name", + State::GetNewName, + get_account_name +); diff --git a/src/commands/delete.rs b/src/commands/delete.rs index 0519365..03fb5f2 100644 --- a/src/commands/delete.rs +++ b/src/commands/delete.rs @@ -7,10 +7,13 @@ async fn get_master_pass( msg: Message, db: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, name: String, _: String, ) -> crate::Result<()> { + ids.delete(&bot).await; dialogue.exit().await?; + let user_id = msg.from().ok_or(NoUserInfo)?.id.0; Account::delete_by_id((user_id, name)).exec(&db).await?; bot.send_message(msg.chat.id, "The account is successfully deleted") @@ -25,4 +28,8 @@ handler!( State::GetMasterPass, get_master_pass ); -ask_name_handler!(pub delete(), "Send the name of the account to delete", get_account_name); +ask_name_handler!( + delete, + "Send the name of the account to delete", + get_account_name +); diff --git a/src/commands/delete_all.rs b/src/commands/delete_all.rs index 59a90e7..7abf40d 100644 --- a/src/commands/delete_all.rs +++ b/src/commands/delete_all.rs @@ -10,8 +10,10 @@ async fn get_master_pass( msg: Message, db: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, _: String, ) -> crate::Result<()> { + ids.delete(&bot).await; dialogue.exit().await?; let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let txn = db.begin().await?; @@ -36,8 +38,8 @@ async fn get_master_pass( Ok(()) } -handler!( - pub delete_all(), +first_handler!( + delete_all, "Send master password to delete EVERYTHING.\nTHIS ACTION IS IRREVERSIBLE", State::GetMasterPass, get_master_pass diff --git a/src/commands/export.rs b/src/commands/export.rs index 4ee67c0..50cc2bf 100644 --- a/src/commands/export.rs +++ b/src/commands/export.rs @@ -23,8 +23,10 @@ async fn get_master_pass( msg: Message, db: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, master_pass: String, ) -> crate::Result<()> { + ids.delete(&bot).await; let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let mut accounts = Vec::new(); @@ -53,4 +55,9 @@ async fn get_master_pass( Ok(()) } -handler!(pub export(), "Send the master password to export your accounts", State::GetMasterPass, get_master_pass); +first_handler!( + export, + "Send the master password to export your accounts", + State::GetMasterPass, + get_master_pass +); diff --git a/src/commands/get_account.rs b/src/commands/get_account.rs index 3cfd8c8..7253913 100644 --- a/src/commands/get_account.rs +++ b/src/commands/get_account.rs @@ -8,10 +8,13 @@ async fn get_master_pass( msg: Message, db: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, name: String, master_pass: String, ) -> crate::Result<()> { + ids.delete(&bot).await; dialogue.exit().await?; + let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let account = match Account::get(user_id, &name, &db).await? { Some(account) => account, @@ -32,4 +35,8 @@ async fn get_master_pass( } handler!(get_account_name(name:String), "Send master password", State::GetMasterPass, get_master_pass); -ask_name_handler!(pub get_account(), "Send the name of the account to get", get_account_name); +ask_name_handler!( + get_account, + "Send the name of the account to get", + get_account_name +); diff --git a/src/commands/import.rs b/src/commands/import.rs index 5942d9c..923cd5b 100644 --- a/src/commands/import.rs +++ b/src/commands/import.rs @@ -31,9 +31,11 @@ async fn get_master_pass( msg: Message, db: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, user: User, master_pass: String, ) -> crate::Result<()> { + ids.delete(&bot).await; let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let mut failed = Vec::new(); @@ -64,4 +66,9 @@ async fn get_master_pass( } handler!(get_user(user: User), "Send master password", State::GetMasterPass, get_master_pass); -handler!(pub import(), "Send a json document with the same format as created by /export", State::GetUser, get_user); +first_handler!( + import, + "Send a json document with the same format as created by /export", + State::GetUser, + get_user +); diff --git a/src/commands/menu.rs b/src/commands/menu.rs index 9d44f32..3082653 100644 --- a/src/commands/menu.rs +++ b/src/commands/menu.rs @@ -15,7 +15,7 @@ pub async fn menu(bot: Throttle, msg: Message, db: DatabaseConnection) -> c .await?; } - let markup = spawn_blocking(|| menu_markup(names)).await?; + let markup = spawn_blocking(|| menu_markup_sync(names)).await?; bot.send_message(msg.chat.id, "Choose your account") .reply_markup(markup) .await?; diff --git a/src/commands/set_master_pass.rs b/src/commands/set_master_pass.rs index 0b6133e..7e4b1b5 100644 --- a/src/commands/set_master_pass.rs +++ b/src/commands/set_master_pass.rs @@ -7,9 +7,12 @@ async fn get_master_pass( msg: Message, db: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, master_pass: String, ) -> crate::Result<()> { + ids.delete(&bot).await; dialogue.exit().await?; + let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let model = spawn_blocking(move || master_pass::ActiveModel::from_unencrypted(user_id, &master_pass)) @@ -40,10 +43,10 @@ pub async fn set_master_pass( .await?; dialogue .update(State::GetNewMasterPass(Handler::new( - |bot, msg, db, dialogue, master_pass| { - Box::pin(get_master_pass(bot, msg, db, dialogue, master_pass)) + |bot, msg, db, dialogue, ids, master_pass| { + Box::pin(get_master_pass(bot, msg, db, dialogue, ids, master_pass)) }, - &previous, + MessageIds::from(&previous), ))) .await?; Ok(()) diff --git a/src/errors.rs b/src/errors.rs index d4c0f99..bae9762 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -11,7 +11,7 @@ pub enum InvalidCommand { #[error("Invalid params")] InvalidParams, #[error("Not enough bytes in the name's hash")] - OutputTooShort, + InvalidOutputLength, #[error("Error decoding the values: {0}")] - NameDecodingError(#[from] base64::DecodeSliceError), + NameDecodingError(#[from] base64::DecodeError), } diff --git a/src/macros.rs b/src/macros.rs index 9723bb8..f335bb9 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -1,39 +1,66 @@ +#[macro_export] +macro_rules! change_state { + ($dialogue: expr, $previous: expr, ($($param: ident),*), $message: literal, $next_state: expr, $next_func: ident) => {{ + $dialogue + .update($next_state(Handler::new( + move |bot, msg, db, dialogue, ids, param| Box::pin($next_func(bot, msg, db, dialogue, ids, $($param,)* param)), + &$previous, + ))) + .await?; + Ok(()) + }}; +} + +#[macro_export] +macro_rules! first_handler { + ($function_name: ident, $message: expr, $next_state: expr, $next_func: ident) => { + #[inline] + pub async fn $function_name( + bot: Throttle, + msg: Message, + dialogue: MainDialogue, + ) -> $crate::Result<()> { + let previous = bot.send_message(msg.chat.id, $message).await?; + + $crate::change_state!(dialogue, previous, (), $message, $next_state, $next_func) + } + }; +} + #[macro_export] macro_rules! handler { - ($v: vis $function_name: ident ($($param: ident: $type: ty),*), $message: literal, $next_state: expr, $next_func: ident) => { + ($function_name: ident ($($param: ident: $type: ty),*), $message: literal, $next_state: expr, $next_func: ident) => { + #[allow(clippy::too_many_arguments)] #[inline] - $v async fn $function_name( + async fn $function_name( bot: Throttle, msg: Message, _: DatabaseConnection, dialogue: MainDialogue, + ids: MessageIds, $($param: $type),* ) -> $crate::Result<()> { + ids.delete(&bot).await; + let previous = bot.send_message(msg.chat.id, $message).await?; - dialogue - .update($next_state(Handler::new( - move |bot, msg, db, dialogue, param| Box::pin($next_func(bot, msg, db, dialogue, $($param,)* param)), - &previous, - ))) - .await?; - Ok(()) + + $crate::change_state!(dialogue, previous, ($($param),*), $message, $next_state, $next_func) } }; } #[macro_export] macro_rules! ask_name_handler { - ($v: vis $function_name: ident ($($param: ident: $type: ty),*), $message: literal, $next_func: ident) => { + ($function_name: ident, $message: literal, $next_func: ident) => { #[inline] - $v async fn $function_name( + pub async fn $function_name( bot: Throttle, msg: Message, dialogue: MainDialogue, db: DatabaseConnection, - $($param: $type),* ) -> $crate::Result<()> { let user_id = msg.from().ok_or(NoUserInfo)?.id.0; - let markup = account_markup(user_id, &db).await?; + let markup = account_list_markup(user_id, &db).await?; if markup.keyboard.is_empty() { bot.send_message(msg.chat.id, "No accounts found") .reply_markup(deletion_markup()) @@ -44,13 +71,15 @@ macro_rules! ask_name_handler { .send_message(msg.chat.id, $message) .reply_markup(markup) .await?; - dialogue - .update(State::GetExistingName(Handler::new( - move |bot, msg, db, dialogue, param| Box::pin($next_func(bot, msg, db, dialogue, $($param,)* param)), - &previous, - ))) - .await?; - Ok(()) + + $crate::change_state!( + dialogue, + previous, + (), + $message, + State::GetExistingName, + $next_func + ) } }; } @@ -71,7 +100,7 @@ macro_rules! simple_state_handler { msg, db, dialogue, - |bot, msg, db, param| Box::pin($check(bot, msg, db, param)), + |msg, db, param| Box::pin($check(msg, db, param)), $no_text_message, next, ) diff --git a/src/main.rs b/src/main.rs index 35ebe69..42bc2e3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -63,7 +63,8 @@ fn get_dispatcher( let callback_handler = Update::filter_callback_query() .filter_map(CallbackCommand::from_query) .branch(case![CallbackCommand::GetMenu].endpoint(callbacks::get_menu)) - .branch(case![CallbackCommand::DeleteMessage].endpoint(callbacks::delete_message)); + .branch(case![CallbackCommand::DeleteMessage].endpoint(callbacks::delete_message)) + .branch(case![CallbackCommand::Get(hash)].endpoint(callbacks::get)); let handler = dptree::entry() .branch(message_handler) diff --git a/src/markups.rs b/src/markups.rs index e5183dc..6f5c730 100644 --- a/src/markups.rs +++ b/src/markups.rs @@ -1,13 +1,13 @@ use crate::prelude::*; -use base64::{engine::general_purpose::STANDARD as B64_ENGINE, Engine as _}; -use futures::future; +use base64::{engine::general_purpose::STANDARD_NO_PAD as B64_ENGINE, Engine as _}; use itertools::Itertools; use sha2::{Digest, Sha256}; use teloxide::types::{InlineKeyboardButton, InlineKeyboardMarkup, KeyboardButton, KeyboardMarkup}; +use tokio::task::spawn_blocking; /// Creates a markup of all user's account names #[inline] -pub async fn account_markup( +pub async fn account_list_markup( user_id: u64, db: &DatabaseConnection, ) -> crate::Result { @@ -25,7 +25,7 @@ pub async fn account_markup( } #[inline] -pub fn menu_markup(names: impl IntoIterator) -> InlineKeyboardMarkup { +pub fn menu_markup_sync(names: impl IntoIterator) -> InlineKeyboardMarkup { let names = names .into_iter() .map(|name| { @@ -40,6 +40,57 @@ pub fn menu_markup(names: impl IntoIterator) -> InlineKeyboardMar InlineKeyboardMarkup::new(&names) } +#[inline] +pub async fn menu_markup( + user_id: u64, + db: &DatabaseConnection, +) -> crate::Result { + let names: Vec = Account::get_names(user_id, db).await?.try_collect().await?; + + spawn_blocking(|| menu_markup_sync(names)) + .await + .map_err(Into::into) +} + +#[inline] +fn make_button(text: &str, command: &str, hash: &str) -> InlineKeyboardButton { + let mut data = command.to_owned(); + data.reserve(44); + data.push(' '); + data.push_str(hash); + InlineKeyboardButton::callback(text, data) +} + +#[inline] +pub fn account_markup(name: &str, is_encrypted: bool) -> InlineKeyboardMarkup { + let mut hash = [0; 43]; + B64_ENGINE + .encode_slice(::digest(name), &mut hash) + .unwrap(); + let hash = std::str::from_utf8(&hash).unwrap(); + + let alter_buttons = [ + ("Alter name", "an"), + ("Alter login", "al"), + ("Alter password", "ap"), + ] + .map(|(text, command)| make_button(text, command, hash)); + + let mut second_raw = Vec::new(); + if is_encrypted { + second_raw.push(make_button("Decrypt", "decrypt", hash)) + } else { + second_raw.push(make_button("Hide", "hide", hash)); + } + second_raw.push(make_button("Delete account", "delete", hash)); + + let menu_button = InlineKeyboardButton::callback("Back to the menu", "get_menu"); + + InlineKeyboardMarkup::new([alter_buttons]) + .append_row(second_raw) + .append_row([menu_button]) +} + /// Creates a markup with a "Delete message" button. /// This markup should be added for all messages that won't be deleted afterwards #[inline] diff --git a/src/prelude.rs b/src/prelude.rs index 04ff7f6..626f825 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -2,7 +2,7 @@ pub(crate) use crate::{ ask_name_handler, commands::Command, errors::*, - handler, + first_handler, handler, markups::*, models::*, state::State, diff --git a/src/state/generic.rs b/src/state/generic.rs index 6f382f9..d0569dc 100644 --- a/src/state/generic.rs +++ b/src/state/generic.rs @@ -14,11 +14,10 @@ pub async fn generic( ) -> crate::Result<()> where for<'a> F: FnOnce( - &'a Throttle, &'a Message, &'a DatabaseConnection, &'a str, - ) -> BoxFuture<'a, crate::Result>>, + ) -> BoxFuture<'a, crate::Result>>, { let mut handler = next.lock().await; if handler.func.is_none() { @@ -26,8 +25,6 @@ where return Err(HandlerUsed.into()); } - handler.previous.delete(&bot).await; - let text = match msg.text() { Some(text) => text.trim(), None => { @@ -45,16 +42,18 @@ where return Ok(()); } - if let Some(failure_message) = check(&bot, &msg, &db, text).await? { + if let Some(text) = check(&msg, &db, text).await? { + let failure_message = bot.send_message(msg.chat.id, text).await?; handler.previous = MessageIds::from(&failure_message); return Ok(()); } + let previous = handler.previous; let func = handler.func.take().unwrap(); drop(handler); let text = text.to_owned(); - if let Err(err) = func(bot, msg, db, dialogue.clone(), text).await { + if let Err(err) = func(bot, msg, db, dialogue.clone(), previous, text).await { let _ = dialogue.exit().await; return Err(err); } diff --git a/src/state/get_existing_name.rs b/src/state/get_existing_name.rs index 9ff342c..2c8293b 100644 --- a/src/state/get_existing_name.rs +++ b/src/state/get_existing_name.rs @@ -1,24 +1,5 @@ use crate::prelude::*; -/// Checks that the account with that name exists -#[inline] -async fn check_name( - bot: &Throttle, - msg: &Message, - db: &DatabaseConnection, - name: &str, - user_id: u64, -) -> crate::Result> { - if !Account::exists(user_id, name, db).await? { - let msg = bot - .send_message(msg.chat.id, "Account doesn't exists. Try again") - .reply_markup(account_markup(user_id, db).await?) - .await?; - return Ok(Some(msg)); - } - Ok(None) -} - /// Function to handle GetExistingName state pub async fn get_existing_name( bot: Throttle, @@ -35,8 +16,6 @@ pub async fn get_existing_name( return Err(HandlerUsed.into()); } - handler.previous.delete(&bot).await; - let text = match msg.text() { Some(text) => text.trim(), None => { @@ -45,7 +24,7 @@ pub async fn get_existing_name( msg.chat.id, "Couldn't get the text of the message. Send the name again", ) - .reply_markup(account_markup(user_id, &db).await?) + .reply_markup(account_list_markup(user_id, &db).await?) .await?; handler.previous = MessageIds::from(&msg); return Ok(()); @@ -60,16 +39,21 @@ pub async fn get_existing_name( return Ok(()); } - if let Some(failure_message) = check_name(&bot, &msg, &db, text, user_id).await? { - handler.previous = MessageIds::from(&failure_message); + if !Account::exists(user_id, text, &db).await? { + let msg = bot + .send_message(msg.chat.id, "Account doesn't exists. Try again") + .reply_markup(account_list_markup(user_id, &db).await?) + .await?; + handler.previous = MessageIds::from(&msg); return Ok(()); } + let previous = handler.previous; let func = handler.func.take().unwrap(); drop(handler); let text = text.to_owned(); - if let Err(err) = func(bot, msg, db, dialogue.clone(), text).await { + if let Err(err) = func(bot, msg, db, dialogue.clone(), previous, text).await { let _ = dialogue.exit().await; return Err(err); } diff --git a/src/state/get_login.rs b/src/state/get_login.rs index c76782e..39a80d0 100644 --- a/src/state/get_login.rs +++ b/src/state/get_login.rs @@ -2,17 +2,13 @@ use crate::prelude::*; #[inline] async fn check_login( - bot: &Throttle, - msg: &Message, + _: &Message, _: &DatabaseConnection, login: &str, -) -> crate::Result> { +) -> crate::Result> { let is_valid = validate_field(login); if !is_valid { - let msg = bot - .send_message(msg.chat.id, "Invalid login. Try again") - .await?; - return Ok(Some(msg)); + return Ok(Some("Invalid login. Try again".to_owned())); } Ok(None) } diff --git a/src/state/get_master_pass.rs b/src/state/get_master_pass.rs index e32e748..8b4967c 100644 --- a/src/state/get_master_pass.rs +++ b/src/state/get_master_pass.rs @@ -5,11 +5,10 @@ use tokio::task::spawn_blocking; /// Returns true if the provided master password is valid #[inline] async fn check_master_pass( - bot: &Throttle, msg: &Message, db: &DatabaseConnection, master_pass: &str, -) -> crate::Result> { +) -> crate::Result> { let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let model = MasterPass::get(user_id, db).await?; @@ -20,21 +19,15 @@ async fn check_master_pass( } None => { error!("User was put into the GetMasterPass state with no master password set"); - let msg = bot - .send_message( - msg.chat.id, - "No master password set. Use /cancel and set it by using /set_master_pass", - ) - .await?; - return Ok(Some(msg)); + return Ok(Some( + "No master password set. Use /cancel and set it by using /set_master_pass" + .to_owned(), + )); } }; if !is_valid { - let msg = bot - .send_message(msg.chat.id, "Wrong master password. Try again") - .await?; - return Ok(Some(msg)); + return Ok(Some("Wrong master password. Try again".to_owned())); } Ok(None) } diff --git a/src/state/get_new_master_pass.rs b/src/state/get_new_master_pass.rs index 9c0bbc1..539e5a7 100644 --- a/src/state/get_new_master_pass.rs +++ b/src/state/get_new_master_pass.rs @@ -33,20 +33,13 @@ fn process_validity(validity: PasswordValidity) -> Result<(), String> { /// Checks that the account with that name exists #[inline] async fn check_new_master_pass( - bot: &Throttle, - msg: &Message, + _: &Message, _: &DatabaseConnection, password: &str, -) -> crate::Result> { +) -> crate::Result> { let validity = check_master_pass(password); - match process_validity(validity) { - Ok(()) => Ok(None), - Err(error_text) => { - let msg = bot.send_message(msg.chat.id, error_text).await?; - Ok(Some(msg)) - } - } + Ok(process_validity(validity).err()) } crate::simple_state_handler!( diff --git a/src/state/get_new_name.rs b/src/state/get_new_name.rs index b40a66e..0916674 100644 --- a/src/state/get_new_name.rs +++ b/src/state/get_new_name.rs @@ -3,26 +3,19 @@ use crate::prelude::*; /// Validates a new account #[inline] async fn check_new_account_name( - bot: &Throttle, msg: &Message, db: &DatabaseConnection, name: &str, -) -> crate::Result> { +) -> crate::Result> { let user_id = msg.from().ok_or(NoUserInfo)?.id.0; - if Account::exists(user_id, name, db).await? { - let msg = bot - .send_message(msg.chat.id, "Account already exists") - .await?; - return Ok(Some(msg)); - } - if !validate_field(name) { - let msg = bot - .send_message(msg.chat.id, "Invalid account name. Try again") - .await?; - return Ok(Some(msg)); - } - Ok(None) + if Account::exists(user_id, name, db).await? { + Ok(Some("Account already exists".to_owned())) + } else if !validate_field(name) { + Ok(Some("Invalid account name. Try again".to_owned())) + } else { + Ok(None) + } } crate::simple_state_handler!( diff --git a/src/state/get_password.rs b/src/state/get_password.rs index 70c8699..a883ea3 100644 --- a/src/state/get_password.rs +++ b/src/state/get_password.rs @@ -2,17 +2,13 @@ use crate::prelude::*; #[inline] async fn check_password( - bot: &Throttle, - msg: &Message, + _: &Message, _: &DatabaseConnection, - login: &str, -) -> crate::Result> { - let is_valid = validate_field(login); + password: &str, +) -> crate::Result> { + let is_valid = validate_field(password); if !is_valid { - let msg = bot - .send_message(msg.chat.id, "Invalid password. Try again") - .await?; - return Ok(Some(msg)); + return Ok(Some("Invalid password. Try again".to_owned())); } Ok(None) } diff --git a/src/state/get_user.rs b/src/state/get_user.rs index 30de797..ae286c0 100644 --- a/src/state/get_user.rs +++ b/src/state/get_user.rs @@ -145,8 +145,6 @@ pub async fn get_user( return Err(HandlerUsed.into()); } - handler.previous.delete(&bot).await; - if let Some("/cancel") = msg.text().map(str::trim) { dialogue.exit().await?; bot.send_message(msg.chat.id, "Successfully cancelled") @@ -190,10 +188,11 @@ pub async fn get_user( } }; + let previous = handler.previous; let func = handler.func.take().unwrap(); drop(handler); - if let Err(err) = func(bot, msg, db, dialogue.clone(), user).await { + if let Err(err) = func(bot, msg, db, dialogue.clone(), previous, user).await { let _ = dialogue.exit().await; return Err(err); } diff --git a/src/state/handler.rs b/src/state/handler.rs index ab58e18..36f713f 100644 --- a/src/state/handler.rs +++ b/src/state/handler.rs @@ -1,7 +1,7 @@ use crate::prelude::*; use futures::future::BoxFuture; use std::sync::Arc; -use teloxide::types::MessageId; +use teloxide::types::{InlineKeyboardMarkup, MessageId}; use tokio::sync::Mutex; #[derive(Clone, Copy)] @@ -27,6 +27,7 @@ type DynHanlder = Box< Message, DatabaseConnection, MainDialogue, + MessageIds, T, ) -> BoxFuture<'static, crate::Result<()>> + Send, @@ -50,6 +51,7 @@ impl Handler { Message, DatabaseConnection, MainDialogue, + MessageIds, T, ) -> BoxFuture<'static, crate::Result<()>> + Send