StNicolay/pass_manager
1
0
Fork 0
Code Issues Pull Requests Releases Activity

Added state persistance

Browse Source
This commit is contained in:
StNicolay 2023-05-28 02:21:50 +03:00
parent ef0ccfb457
commit 40f7194cbe
Signed by: StNicolay
GPG Key ID: 9693D04DCD962B0D
14 changed files with 265 additions and 221 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
src/handlers/commands/add_account.rs
View File

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

19
src/handlers/commands/delete.rs
View File

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

19
src/handlers/commands/get_account.rs
View File

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

48
src/handlers/commands/import.rs
View File

@ -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<DecryptedAccount>,
user: User,
) -> crate::Result<()> {
let user_id = msg.from().ok_or(NoUserInfo)?.id.0;
let accounts = user.accounts;
let failed: Vec<String> = {
let master_pass: Arc<str> = 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::<User>(&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<Bot>, msg: Message, dialogue: MainDialogue) ->
)
.await?;
dialogue
.update(State::GetDocument(Handler::new(get_document, previous)))
.update(State::GetUser(Handler::new(get_document, previous)))
.await?;
Ok(())
}

19
src/handlers/mod.rs
View File

@ -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<State, InMemStorage<State>>;
#[derive(Default, Clone)]
pub enum State {
#[default]
Start,
GetAccountName(PackagedHandler<String>),
GetMasterPass(PackagedHandler<String>),
GetLogin(PackagedHandler<String>),
GetPassword(PackagedHandler<String>),
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::<Update, InMemStorage<State>, 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);

38
src/handlers/state/generic.rs
View File

@ -22,19 +22,23 @@ where
&'a Message,
&'a DatabaseConnection,
&'a str,
) -> PinnedFuture<'a, crate::Result<bool>>,
) -> PinnedFuture<'a, crate::Result<Option<Message>>>,
{
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")
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 => {
if handler.func.is_none() {
let _ = dialogue.exit().await;
return Err(HandlerUsed.into());
}
};
match check(&bot, &msg, &db, text).await {
Ok(true) => (),
Ok(false) => {
dialogue.exit().await?;
if let Some(failure_message) = check(&bot, &msg, &db, text).await? {
handler.previous = Some(failure_message);
return Ok(());
}
Err(err) => {
let _ = dialogue.exit().await;
return Err(err);
}
};
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);
}

66
src/handlers/state/get_account_name.rs
View File

@ -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<Bot>,
msg: &Message,
db: &DatabaseConnection,
name: &str,
check_kind: NameCheckKind,
) -> crate::Result<Option<Message>> {
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<Bot>,
msg: Message,
db: DatabaseConnection,
dialogue: MainDialogue,
next: super::PackagedHandler<String>,
(next, check_kind): (super::PackagedHandler<String>, 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

45
src/handlers/state/get_document.rs
View File

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

9
src/handlers/state/get_login.rs
View File

@ -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,

38
src/handlers/state/get_master_pass.rs
View File

@ -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<Bot>,
msg: &'a Message,
db: &'a DatabaseConnection,
master_pass: &'a str,
) -> crate::Result<bool> {
pub async fn check_master_pass(
bot: &Throttle<Bot>,
msg: &Message,
db: &DatabaseConnection,
master_pass: &str,
) -> crate::Result<Option<Message>> {
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(

9
src/handlers/state/get_password.rs
View File

@ -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,

88
src/handlers/state/get_user.rs Normal file
View File

@ -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<Bot>,
msg: Message,
db: DatabaseConnection,
dialogue: MainDialogue,
next: super::PackagedHandler<User>,
) -> 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(())
}

17
src/handlers/state/handler.rs
View File

@ -4,9 +4,7 @@ use std::{future::Future, sync::Arc};
use teloxide::{adaptors::Throttle, prelude::*};
use tokio::sync::Mutex;
#[allow(clippy::type_complexity)]
pub struct Handler<T> {
pub handler: Box<
type DynHanlder<T> = Box<
dyn FnOnce(
Throttle<Bot>,
Message,
@ -16,12 +14,15 @@ pub struct Handler<T> {
) -> crate::PinnedFuture<'static, crate::Result<()>>
+ Send
+ Sync,
>,
>;
pub struct Handler<T> {
pub func: Option<DynHanlder<T>>,
pub previous: Option<Message>,
}
pub type PackagedHandler<T> = Arc<Mutex<Option<Handler<T>>>>;
pub type PackagedHandler<T> = Arc<Mutex<Handler<T>>>;
impl<T> Handler<T> {
/// Convinience method to convert a simple async function and a previous message into PackagedHandler
@ -35,11 +36,11 @@ impl<T> Handler<T> {
F: Future<Output = crate::Result<()>> + 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))
}
}

22
src/handlers/state/mod.rs
View File

@ -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<String>, NameCheckKind),
GetMasterPass(PackagedHandler<String>),
GetLogin(PackagedHandler<String>),
GetPassword(PackagedHandler<String>),
GetUser(PackagedHandler<User>),
}
pub type MainDialogue = Dialogue<State, InMemStorage<State>>;