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

Cargo.lock

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

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"] }

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<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)
+    }
+}

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<impl Stream<Item = crate::Result<String>> + '_> {
+    ) -> crate::Result<impl Stream<Item = crate::Result<Model>> + '_> {
         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<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())
+    }
 }

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<Bot>,
     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(

src/handlers/commands/delete.rs

@@ -25,12 +25,18 @@ async fn get_master_pass(
 async fn get_account_name(
     bot: Throttle<Bot>,
     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?;

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<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(())
+}

src/handlers/commands/get_account.rs

@@ -35,12 +35,18 @@ async fn get_master_pass(
 async fn get_account_name(
     bot: Throttle<Bot>,
     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?;

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 => {

src/handlers/commands/import.rs

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

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;

src/handlers/markups.rs

@@ -8,7 +8,7 @@ pub async fn account_markup(
     user_id: u64,
     db: &DatabaseConnection,
 ) -> 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?
         .map_ok(|account| KeyboardButton::new(account))
         .try_chunks(3)

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<State, InMemStorage<State>>;
 type PackagedHandler<T> = Arc<
     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)]
 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)

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

src/main.rs

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