Added import command, created User struct to remove the usage of a json! macro

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

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"

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?;

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<Bot>, msg: Message, dialogue: MainDialogue) -> crate::Result<()> {
+async fn get_master_pass(
+    bot: Throttle<Bot>,
+    msg: Message,
+    db: DatabaseConnection,
+    dialogue: MainDialogue,
+    master_pass: String,
+    previous: Message,
+    accounts: Vec<DecryptedAccount>,
+) -> crate::Result<()> {
+    let _ = bot.delete_message(previous.chat.id, previous.id).await;
+    let user_id = msg.from().unwrap().id.0;
+    let master_pass: Arc<str> = 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<Bot>,
+    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::<User>(&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<Bot>, 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(())
 }

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<String>),
     GetLogin(PackagedHandler<String>),
     GetPassword(PackagedHandler<String>),
+    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));
 

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<Bot>,
+    msg: Message,
+    db: DatabaseConnection,
+    dialogue: MainDialogue,
+    next: PackagedHandler<()>,
+) -> crate::Result<()> {
+    next.lock().await.take().unwrap()(bot, msg, db, dialogue, ()).await
+}

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;

src/main.rs

@@ -1,6 +1,6 @@
-mod decrypted_account;
 mod entity;
 mod handlers;
+mod models;
 
 use anyhow::{Error, Result};
 use dotenv::dotenv;

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<DecryptedAccount>,
+}