diff --git a/src/handlers/commands/add_account.rs b/src/handlers/commands/add_account.rs index 2263357..e1d778c 100644 --- a/src/handlers/commands/add_account.rs +++ b/src/handlers/commands/add_account.rs @@ -1,7 +1,6 @@ use crate::{ - entity::prelude::Account, errors::NoUserInfo, - handlers::{markups::deletion_markup, Handler, MainDialogue, State}, + handlers::{markups::deletion_markup, state::NameCheckKind, Handler, MainDialogue, State}, models::DecryptedAccount, }; use sea_orm::prelude::*; @@ -81,18 +80,10 @@ async fn get_login( async fn get_account_name( bot: Throttle, msg: Message, - db: DatabaseConnection, + _: DatabaseConnection, dialogue: MainDialogue, name: String, ) -> crate::Result<()> { - let user_id = msg.from().ok_or(NoUserInfo)?.id.0; - if Account::exists(user_id, &name, &db).await? { - dialogue.exit().await?; - bot.send_message(msg.chat.id, "Account alreay exists") - .reply_markup(deletion_markup()) - .await?; - return Ok(()); - } let previous = bot.send_message(msg.chat.id, "Send login").await?; dialogue .update(State::GetLogin(Handler::new( @@ -111,10 +102,10 @@ pub async fn add_account( ) -> crate::Result<()> { let previous = bot.send_message(msg.chat.id, "Send account name").await?; dialogue - .update(State::GetAccountName(Handler::new( - get_account_name, - previous, - ))) + .update(State::GetAccountName( + Handler::new(get_account_name, previous), + NameCheckKind::NewAccountName, + )) .await?; Ok(()) } diff --git a/src/handlers/commands/delete.rs b/src/handlers/commands/delete.rs index f36c54f..e525461 100644 --- a/src/handlers/commands/delete.rs +++ b/src/handlers/commands/delete.rs @@ -3,6 +3,7 @@ use crate::{ errors::NoUserInfo, handlers::{ markups::{self, deletion_markup}, + state::NameCheckKind, Handler, MainDialogue, State, }, }; @@ -31,18 +32,10 @@ async fn get_master_pass( async fn get_account_name( bot: Throttle, msg: Message, - db: DatabaseConnection, + _: DatabaseConnection, dialogue: MainDialogue, name: String, ) -> crate::Result<()> { - let user_id = msg.from().ok_or(NoUserInfo)?.id.0; - if !Account::exists(user_id, &name, &db).await? { - dialogue.exit().await?; - bot.send_message(msg.chat.id, "Account doesn't exists") - .reply_markup(deletion_markup()) - .await?; - return Ok(()); - } let previous = bot .send_message(msg.chat.id, "Send master password. Once you send correct master password the account is unrecoverable") .await?; @@ -69,10 +62,10 @@ pub async fn delete( .reply_markup(markup) .await?; dialogue - .update(State::GetAccountName(Handler::new( - get_account_name, - previous, - ))) + .update(State::GetAccountName( + Handler::new(get_account_name, previous), + NameCheckKind::MustExist, + )) .await?; Ok(()) } diff --git a/src/handlers/commands/get_account.rs b/src/handlers/commands/get_account.rs index 6a457bb..e6b4d36 100644 --- a/src/handlers/commands/get_account.rs +++ b/src/handlers/commands/get_account.rs @@ -3,6 +3,7 @@ use crate::{ errors::NoUserInfo, handlers::{ markups::{self, deletion_markup}, + state::NameCheckKind, Handler, MainDialogue, State, }, }; @@ -43,18 +44,10 @@ async fn get_master_pass( async fn get_account_name( bot: Throttle, msg: Message, - db: DatabaseConnection, + _: DatabaseConnection, dialogue: MainDialogue, name: String, ) -> crate::Result<()> { - let user_id = msg.from().ok_or(NoUserInfo)?.id.0; - if !Account::exists(user_id, &name, &db).await? { - dialogue.exit().await?; - bot.send_message(msg.chat.id, "Account doesn't exists") - .reply_markup(deletion_markup()) - .await?; - return Ok(()); - } let previous = bot .send_message(msg.chat.id, "Send master password") .await?; @@ -83,10 +76,10 @@ pub async fn get_account( .reply_markup(markup) .await?; dialogue - .update(State::GetAccountName(Handler::new( - get_account_name, - previous, - ))) + .update(State::GetAccountName( + Handler::new(get_account_name, previous), + NameCheckKind::MustExist, + )) .await?; Ok(()) } diff --git a/src/handlers/commands/import.rs b/src/handlers/commands/import.rs index 0188de6..a3e698e 100644 --- a/src/handlers/commands/import.rs +++ b/src/handlers/commands/import.rs @@ -1,14 +1,13 @@ use crate::{ errors::NoUserInfo, handlers::{markups::deletion_markup, Handler, MainDialogue, State}, - models::{DecryptedAccount, User}, + models::User, }; -use futures::{future, stream::FuturesUnordered, StreamExt, TryStreamExt}; +use futures::{future, stream::FuturesUnordered, StreamExt}; use itertools::Itertools; use sea_orm::prelude::*; -use serde_json::from_slice; use std::sync::Arc; -use teloxide::{adaptors::Throttle, net::Download, prelude::*}; +use teloxide::{adaptors::Throttle, prelude::*}; use tokio::task::spawn_blocking; /// Gets the master password, encryptes and adds the accounts to the DB @@ -18,12 +17,15 @@ async fn get_master_pass( db: DatabaseConnection, dialogue: MainDialogue, master_pass: String, - accounts: Vec, + user: User, ) -> crate::Result<()> { let user_id = msg.from().ok_or(NoUserInfo)?.id.0; + let accounts = user.accounts; + let failed: Vec = { let master_pass: Arc = master_pass.into(); let db = &db; + let futures: FuturesUnordered<_> = accounts .into_iter() .map(|account| { @@ -44,6 +46,7 @@ async fn get_master_pass( } }) .collect(); + futures .filter_map(|result| future::ready(result.err())) .collect() @@ -70,44 +73,15 @@ async fn get_document( msg: Message, _: DatabaseConnection, dialogue: MainDialogue, - (): (), + user: User, ) -> crate::Result<()> { - let document = msg.document().unwrap(); - match document.file_name { - Some(ref name) if name.trim_end().ends_with(".json") => (), - _ => { - bot.send_message(msg.chat.id, "Invalid file name") - .reply_markup(deletion_markup()) - .await?; - dialogue.exit().await?; - return Ok(()); - } - } - let mut data = Vec::with_capacity(document.file.size as usize); - let file = bot.get_file(&document.file.id).await?; - bot.download_file_stream(&file.path) - .try_for_each(|bytes| { - data.extend(bytes); - async { Ok(()) } - }) - .await?; - let accounts = match spawn_blocking(move || from_slice::(&data)).await? { - Ok(user) => user.accounts, - Err(_) => { - bot.send_message(msg.chat.id, "Error parsing the json file") - .reply_markup(deletion_markup()) - .await?; - dialogue.exit().await?; - return Ok(()); - } - }; let previous = bot .send_message(msg.chat.id, "Send master password") .await?; dialogue .update(State::GetMasterPass(Handler::new( move |bot, msg, db, dialogue, master_pass| { - get_master_pass(bot, msg, db, dialogue, master_pass, accounts) + get_master_pass(bot, msg, db, dialogue, master_pass, user) }, previous, ))) @@ -124,7 +98,7 @@ pub async fn import(bot: Throttle, msg: Message, dialogue: MainDialogue) -> ) .await?; dialogue - .update(State::GetDocument(Handler::new(get_document, previous))) + .update(State::GetUser(Handler::new(get_document, previous))) .await?; Ok(()) } diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index 69b911d..81e161d 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -8,7 +8,7 @@ mod utils; use commands::Command; use sea_orm::prelude::*; -use state::{Handler, PackagedHandler}; +use state::{Handler, MainDialogue, State}; use teloxide::{ adaptors::{throttle::Limits, Throttle}, dispatching::dialogue::InMemStorage, @@ -16,19 +16,6 @@ use teloxide::{ prelude::*, }; -type MainDialogue = Dialogue>; - -#[derive(Default, Clone)] -pub enum State { - #[default] - Start, - GetAccountName(PackagedHandler), - GetMasterPass(PackagedHandler), - GetLogin(PackagedHandler), - GetPassword(PackagedHandler), - GetDocument(PackagedHandler<()>), -} - pub fn get_dispatcher( token: String, db: DatabaseConnection, @@ -56,11 +43,11 @@ pub fn get_dispatcher( let message_handler = Update::filter_message() .map_async(utils::delete_message) .enter_dialogue::, State>() - .branch(case![State::GetAccountName(next)].endpoint(state::get_account_name)) + .branch(case![State::GetAccountName(next, check)].endpoint(state::get_account_name)) .branch(case![State::GetMasterPass(next)].endpoint(state::get_master_pass)) .branch(case![State::GetLogin(next)].endpoint(state::get_login)) .branch(case![State::GetPassword(next)].endpoint(state::get_password)) - .branch(case![State::GetDocument(next)].endpoint(state::get_document)) + .branch(case![State::GetUser(next)].endpoint(state::get_user)) .branch(command_handler) .endpoint(default::default); diff --git a/src/handlers/state/generic.rs b/src/handlers/state/generic.rs index a60c38b..a8d3277 100644 --- a/src/handlers/state/generic.rs +++ b/src/handlers/state/generic.rs @@ -22,19 +22,23 @@ where &'a Message, &'a DatabaseConnection, &'a str, - ) -> PinnedFuture<'a, crate::Result>, + ) -> PinnedFuture<'a, crate::Result>>, { - let handler = next.lock().await.take(); - let previous = handler.as_ref().and_then(|h| h.previous.as_ref()); - delete_optional(&bot, previous).await; + let mut handler = next.lock().await; + delete_optional(&bot, handler.previous.as_ref()).await; + let text = match msg.text() { Some(text) => text.trim_end(), None => { - bot.send_message(msg.chat.id, "Couldn't get the text of the message") - .await?; + bot.send_message( + msg.chat.id, + "Couldn't get the text of the message. Send the message again", + ) + .await?; return Ok(()); } }; + if text == "/cancel" { dialogue.exit().await?; bot.send_message(msg.chat.id, "Successfully cancelled") @@ -42,26 +46,22 @@ where .await?; return Ok(()); } - let handler = match handler { - Some(handler) => handler, - None => { - let _ = dialogue.exit().await; - return Err(HandlerUsed.into()); - } - }; - match check(&bot, &msg, &db, text).await { - Ok(true) => (), - Ok(false) => { - dialogue.exit().await?; - return Ok(()); - } - Err(err) => { - let _ = dialogue.exit().await; - return Err(err); - } - }; + + if handler.func.is_none() { + let _ = dialogue.exit().await; + return Err(HandlerUsed.into()); + } + + if let Some(failure_message) = check(&bot, &msg, &db, text).await? { + handler.previous = Some(failure_message); + return Ok(()); + } + + let func = handler.func.take().unwrap(); + drop(handler); let text = text.to_owned(); - if let Err(err) = (handler.handler)(bot, msg, db, dialogue.clone(), text).await { + + if let Err(err) = func(bot, msg, db, dialogue.clone(), text).await { let _ = dialogue.exit().await; return Err(err); } diff --git a/src/handlers/state/get_account_name.rs b/src/handlers/state/get_account_name.rs index ca35245..3735802 100644 --- a/src/handlers/state/get_account_name.rs +++ b/src/handlers/state/get_account_name.rs @@ -1,31 +1,71 @@ -use crate::handlers::{markups::deletion_markup, utils::validate_field, MainDialogue}; +use crate::{ + entity::prelude::Account, + errors::NoUserInfo, + handlers::{markups::account_markup, utils::validate_field, MainDialogue}, +}; use sea_orm::prelude::*; use teloxide::{adaptors::Throttle, prelude::*}; +/// Specifies the kind of checks to be run during the account name validation +#[derive(Clone, Copy)] +pub enum NameCheckKind { + NewAccountName, + MustExist, +} + +/// Validates the account name +#[inline] +pub async fn check_name( + bot: &Throttle, + msg: &Message, + db: &DatabaseConnection, + name: &str, + check_kind: NameCheckKind, +) -> crate::Result> { + let user_id = msg.from().ok_or(NoUserInfo)?.id.0; + match check_kind { + NameCheckKind::MustExist => { + 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) + } + NameCheckKind::NewAccountName => { + if Account::exists(user_id, name, db).await? { + let msg = bot + .send_message(msg.chat.id, "Account alreay 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) + } + } +} + /// Function to handle GetAccountName state pub async fn get_account_name( bot: Throttle, msg: Message, db: DatabaseConnection, dialogue: MainDialogue, - next: super::PackagedHandler, + (next, check_kind): (super::PackagedHandler, NameCheckKind), ) -> crate::Result<()> { super::generic::generic( bot, msg, db, dialogue, - |bot, msg, _, name| { - Box::pin(async move { - let is_valid = validate_field(name); - if !is_valid { - bot.send_message(msg.chat.id, "Invalid account name") - .reply_markup(deletion_markup()) - .await?; - } - Ok(is_valid) - }) - }, + |bot, msg, db, name| Box::pin(check_name(bot, msg, db, name, check_kind)), next, ) .await diff --git a/src/handlers/state/get_document.rs b/src/handlers/state/get_document.rs deleted file mode 100644 index 494329f..0000000 --- a/src/handlers/state/get_document.rs +++ /dev/null @@ -1,45 +0,0 @@ -use crate::{ - errors::HandlerUsed, - handlers::{markups::deletion_markup, utils::delete_optional, MainDialogue}, -}; -use sea_orm::prelude::*; -use teloxide::{adaptors::Throttle, prelude::*}; - -/// Function to handle get_document state. It doesn't actually validate anything -pub async fn get_document( - bot: Throttle, - msg: Message, - db: DatabaseConnection, - dialogue: MainDialogue, - next: super::PackagedHandler<()>, -) -> crate::Result<()> { - let handler = next.lock().await.take(); - let previous = handler.as_ref().and_then(|h| h.previous.as_ref()); - delete_optional(&bot, previous).await; - if let Some("/cancel") = msg.text().map(str::trim_end) { - dialogue.exit().await?; - bot.send_message(msg.chat.id, "Successfully cancelled") - .reply_markup(deletion_markup()) - .await?; - return Ok(()); - } - let handler = match handler { - Some(handler) => handler, - None => { - let _ = dialogue.exit().await; - return Err(HandlerUsed.into()); - } - }; - if msg.document().is_none() { - bot.send_message(msg.chat.id, "You didn't send a file") - .reply_markup(deletion_markup()) - .await?; - dialogue.exit().await?; - return Ok(()); - } - if let Err(err) = (handler.handler)(bot, msg, db, dialogue.clone(), ()).await { - let _ = dialogue.exit().await; - return Err(err); - } - Ok(()) -} diff --git a/src/handlers/state/get_login.rs b/src/handlers/state/get_login.rs index 0605977..0294e57 100644 --- a/src/handlers/state/get_login.rs +++ b/src/handlers/state/get_login.rs @@ -1,4 +1,4 @@ -use crate::handlers::{markups::deletion_markup, utils::validate_field, MainDialogue}; +use crate::handlers::{utils::validate_field, MainDialogue}; use sea_orm::prelude::*; use teloxide::{adaptors::Throttle, prelude::*}; @@ -19,11 +19,12 @@ pub async fn get_login( Box::pin(async move { let is_valid = validate_field(login); if !is_valid { - bot.send_message(msg.chat.id, "Invalid login") - .reply_markup(deletion_markup()) + let msg = bot + .send_message(msg.chat.id, "Invalid login. Try again") .await?; + return Ok(Some(msg)); } - Ok(is_valid) + Ok(None) }) }, next, diff --git a/src/handlers/state/get_master_pass.rs b/src/handlers/state/get_master_pass.rs index 6891e13..94795cd 100644 --- a/src/handlers/state/get_master_pass.rs +++ b/src/handlers/state/get_master_pass.rs @@ -1,40 +1,44 @@ -use crate::{ - entity::prelude::MasterPass, - errors::NoUserInfo, - handlers::{markups::deletion_markup, MainDialogue}, -}; +use crate::{entity::prelude::MasterPass, errors::NoUserInfo, handlers::MainDialogue}; +use log::error; use sea_orm::prelude::*; use teloxide::{adaptors::Throttle, prelude::*}; use tokio::task::spawn_blocking; /// Returns true if the provided master password is valid #[inline] -pub async fn check_master_pass<'a>( - bot: &'a Throttle, - msg: &'a Message, - db: &'a DatabaseConnection, - master_pass: &'a str, -) -> crate::Result { +pub async fn check_master_pass( + bot: &Throttle, + msg: &Message, + db: &DatabaseConnection, + master_pass: &str, +) -> crate::Result> { let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let model = MasterPass::get(user_id, db).await?; + let is_valid = match model { Some(model) => { let master_pass = master_pass.to_owned(); spawn_blocking(move || model.verify(&master_pass)).await?? } None => { - bot.send_message(msg.chat.id, "No master password set") - .reply_markup(deletion_markup()) + 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(false); + return Ok(Some(msg)); } }; + if !is_valid { - bot.send_message(msg.chat.id, "Wrong master password") - .reply_markup(deletion_markup()) + let msg = bot + .send_message(msg.chat.id, "Wrong master password. Try again") .await?; + return Ok(Some(msg)); } - Ok(is_valid) + Ok(None) } pub async fn get_master_pass( diff --git a/src/handlers/state/get_password.rs b/src/handlers/state/get_password.rs index 98f67f3..ece1e56 100644 --- a/src/handlers/state/get_password.rs +++ b/src/handlers/state/get_password.rs @@ -1,7 +1,7 @@ use sea_orm::prelude::*; use teloxide::{adaptors::Throttle, prelude::*}; -use crate::handlers::{markups::deletion_markup, utils::validate_field, MainDialogue}; +use crate::handlers::{utils::validate_field, MainDialogue}; /// Function to handle GetPassword state pub async fn get_password( @@ -20,11 +20,12 @@ pub async fn get_password( Box::pin(async move { let is_valid = validate_field(password); if !is_valid { - bot.send_message(msg.chat.id, "Invalid password") - .reply_markup(deletion_markup()) + let msg = bot + .send_message(msg.chat.id, "Invalid password. Try again") .await?; + return Ok(Some(msg)); } - Ok(is_valid) + Ok(None) }) }, next, diff --git a/src/handlers/state/get_user.rs b/src/handlers/state/get_user.rs new file mode 100644 index 0000000..bb80888 --- /dev/null +++ b/src/handlers/state/get_user.rs @@ -0,0 +1,88 @@ +use crate::{ + errors::HandlerUsed, + handlers::{markups::deletion_markup, utils::delete_optional, MainDialogue}, + models::User, +}; +use futures::TryStreamExt; +use sea_orm::prelude::*; +use teloxide::{adaptors::Throttle, net::Download, prelude::*}; +use tokio::task::spawn_blocking; + +/// Function to handle GetUser state. It doesn't actually validate anything +pub async fn get_user( + bot: Throttle, + msg: Message, + db: DatabaseConnection, + dialogue: MainDialogue, + next: super::PackagedHandler, +) -> crate::Result<()> { + let mut handler = next.lock().await; + delete_optional(&bot, handler.previous.as_ref()).await; + + if let Some("/cancel") = msg.text().map(str::trim_end) { + dialogue.exit().await?; + bot.send_message(msg.chat.id, "Successfully cancelled") + .reply_markup(deletion_markup()) + .await?; + return Ok(()); + } + + if handler.func.is_none() { + let _ = dialogue.exit().await; + return Err(HandlerUsed.into()); + } + + let document = match msg.document() { + Some(document) => document, + None => { + let msg = bot + .send_message(msg.chat.id, "You didn't send a file. Try again") + .await?; + handler.previous = Some(msg); + return Ok(()); + } + }; + + match document.file_name.as_deref() { + Some(name) if name.trim_end().ends_with(".json") => (), + _ => { + let msg = bot + .send_message( + msg.chat.id, + "Invalid file name. You need to send a json file. Try again", + ) + .await?; + handler.previous = Some(msg); + return Ok(()); + } + } + + let file = bot.get_file(&document.file.id).await?; + let mut data = Vec::with_capacity(document.file.size as usize); + bot.download_file_stream(&file.path) + .try_for_each(|bytes| { + data.extend(bytes); + async { Ok(()) } + }) + .await?; + + let user: User = match spawn_blocking(move || serde_json::from_slice(&data)).await? { + Ok(user) => user, + Err(_) => { + let msg = bot + .send_message(msg.chat.id, "Error parsing the json file. Try again") + .await?; + handler.previous = Some(msg); + return Ok(()); + } + }; + + let func = handler.func.take().unwrap(); + drop(handler); + + if let Err(err) = func(bot, msg, db, dialogue.clone(), user).await { + let _ = dialogue.exit().await; + return Err(err); + } + Ok(()) +} diff --git a/src/handlers/state/handler.rs b/src/handlers/state/handler.rs index 6e61dbc..e3c5c5b 100644 --- a/src/handlers/state/handler.rs +++ b/src/handlers/state/handler.rs @@ -4,24 +4,25 @@ use std::{future::Future, sync::Arc}; use teloxide::{adaptors::Throttle, prelude::*}; use tokio::sync::Mutex; -#[allow(clippy::type_complexity)] +type DynHanlder = Box< + dyn FnOnce( + Throttle, + Message, + DatabaseConnection, + MainDialogue, + T, + ) -> crate::PinnedFuture<'static, crate::Result<()>> + + Send + + Sync, +>; + pub struct Handler { - pub handler: Box< - dyn FnOnce( - Throttle, - Message, - DatabaseConnection, - MainDialogue, - T, - ) -> crate::PinnedFuture<'static, crate::Result<()>> - + Send - + Sync, - >, + pub func: Option>, pub previous: Option, } -pub type PackagedHandler = Arc>>>; +pub type PackagedHandler = Arc>>; impl Handler { /// Convinience method to convert a simple async function and a previous message into PackagedHandler @@ -35,11 +36,11 @@ impl Handler { F: Future> + Send + 'static, { let handler = Self { - handler: Box::new(move |bot, msg, db, dialogue, val| { + func: Some(Box::new(move |bot, msg, db, dialogue, val| { Box::pin(f(bot, msg, db, dialogue, val)) - }), + })), previous: previous.into(), }; - Arc::new(Mutex::new(Some(handler))) + Arc::new(Mutex::new(handler)) } } diff --git a/src/handlers/state/mod.rs b/src/handlers/state/mod.rs index 81012e3..da7c24b 100644 --- a/src/handlers/state/mod.rs +++ b/src/handlers/state/mod.rs @@ -2,15 +2,31 @@ mod generic; mod get_account_name; -mod get_document; mod get_login; mod get_master_pass; mod get_password; +mod get_user; mod handler; -pub use get_account_name::get_account_name; -pub use get_document::get_document; +pub use get_account_name::{get_account_name, NameCheckKind}; pub use get_login::get_login; pub use get_master_pass::get_master_pass; pub use get_password::get_password; +pub use get_user::get_user; pub use handler::{Handler, PackagedHandler}; + +use crate::models::User; +use teloxide::{dispatching::dialogue::InMemStorage, prelude::*}; + +#[derive(Default, Clone)] +pub enum State { + #[default] + Start, + GetAccountName(PackagedHandler, NameCheckKind), + GetMasterPass(PackagedHandler), + GetLogin(PackagedHandler), + GetPassword(PackagedHandler), + GetUser(PackagedHandler), +} + +pub type MainDialogue = Dialogue>;