From fc76d7f65c08f6187d3ec4040d1662ef535d0da4 Mon Sep 17 00:00:00 2001 From: StNicolay Date: Fri, 5 May 2023 21:29:58 +0300 Subject: [PATCH] Added import command, created User struct to remove the usage of a json! macro --- Cargo.lock | 9 +- Cargo.toml | 1 + src/handlers/commands/export.rs | 6 +- src/handlers/commands/import.rs | 128 +++++++++++++++++++++++- src/handlers/mod.rs | 7 +- src/handlers/state/get_document.rs | 13 +++ src/handlers/state/mod.rs | 2 + src/main.rs | 2 +- src/{decrypted_account.rs => models.rs} | 5 + 9 files changed, 161 insertions(+), 12 deletions(-) create mode 100644 src/handlers/state/get_document.rs rename src/{decrypted_account.rs => models.rs} (89%) diff --git a/Cargo.lock b/Cargo.lock index 84bcabe..bb1b9dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1570,6 +1570,7 @@ dependencies = [ "chacha20poly1305", "dotenv", "futures", + "itertools 0.10.5", "log", "migration", "pbkdf2", @@ -2285,18 +2286,18 @@ checksum = "bebd363326d05ec3e2f532ab7660680f3b02130d780c299bca73469d521bc0ed" [[package]] name = "serde" -version = "1.0.160" +version = "1.0.162" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb2f3770c8bce3bcda7e149193a069a0f4365bda1fa5cd88e03bca26afc1216c" +checksum = "71b2f6e1ab5c2b98c05f0f35b236b22e8df7ead6ffbf51d7808da7f8817e7ab6" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.160" +version = "1.0.162" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "291a097c63d8497e00160b166a967a4a79c64f3facdd01cbd7502231688d77df" +checksum = "a2a0814352fd64b58489904a44ea8d90cb1a91dcb6b4f5ebabc32c8318e93cb6" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 6ad8548..037ffe0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ anyhow = "1.0.70" chacha20poly1305 = { version = "0.10.1", features = ["std"] } dotenv = "0.15.0" futures = "0.3.28" +itertools = "0.10.5" log = "0.4.17" migration = { version = "0.1.0", path = "migration" } pbkdf2 = "0.12.1" diff --git a/src/handlers/commands/export.rs b/src/handlers/commands/export.rs index 4d5d716..2bf732c 100644 --- a/src/handlers/commands/export.rs +++ b/src/handlers/commands/export.rs @@ -1,11 +1,11 @@ use crate::{ - decrypted_account::DecryptedAccount, entity::prelude::Account, handlers::{utils::package_handler, MainDialogue, State}, + models::{DecryptedAccount, User}, }; use futures::TryStreamExt; use sea_orm::DatabaseConnection; -use serde_json::{json, to_string_pretty}; +use serde_json::to_string_pretty; use std::sync::Arc; use teloxide::{adaptors::Throttle, prelude::*, types::InputFile}; use tokio::task::JoinSet; @@ -36,7 +36,7 @@ async fn get_master_pass( accounts.push(account) } accounts.sort_by(|this, other| this.name.cmp(&other.name)); - let json = to_string_pretty(&json!({ "accounts": accounts }))?; + let json = to_string_pretty(&User { accounts })?; let file = InputFile::memory(json).file_name("accounts.json"); bot.send_document(msg.chat.id, file).await?; dialogue.exit().await?; diff --git a/src/handlers/commands/import.rs b/src/handlers/commands/import.rs index 9b6ea81..4d02cd3 100644 --- a/src/handlers/commands/import.rs +++ b/src/handlers/commands/import.rs @@ -1,6 +1,128 @@ -use crate::handlers::MainDialogue; -use teloxide::{adaptors::Throttle, prelude::*}; +use crate::{ + handlers::{utils::package_handler, MainDialogue, State}, + models::{DecryptedAccount, User}, +}; +use futures::TryStreamExt; +use itertools::Itertools; +use sea_orm::prelude::*; +use serde_json::from_slice; +use std::sync::Arc; +use teloxide::{adaptors::Throttle, net::Download, prelude::*}; +use tokio::task::{spawn_blocking, JoinSet}; -pub async fn import(bot: Throttle, msg: Message, dialogue: MainDialogue) -> crate::Result<()> { +async fn get_master_pass( + bot: Throttle, + msg: Message, + db: DatabaseConnection, + dialogue: MainDialogue, + master_pass: String, + previous: Message, + accounts: Vec, +) -> crate::Result<()> { + let _ = bot.delete_message(previous.chat.id, previous.id).await; + let user_id = msg.from().unwrap().id.0; + let master_pass: Arc = master_pass.into(); + let mut join_set = JoinSet::new(); + let mut failed = Vec::new(); + for account in accounts { + let master_pass = Arc::clone(&master_pass); + let db = db.clone(); + join_set.spawn(async move { + let name = account.name.clone(); + let account = + match spawn_blocking(move || account.into_account(user_id, &master_pass)).await { + Ok(Ok(account)) => account, + _ => return Err(name), + }; + account.insert(&db).await.map_err(|_| name)?; + Ok(()) + }); + } + while let Some(result) = join_set.join_next().await.transpose()? { + match result { + Ok(()) => (), + Err(name) => failed.push(name), + } + } + if failed.is_empty() { + bot.send_message(msg.chat.id, "Success").await?; + } else { + let text = format!( + "Failed to create the following accounts:\n{}", + failed.into_iter().format("\n") + ); + bot.send_message(msg.chat.id, text).await?; + } + dialogue.exit().await?; + Ok(()) +} + +async fn get_document( + bot: Throttle, + msg: Message, + _: DatabaseConnection, + dialogue: MainDialogue, + previous: Message, +) -> crate::Result<()> { + let _ = bot.delete_message(previous.chat.id, previous.id).await; + let document = match msg.document() { + Some(doc) => doc, + None => { + bot.send_message(msg.chat.id, "You didn't send a file") + .await?; + dialogue.exit().await?; + return Ok(()); + } + }; + match document.file_name { + Some(ref name) if name.trim().ends_with(".json") => (), + _ => { + bot.send_message(msg.chat.id, "Invalid file name").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") + .await?; + dialogue.exit().await?; + return Ok(()); + } + }; + let previous = bot + .send_message(msg.chat.id, "Send a new master password") + .await?; + dialogue + .update(State::GetMasterPass(package_handler( + move |bot, msg, db, dialogue, master_pass| { + get_master_pass(bot, msg, db, dialogue, master_pass, previous, accounts) + }, + ))) + .await?; + Ok(()) +} + +pub async fn import(bot: Throttle, msg: Message, dialogue: MainDialogue) -> crate::Result<()> { + let previous = bot + .send_message( + msg.chat.id, + "Send a json document with the same format as created by /export", + ) + .await?; + dialogue + .update(State::GetDocument(package_handler( + move |bot, msg, db, dialogue, ()| get_document(bot, msg, db, dialogue, previous), + ))) + .await?; Ok(()) } diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index e83ed17..eb72e43 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -52,6 +52,8 @@ enum Command { DeleteAll, #[command()] Export, + #[command()] + Import, } #[derive(Default, Clone)] @@ -62,6 +64,7 @@ pub enum State { GetMasterPass(PackagedHandler), GetLogin(PackagedHandler), GetPassword(PackagedHandler), + GetDocument(PackagedHandler<()>), } pub fn get_dispatcher( @@ -79,7 +82,8 @@ pub fn get_dispatcher( .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::Export].endpoint(commands::export)); + .branch(case![Command::Export].endpoint(commands::export)) + .branch(case![Command::Import].endpoint(commands::import)); let message_handler = Update::filter_message() .map_async(utils::delete_message) @@ -88,6 +92,7 @@ pub fn get_dispatcher( .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(command_handler) .branch(endpoint(commands::default)); diff --git a/src/handlers/state/get_document.rs b/src/handlers/state/get_document.rs new file mode 100644 index 0000000..07c8665 --- /dev/null +++ b/src/handlers/state/get_document.rs @@ -0,0 +1,13 @@ +use crate::handlers::{MainDialogue, PackagedHandler}; +use sea_orm::prelude::*; +use teloxide::{adaptors::Throttle, prelude::*}; + +pub async fn get_document( + bot: Throttle, + msg: Message, + db: DatabaseConnection, + dialogue: MainDialogue, + next: PackagedHandler<()>, +) -> crate::Result<()> { + next.lock().await.take().unwrap()(bot, msg, db, dialogue, ()).await +} diff --git a/src/handlers/state/mod.rs b/src/handlers/state/mod.rs index c431648..db0e503 100644 --- a/src/handlers/state/mod.rs +++ b/src/handlers/state/mod.rs @@ -1,10 +1,12 @@ mod generic; mod get_account_name; +mod get_document; mod get_login; mod get_master_pass; mod get_password; pub use get_account_name::get_account_name; +pub use get_document::get_document; pub use get_login::get_login; pub use get_master_pass::get_master_pass; pub use get_password::get_password; diff --git a/src/main.rs b/src/main.rs index 1fdf9a9..f88b544 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,6 @@ -mod decrypted_account; mod entity; mod handlers; +mod models; use anyhow::{Error, Result}; use dotenv::dotenv; diff --git a/src/decrypted_account.rs b/src/models.rs similarity index 89% rename from src/decrypted_account.rs rename to src/models.rs index 1b83a3f..e497920 100644 --- a/src/decrypted_account.rs +++ b/src/models.rs @@ -28,3 +28,8 @@ impl DecryptedAccount { account::ActiveModel::from_unencrypted(user_id, name, &login, &password, master_pass) } } + +#[derive(Serialize, Deserialize)] +pub struct User { + pub accounts: Vec, +}