Now checking that account exists before getting/deleting it, added export command

This commit is contained in:
StNicolay 2023-05-05 16:45:56 +03:00
parent b5e003e1d7
commit 957bcfb952
Signed by: StNicolay
GPG Key ID: 9693D04DCD962B0D
15 changed files with 193 additions and 39 deletions

2
Cargo.lock generated
View File

@ -1577,6 +1577,8 @@ dependencies = [
"rand", "rand",
"scrypt", "scrypt",
"sea-orm", "sea-orm",
"serde",
"serde_json",
"sha2", "sha2",
"teloxide", "teloxide",
"tokio", "tokio",

View File

@ -23,6 +23,8 @@ pretty_env_logger = "0.4.0"
rand = { version = "0.8.5", default-features = false, features = ["std_rng"] } rand = { version = "0.8.5", default-features = false, features = ["std_rng"] }
scrypt = { version = "0.11.0", default-features = false, features = ["std"] } scrypt = { version = "0.11.0", default-features = false, features = ["std"] }
sea-orm = { version = "0.11.2", features = ["sqlx-mysql", "runtime-tokio-rustls"] } 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" sha2 = "0.10.6"
teloxide = { version = "0.12.2", features = ["macros", "ctrlc_handler", "rustls", "throttle"], default-features = false } teloxide = { version = "0.12.2", features = ["macros", "ctrlc_handler", "rustls", "throttle"], default-features = false }
tokio = { version = "1.27.0", features = ["macros", "rt-multi-thread"] } tokio = { version = "1.27.0", features = ["macros", "rt-multi-thread"] }

30
src/decrypted_account.rs Normal file
View File

@ -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<Self> {
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<account::ActiveModel> {
let (name, login, password) = (self.name, self.login, self.password);
account::ActiveModel::from_unencrypted(user_id, name, &login, &password, master_pass)
}
}

View File

@ -4,7 +4,7 @@ use chacha20poly1305::{aead::Aead, AeadCore, ChaCha20Poly1305, KeyInit};
use futures::{Stream, TryStreamExt}; use futures::{Stream, TryStreamExt};
use pbkdf2::pbkdf2_hmac_array; use pbkdf2::pbkdf2_hmac_array;
use rand::{rngs::OsRng, RngCore}; use rand::{rngs::OsRng, RngCore};
use sea_orm::{prelude::*, ActiveValue::Set, QuerySelect}; use sea_orm::{prelude::*, ActiveValue::Set, QueryOrder, QuerySelect};
use sha2::Sha256; use sha2::Sha256;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)] #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
@ -88,18 +88,45 @@ impl Model {
} }
impl Entity { impl Entity {
/// Gets a list of account names of a user pub async fn get_all(
pub async fn get_names(
user_id: u64, user_id: u64,
db: &DatabaseConnection, db: &DatabaseConnection,
) -> crate::Result<impl Stream<Item = crate::Result<String>> + '_> { ) -> crate::Result<impl Stream<Item = crate::Result<Model>> + '_> {
let result = Self::find() let result = Self::find()
.select_only()
.column(Column::Name)
.filter(Column::UserId.eq(user_id)) .filter(Column::UserId.eq(user_id))
.into_tuple()
.stream(db) .stream(db)
.await?; .await?;
Ok(result.map_err(Into::into)) 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<impl Stream<Item = crate::Result<String>> + '_> {
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<String>,
db: &DatabaseConnection,
) -> crate::Result<bool> {
let result = Self::find_by_id((user_id, account_name.into()))
.select_only()
.column(Column::UserId)
.into_tuple::<u64>()
.one(db)
.await?;
Ok(result.is_some())
}
} }

View File

@ -1,4 +1,4 @@
use crate::entity::account; use crate::entity::{account, prelude::Account};
use crate::handlers::{utils::package_handler, MainDialogue, State}; use crate::handlers::{utils::package_handler, MainDialogue, State};
use sea_orm::prelude::*; use sea_orm::prelude::*;
use teloxide::{adaptors::Throttle, prelude::*}; use teloxide::{adaptors::Throttle, prelude::*};
@ -85,12 +85,18 @@ async fn get_login(
async fn get_account_name( async fn get_account_name(
bot: Throttle<Bot>, bot: Throttle<Bot>,
msg: Message, msg: Message,
_: DatabaseConnection, db: DatabaseConnection,
dialogue: MainDialogue, dialogue: MainDialogue,
previous: Message, previous: Message,
name: String, name: String,
) -> crate::Result<()> { ) -> crate::Result<()> {
let _ = bot.delete_message(previous.chat.id, previous.id).await; 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?; let previous = bot.send_message(msg.chat.id, "Send login").await?;
dialogue dialogue
.update(State::GetLogin(package_handler( .update(State::GetLogin(package_handler(

View File

@ -25,12 +25,18 @@ async fn get_master_pass(
async fn get_account_name( async fn get_account_name(
bot: Throttle<Bot>, bot: Throttle<Bot>,
msg: Message, msg: Message,
_: DatabaseConnection, db: DatabaseConnection,
dialogue: MainDialogue, dialogue: MainDialogue,
previous: Message, previous: Message,
name: String, name: String,
) -> crate::Result<()> { ) -> crate::Result<()> {
let _ = bot.delete_message(previous.chat.id, previous.id).await; 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 let previous = bot
.send_message(msg.chat.id, "Send master password. Once you send correct master password the account is unrecoverable") .send_message(msg.chat.id, "Send master password. Once you send correct master password the account is unrecoverable")
.await?; .await?;

View File

@ -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<Bot>,
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<str> = 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<Bot>, 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(())
}

View File

@ -35,12 +35,18 @@ async fn get_master_pass(
async fn get_account_name( async fn get_account_name(
bot: Throttle<Bot>, bot: Throttle<Bot>,
msg: Message, msg: Message,
_: DatabaseConnection, db: DatabaseConnection,
dialogue: MainDialogue, dialogue: MainDialogue,
previous: Message, previous: Message,
name: String, name: String,
) -> crate::Result<()> { ) -> crate::Result<()> {
let _ = bot.delete_message(previous.chat.id, previous.id).await; 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 let previous = bot
.send_message(msg.chat.id, "Send master password") .send_message(msg.chat.id, "Send master password")
.await?; .await?;

View File

@ -9,7 +9,7 @@ pub async fn get_accounts(
db: DatabaseConnection, db: DatabaseConnection,
) -> crate::Result<()> { ) -> crate::Result<()> {
let user_id = msg.from().unwrap().id.0; 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? { let mut result = match account_names.try_next().await? {
Some(name) => format!("Accounts:\n`{name}`"), Some(name) => format!("Accounts:\n`{name}`"),
None => { None => {

View File

@ -0,0 +1,6 @@
use crate::handlers::MainDialogue;
use teloxide::{adaptors::Throttle, prelude::*};
pub async fn import(bot: Throttle<Bot>, msg: Message, dialogue: MainDialogue) -> crate::Result<()> {
Ok(())
}

View File

@ -2,16 +2,20 @@ mod add_account;
mod default; mod default;
mod delete; mod delete;
mod delete_all; mod delete_all;
mod export;
mod get_account; mod get_account;
mod get_accounts; mod get_accounts;
mod help; mod help;
mod import;
mod set_master_pass; mod set_master_pass;
pub use add_account::add_account; pub use add_account::add_account;
pub use default::default; pub use default::default;
pub use delete::delete; pub use delete::delete;
pub use delete_all::delete_all; pub use delete_all::delete_all;
pub use export::export;
pub use get_account::get_account; pub use get_account::get_account;
pub use get_accounts::get_accounts; pub use get_accounts::get_accounts;
pub use help::help; pub use help::help;
pub use import::import;
pub use set_master_pass::set_master_pass; pub use set_master_pass::set_master_pass;

View File

@ -8,7 +8,7 @@ pub async fn account_markup(
user_id: u64, user_id: u64,
db: &DatabaseConnection, db: &DatabaseConnection,
) -> crate::Result<KeyboardMarkup> { ) -> crate::Result<KeyboardMarkup> {
let account_names: Vec<Vec<KeyboardButton>> = Account::get_names(user_id, db) let account_names: Vec<Vec<KeyboardButton>> = Account::get_names(user_id, db, true)
.await? .await?
.map_ok(|account| KeyboardButton::new(account)) .map_ok(|account| KeyboardButton::new(account))
.try_chunks(3) .try_chunks(3)

View File

@ -1,3 +1,8 @@
mod commands;
mod markups;
mod state;
mod utils;
use sea_orm::prelude::*; use sea_orm::prelude::*;
use std::sync::Arc; use std::sync::Arc;
use teloxide::{ use teloxide::{
@ -9,30 +14,6 @@ use teloxide::{
}; };
use tokio::sync::Mutex; 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<State, InMemStorage<State>>; type MainDialogue = Dialogue<State, InMemStorage<State>>;
type PackagedHandler<T> = Arc< type PackagedHandler<T> = Arc<
Mutex< Mutex<
@ -52,6 +33,27 @@ type PackagedHandler<T> = 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)] #[derive(Default, Clone)]
pub enum State { pub enum State {
#[default] #[default]
@ -76,7 +78,8 @@ pub fn get_dispatcher(
.branch(case![Command::GetAccounts].endpoint(commands::get_accounts)) .branch(case![Command::GetAccounts].endpoint(commands::get_accounts))
.branch(case![Command::SetMasterPass].endpoint(commands::set_master_pass)) .branch(case![Command::SetMasterPass].endpoint(commands::set_master_pass))
.branch(case![Command::Delete].endpoint(commands::delete)) .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() let message_handler = Update::filter_message()
.map_async(utils::delete_message) .map_async(utils::delete_message)

View File

@ -28,7 +28,10 @@ where
} }
match check(&bot, &msg, &db, &text).await { match check(&bot, &msg, &db, &text).await {
Ok(true) => (), Ok(true) => (),
Ok(false) => dialogue.exit().await?, Ok(false) => {
dialogue.exit().await?;
return Ok(());
}
Err(err) => { Err(err) => {
let _ = dialogue.exit().await; let _ = dialogue.exit().await;
return Err(err); return Err(err);

View File

@ -1,3 +1,4 @@
mod decrypted_account;
mod entity; mod entity;
mod handlers; mod handlers;