diff --git a/src/callbacks/alter.rs b/src/callbacks/alter.rs index ab865d6..f0034b9 100644 --- a/src/callbacks/alter.rs +++ b/src/callbacks/alter.rs @@ -51,7 +51,7 @@ async fn get_master_pass( master_pass: String, ) -> crate::Result<()> { dialogue.exit().await?; - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let text = match update_account(user_id, &db, name, field, field_value, master_pass).await { Ok(true) => locale.success.as_str(), diff --git a/src/callbacks/decrypt.rs b/src/callbacks/decrypt.rs index 3f93436..33c2ea7 100644 --- a/src/callbacks/decrypt.rs +++ b/src/callbacks/decrypt.rs @@ -15,7 +15,7 @@ async fn get_master_pass( ) -> crate::Result<()> { dialogue.exit().await?; - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let Some(account) = Account::get(user_id, &name, &db).await? else { bot.send_message(msg.chat.id, &locale.no_accounts_found) diff --git a/src/callbacks/delete.rs b/src/callbacks/delete.rs index ec44c6b..2420206 100644 --- a/src/callbacks/delete.rs +++ b/src/callbacks/delete.rs @@ -13,7 +13,7 @@ async fn get_master_pass( ) -> crate::Result<()> { dialogue.exit().await?; - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; Account::delete(user_id, &name, &db).await?; ids.alter_message(&bot, &locale.success, deletion_markup(locale), None) diff --git a/src/commands/add_account.rs b/src/commands/add_account.rs index f07cfcf..cbfb4f5 100644 --- a/src/commands/add_account.rs +++ b/src/commands/add_account.rs @@ -17,7 +17,7 @@ async fn get_master_pass( ) -> crate::Result<()> { dialogue.exit().await?; - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let account = spawn_blocking(move || { DecryptedAccount { diff --git a/src/commands/delete.rs b/src/commands/delete.rs index adc7394..137e1d3 100644 --- a/src/commands/delete.rs +++ b/src/commands/delete.rs @@ -6,7 +6,7 @@ pub async fn delete( db: Pool, locale: LocaleRef, ) -> crate::Result<()> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let markup = menu_markup("delete1", user_id, &db).await?; diff --git a/src/commands/delete_all.rs b/src/commands/delete_all.rs index c3f282c..3087dd3 100644 --- a/src/commands/delete_all.rs +++ b/src/commands/delete_all.rs @@ -12,7 +12,7 @@ async fn get_master_pass( _: String, ) -> crate::Result<()> { dialogue.exit().await?; - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let mut txn = db.begin().await?; let result = ( diff --git a/src/commands/export.rs b/src/commands/export.rs index a9dcf81..975144c 100644 --- a/src/commands/export.rs +++ b/src/commands/export.rs @@ -29,7 +29,7 @@ async fn get_master_pass( dialogue.exit().await?; ids.delete(&bot).await?; - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let mut accounts = Vec::new(); { diff --git a/src/commands/get_account.rs b/src/commands/get_account.rs index 9b0efc3..2b30ea8 100644 --- a/src/commands/get_account.rs +++ b/src/commands/get_account.rs @@ -6,7 +6,7 @@ pub async fn get_account( db: Pool, locale: LocaleRef, ) -> crate::Result<()> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let markup = menu_markup("decrypt", user_id, &db).await?; diff --git a/src/commands/get_accounts.rs b/src/commands/get_accounts.rs index 5787c1d..8acfe92 100644 --- a/src/commands/get_accounts.rs +++ b/src/commands/get_accounts.rs @@ -9,7 +9,7 @@ pub async fn get_accounts( db: Pool, locale: LocaleRef, ) -> crate::Result<()> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let mut account_names = Account::get_names(user_id, &db); let Some(mut text) = account_names.try_next().await? else { diff --git a/src/commands/import.rs b/src/commands/import.rs index ef42e27..77708a2 100644 --- a/src/commands/import.rs +++ b/src/commands/import.rs @@ -38,7 +38,7 @@ async fn get_master_pass( ) -> crate::Result<()> { dialogue.exit().await?; - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let mut failed = Vec::new(); { diff --git a/src/commands/menu.rs b/src/commands/menu.rs index 6e6d868..bc300b0 100644 --- a/src/commands/menu.rs +++ b/src/commands/menu.rs @@ -6,7 +6,7 @@ pub async fn menu( db: Pool, locale: LocaleRef, ) -> crate::Result<()> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let markup = menu_markup("get", user_id, &db).await?; diff --git a/src/commands/set_master_pass.rs b/src/commands/set_master_pass.rs index 6557a32..318c239 100644 --- a/src/commands/set_master_pass.rs +++ b/src/commands/set_master_pass.rs @@ -16,8 +16,8 @@ async fn get_master_pass2( master_pass: String, ) -> crate::Result<()> { dialogue.exit().await?; - let from = msg.from.as_ref().ok_or(NoUserInfo)?; - let user_id = from.id.0; + let from = msg.get_user()?; + let user_id = msg.get_user_id()?; if !hash.verify(master_pass.as_bytes()) { ids.alter_message( @@ -82,7 +82,7 @@ pub async fn set_master_pass( db: Pool, locale: LocaleRef, ) -> crate::Result<()> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; if MasterPass::exists(user_id, &db).await? { bot.send_message(msg.chat.id, &locale.master_password_is_set) .reply_markup(deletion_markup(locale)) diff --git a/src/dispatching/master_password_check.rs b/src/dispatching/master_password_check.rs index cba7ea5..3d29776 100644 --- a/src/dispatching/master_password_check.rs +++ b/src/dispatching/master_password_check.rs @@ -1,4 +1,4 @@ -use crate::prelude::*; +use crate::{errors::NoUserInfo, prelude::*}; use std::sync::Arc; use teloxide::{ dispatching::{dialogue::GetChatId, DpHandlerDescription}, diff --git a/src/main.rs b/src/main.rs index 4052317..b7f9aad 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ mod markups; mod models; mod prelude; mod state; +mod utils; use anyhow::{Error, Result}; diff --git a/src/prelude.rs b/src/prelude.rs index f6e9663..98a4e85 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -1,11 +1,12 @@ pub use crate::{ cryptography::{account::Decrypted as DecryptedAccount, validate_field}, entity::{account::Account, master_pass::MasterPass, Pool}, - errors::{handle_error, NoUserInfo}, + errors::handle_error, first_handler, handler, locales::LocaleRef, markups::{deletion_markup, menu_markup}, state::{CallbackMessageIdExt as _, Handler, MainDialogue, MessageIds, PackagedHandler, State}, + utils::GetUserExt as _, }; pub use futures::{StreamExt as _, TryStreamExt as _}; pub use teloxide::{adaptors::Throttle, prelude::*}; diff --git a/src/state/get_master_pass.rs b/src/state/get_master_pass.rs index 8401a82..5bc01dc 100644 --- a/src/state/get_master_pass.rs +++ b/src/state/get_master_pass.rs @@ -8,7 +8,7 @@ async fn check_master_pass( locale: LocaleRef, master_pass: &str, ) -> crate::Result> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let Some(model) = MasterPass::get(user_id, db).await? else { handle_error(anyhow::anyhow!( "User was put into the GetMasterPass state with no master password set" diff --git a/src/state/get_new_name.rs b/src/state/get_new_name.rs index 8768b40..8f6b317 100644 --- a/src/state/get_new_name.rs +++ b/src/state/get_new_name.rs @@ -7,7 +7,7 @@ async fn check_new_account_name( locale: LocaleRef, name: &str, ) -> crate::Result> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; if Account::exists(user_id, name, db).await? { Ok(Some(locale.account_already_exists.to_owned())) diff --git a/src/state/get_user.rs b/src/state/get_user.rs index 0009d1b..26031e1 100644 --- a/src/state/get_user.rs +++ b/src/state/get_user.rs @@ -181,7 +181,7 @@ pub async fn get_user( locale: LocaleRef, handler: PackagedHandler, ) -> crate::Result<()> { - let user_id = msg.from.as_ref().ok_or(NoUserInfo)?.id.0; + let user_id = msg.get_user_id()?; let mut handler = handler.lock().await; if handler.func.is_none() { diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..b6b98d9 --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,17 @@ +use teloxide::types::User; + +use crate::{errors::NoUserInfo, prelude::*}; + +pub trait GetUserExt { + fn get_user(&self) -> Result<&User, NoUserInfo>; + + fn get_user_id(&self) -> Result { + self.get_user().map(|user| user.id.0) + } +} + +impl GetUserExt for Message { + fn get_user(&self) -> Result<&User, NoUserInfo> { + self.from.as_ref().ok_or(NoUserInfo) + } +}