From 957bcfb952378c82ea04fb54a37d91ebea7b84a6 Mon Sep 17 00:00:00 2001 From: StNicolay Date: Fri, 5 May 2023 16:45:56 +0300 Subject: [PATCH] Now checking that account exists before getting/deleting it, added export command --- Cargo.lock | 2 + Cargo.toml | 2 + src/decrypted_account.rs | 30 ++++++++++++++ src/entity/account.rs | 41 +++++++++++++++---- src/handlers/commands/add_account.rs | 10 ++++- src/handlers/commands/delete.rs | 8 +++- src/handlers/commands/export.rs | 58 +++++++++++++++++++++++++++ src/handlers/commands/get_account.rs | 8 +++- src/handlers/commands/get_accounts.rs | 2 +- src/handlers/commands/import.rs | 6 +++ src/handlers/commands/mod.rs | 4 ++ src/handlers/markups.rs | 2 +- src/handlers/mod.rs | 53 ++++++++++++------------ src/handlers/state/generic.rs | 5 ++- src/main.rs | 1 + 15 files changed, 193 insertions(+), 39 deletions(-) create mode 100644 src/decrypted_account.rs create mode 100644 src/handlers/commands/export.rs create mode 100644 src/handlers/commands/import.rs diff --git a/Cargo.lock b/Cargo.lock index ff6b1aa..84bcabe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1577,6 +1577,8 @@ dependencies = [ "rand", "scrypt", "sea-orm", + "serde", + "serde_json", "sha2", "teloxide", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 5a236bb..6ad8548 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,8 @@ pretty_env_logger = "0.4.0" rand = { version = "0.8.5", default-features = false, features = ["std_rng"] } scrypt = { version = "0.11.0", default-features = false, features = ["std"] } sea-orm = { version = "0.11.2", features = ["sqlx-mysql", "runtime-tokio-rustls"] } +serde = "1.0.160" +serde_json = "1.0.96" sha2 = "0.10.6" teloxide = { version = "0.12.2", features = ["macros", "ctrlc_handler", "rustls", "throttle"], default-features = false } tokio = { version = "1.27.0", features = ["macros", "rt-multi-thread"] } diff --git a/src/decrypted_account.rs b/src/decrypted_account.rs new file mode 100644 index 0000000..1b83a3f --- /dev/null +++ b/src/decrypted_account.rs @@ -0,0 +1,30 @@ +use crate::entity::account; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize)] +pub struct DecryptedAccount { + pub name: String, + pub login: String, + pub password: String, +} + +impl DecryptedAccount { + pub fn from_account(account: account::Model, master_pass: &str) -> crate::Result { + let name = account.name.clone(); + let (login, password) = account.decrypt(master_pass)?; + Ok(Self { + name, + login, + password, + }) + } + + pub fn into_account( + self, + user_id: u64, + master_pass: &str, + ) -> crate::Result { + let (name, login, password) = (self.name, self.login, self.password); + account::ActiveModel::from_unencrypted(user_id, name, &login, &password, master_pass) + } +} diff --git a/src/entity/account.rs b/src/entity/account.rs index 2a54fa1..771388c 100644 --- a/src/entity/account.rs +++ b/src/entity/account.rs @@ -4,7 +4,7 @@ use chacha20poly1305::{aead::Aead, AeadCore, ChaCha20Poly1305, KeyInit}; use futures::{Stream, TryStreamExt}; use pbkdf2::pbkdf2_hmac_array; use rand::{rngs::OsRng, RngCore}; -use sea_orm::{prelude::*, ActiveValue::Set, QuerySelect}; +use sea_orm::{prelude::*, ActiveValue::Set, QueryOrder, QuerySelect}; use sha2::Sha256; #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)] @@ -88,18 +88,45 @@ impl Model { } impl Entity { - /// Gets a list of account names of a user - pub async fn get_names( + pub async fn get_all( user_id: u64, db: &DatabaseConnection, - ) -> crate::Result> + '_> { + ) -> crate::Result> + '_> { let result = Self::find() - .select_only() - .column(Column::Name) .filter(Column::UserId.eq(user_id)) - .into_tuple() .stream(db) .await?; Ok(result.map_err(Into::into)) } + + /// Gets a list of account names of a user + pub async fn get_names( + user_id: u64, + db: &DatabaseConnection, + ordered: bool, + ) -> crate::Result> + '_> { + let mut select = Self::find() + .select_only() + .column(Column::Name) + .filter(Column::UserId.eq(user_id)); + if ordered { + select = select.order_by_asc(Column::Name); + } + let result = select.into_tuple().stream(db).await?; + Ok(result.map_err(Into::into)) + } + + pub async fn exists( + user_id: u64, + account_name: impl Into, + db: &DatabaseConnection, + ) -> crate::Result { + let result = Self::find_by_id((user_id, account_name.into())) + .select_only() + .column(Column::UserId) + .into_tuple::() + .one(db) + .await?; + Ok(result.is_some()) + } } diff --git a/src/handlers/commands/add_account.rs b/src/handlers/commands/add_account.rs index c4d9627..a15d052 100644 --- a/src/handlers/commands/add_account.rs +++ b/src/handlers/commands/add_account.rs @@ -1,4 +1,4 @@ -use crate::entity::account; +use crate::entity::{account, prelude::Account}; use crate::handlers::{utils::package_handler, MainDialogue, State}; use sea_orm::prelude::*; use teloxide::{adaptors::Throttle, prelude::*}; @@ -85,12 +85,18 @@ async fn get_login( async fn get_account_name( bot: Throttle, msg: Message, - _: DatabaseConnection, + db: DatabaseConnection, dialogue: MainDialogue, previous: Message, name: String, ) -> crate::Result<()> { let _ = bot.delete_message(previous.chat.id, previous.id).await; + let user_id = msg.from().unwrap().id.0; + if Account::exists(user_id, &name, &db).await? { + bot.send_message(msg.chat.id, "Account alreay exists") + .await?; + return Ok(()); + } let previous = bot.send_message(msg.chat.id, "Send login").await?; dialogue .update(State::GetLogin(package_handler( diff --git a/src/handlers/commands/delete.rs b/src/handlers/commands/delete.rs index 2af1d07..4dff7a3 100644 --- a/src/handlers/commands/delete.rs +++ b/src/handlers/commands/delete.rs @@ -25,12 +25,18 @@ async fn get_master_pass( async fn get_account_name( bot: Throttle, msg: Message, - _: DatabaseConnection, + db: DatabaseConnection, dialogue: MainDialogue, previous: Message, name: String, ) -> crate::Result<()> { let _ = bot.delete_message(previous.chat.id, previous.id).await; + let user_id = msg.from().unwrap().id.0; + if !Account::exists(user_id, &name, &db).await? { + bot.send_message(msg.chat.id, "Account doesn't exists") + .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?; diff --git a/src/handlers/commands/export.rs b/src/handlers/commands/export.rs new file mode 100644 index 0000000..4d5d716 --- /dev/null +++ b/src/handlers/commands/export.rs @@ -0,0 +1,58 @@ +use crate::{ + decrypted_account::DecryptedAccount, + entity::prelude::Account, + handlers::{utils::package_handler, MainDialogue, State}, +}; +use futures::TryStreamExt; +use sea_orm::DatabaseConnection; +use serde_json::{json, to_string_pretty}; +use std::sync::Arc; +use teloxide::{adaptors::Throttle, prelude::*, types::InputFile}; +use tokio::task::JoinSet; + +async fn get_master_pass( + bot: Throttle, + msg: Message, + db: DatabaseConnection, + dialogue: MainDialogue, + previous: Message, + master_pass: String, +) -> crate::Result<()> { + let _ = bot.delete_message(previous.chat.id, previous.id).await; + let master_pass: Arc = master_pass.into(); + let user_id = msg.from().unwrap().id.0; + let mut join_set = JoinSet::new(); + let mut accounts = Vec::new(); + Account::get_all(user_id, &db) + .await? + .try_for_each(|account| { + let master_pass = Arc::clone(&master_pass); + join_set.spawn_blocking(move || DecryptedAccount::from_account(account, &master_pass)); + async { crate::Result::Ok(()) } + }) + .await?; + drop(master_pass); + while let Some(account) = join_set.join_next().await.transpose()?.transpose()? { + accounts.push(account) + } + accounts.sort_by(|this, other| this.name.cmp(&other.name)); + let json = to_string_pretty(&json!({ "accounts": accounts }))?; + let file = InputFile::memory(json).file_name("accounts.json"); + bot.send_document(msg.chat.id, file).await?; + dialogue.exit().await?; + Ok(()) +} + +pub async fn export(bot: Throttle, msg: Message, dialogue: MainDialogue) -> crate::Result<()> { + let previous = bot + .send_message(msg.chat.id, "Send a master password to export the accounts") + .await?; + dialogue + .update(State::GetMasterPass(package_handler( + move |bot, msg, db, dialogue, master_pass| { + get_master_pass(bot, msg, db, dialogue, previous, master_pass) + }, + ))) + .await?; + Ok(()) +} diff --git a/src/handlers/commands/get_account.rs b/src/handlers/commands/get_account.rs index a48aee5..5568833 100644 --- a/src/handlers/commands/get_account.rs +++ b/src/handlers/commands/get_account.rs @@ -35,12 +35,18 @@ async fn get_master_pass( async fn get_account_name( bot: Throttle, msg: Message, - _: DatabaseConnection, + db: DatabaseConnection, dialogue: MainDialogue, previous: Message, name: String, ) -> crate::Result<()> { let _ = bot.delete_message(previous.chat.id, previous.id).await; + let user_id = msg.from().unwrap().id.0; + if !Account::exists(user_id, &name, &db).await? { + bot.send_message(msg.chat.id, "Account doesn't exists") + .await?; + return Ok(()); + } let previous = bot .send_message(msg.chat.id, "Send master password") .await?; diff --git a/src/handlers/commands/get_accounts.rs b/src/handlers/commands/get_accounts.rs index e275740..17204b5 100644 --- a/src/handlers/commands/get_accounts.rs +++ b/src/handlers/commands/get_accounts.rs @@ -9,7 +9,7 @@ pub async fn get_accounts( db: DatabaseConnection, ) -> crate::Result<()> { let user_id = msg.from().unwrap().id.0; - let mut account_names = Account::get_names(user_id, &db).await?; + let mut account_names = Account::get_names(user_id, &db, true).await?; let mut result = match account_names.try_next().await? { Some(name) => format!("Accounts:\n`{name}`"), None => { diff --git a/src/handlers/commands/import.rs b/src/handlers/commands/import.rs new file mode 100644 index 0000000..9b6ea81 --- /dev/null +++ b/src/handlers/commands/import.rs @@ -0,0 +1,6 @@ +use crate::handlers::MainDialogue; +use teloxide::{adaptors::Throttle, prelude::*}; + +pub async fn import(bot: Throttle, msg: Message, dialogue: MainDialogue) -> crate::Result<()> { + Ok(()) +} diff --git a/src/handlers/commands/mod.rs b/src/handlers/commands/mod.rs index 07e89e0..90ee1c8 100644 --- a/src/handlers/commands/mod.rs +++ b/src/handlers/commands/mod.rs @@ -2,16 +2,20 @@ mod add_account; mod default; mod delete; mod delete_all; +mod export; mod get_account; mod get_accounts; mod help; +mod import; mod set_master_pass; pub use add_account::add_account; pub use default::default; pub use delete::delete; pub use delete_all::delete_all; +pub use export::export; pub use get_account::get_account; pub use get_accounts::get_accounts; pub use help::help; +pub use import::import; pub use set_master_pass::set_master_pass; diff --git a/src/handlers/markups.rs b/src/handlers/markups.rs index 15a995e..e42b350 100644 --- a/src/handlers/markups.rs +++ b/src/handlers/markups.rs @@ -8,7 +8,7 @@ pub async fn account_markup( user_id: u64, db: &DatabaseConnection, ) -> crate::Result { - let account_names: Vec> = Account::get_names(user_id, db) + let account_names: Vec> = Account::get_names(user_id, db, true) .await? .map_ok(|account| KeyboardButton::new(account)) .try_chunks(3) diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index 0fc3a7d..e83ed17 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -1,3 +1,8 @@ +mod commands; +mod markups; +mod state; +mod utils; + use sea_orm::prelude::*; use std::sync::Arc; use teloxide::{ @@ -9,30 +14,6 @@ use teloxide::{ }; use tokio::sync::Mutex; -mod commands; -mod markups; -mod state; -mod utils; - -#[derive(BotCommands, Clone, Copy)] -#[command(rename_rule = "snake_case")] -enum Command { - #[command()] - Help, - #[command()] - AddAccount, - #[command()] - GetAccount, - #[command()] - GetAccounts, - #[command()] - SetMasterPass, - #[command()] - Delete, - #[command()] - DeleteAll, -} - type MainDialogue = Dialogue>; type PackagedHandler = Arc< Mutex< @@ -52,6 +33,27 @@ type PackagedHandler = Arc< >, >; +#[derive(BotCommands, Clone, Copy)] +#[command(rename_rule = "snake_case")] +enum Command { + #[command()] + Help, + #[command()] + AddAccount, + #[command()] + GetAccount, + #[command()] + GetAccounts, + #[command()] + SetMasterPass, + #[command()] + Delete, + #[command()] + DeleteAll, + #[command()] + Export, +} + #[derive(Default, Clone)] pub enum State { #[default] @@ -76,7 +78,8 @@ pub fn get_dispatcher( .branch(case![Command::GetAccounts].endpoint(commands::get_accounts)) .branch(case![Command::SetMasterPass].endpoint(commands::set_master_pass)) .branch(case![Command::Delete].endpoint(commands::delete)) - .branch(case![Command::DeleteAll].endpoint(commands::delete_all)); + .branch(case![Command::DeleteAll].endpoint(commands::delete_all)) + .branch(case![Command::Export].endpoint(commands::export)); let message_handler = Update::filter_message() .map_async(utils::delete_message) diff --git a/src/handlers/state/generic.rs b/src/handlers/state/generic.rs index f561f13..85b4145 100644 --- a/src/handlers/state/generic.rs +++ b/src/handlers/state/generic.rs @@ -28,7 +28,10 @@ where } match check(&bot, &msg, &db, &text).await { Ok(true) => (), - Ok(false) => dialogue.exit().await?, + Ok(false) => { + dialogue.exit().await?; + return Ok(()); + } Err(err) => { let _ = dialogue.exit().await; return Err(err); diff --git a/src/main.rs b/src/main.rs index d2b57e7..1fdf9a9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +mod decrypted_account; mod entity; mod handlers;