From 83aa2c90b142ab33c0b6ae5ab8b5a48dd0823d9a Mon Sep 17 00:00:00 2001
From: StNicolay <stnicolay2007@mail.ru>
Date: Thu, 4 May 2023 19:03:03 +0300
Subject: [PATCH] Moved the check of existence of a master password into
 master_pass.rs, added buttons for /get_account

---
 Cargo.lock                               |  8 ++++----
 src/entity/master_pass.rs                | 14 +++++++++++++-
 src/handlers/commands/get_account.rs     | 21 +++++++++++++++++++--
 src/handlers/commands/set_master_pass.rs | 22 ++++++++--------------
 4 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7f50da1..ff6b1aa 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1324,9 +1324,9 @@ dependencies = [
 
 [[package]]
 name = "linux-raw-sys"
-version = "0.3.6"
+version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b64f40e5e03e0d54f03845c8197d0291253cdbedfb1cb46b13c2c117554a9f4c"
+checksum = "ece97ea872ece730aed82664c424eb4c8291e1ff2480247ccf7409044bc6479f"
 
 [[package]]
 name = "lock_api"
@@ -2018,9 +2018,9 @@ dependencies = [
 
 [[package]]
 name = "rustix"
-version = "0.37.18"
+version = "0.37.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8bbfc1d1c7c40c01715f47d71444744a81669ca84e8b63e25a55e169b1f86433"
+checksum = "acf8729d8542766f1b2cf77eb034d52f40d375bb8b615d0b147089946e16613d"
 dependencies = [
  "bitflags",
  "errno",
diff --git a/src/entity/master_pass.rs b/src/entity/master_pass.rs
index ac382af..df86a43 100644
--- a/src/entity/master_pass.rs
+++ b/src/entity/master_pass.rs
@@ -2,7 +2,7 @@
 
 use rand::{rngs::OsRng, RngCore};
 use scrypt::{scrypt, Params};
-use sea_orm::{entity::prelude::*, ActiveValue::Set};
+use sea_orm::{entity::prelude::*, ActiveValue::Set, QuerySelect};
 
 #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
 #[sea_orm(table_name = "master_pass")]
@@ -54,4 +54,16 @@ impl Entity {
         let password_hash = hash_password(master_pass.as_ref(), &model.salt)?;
         Ok(Some(password_hash == model.password_hash))
     }
+
+    /// Checks if the master password for the user exists
+    pub async fn exists(user_id: u64, db: &DatabaseConnection) -> crate::Result<bool> {
+        let id = Self::find()
+            .select_only()
+            .column(Column::UserId)
+            .filter(Column::UserId.eq(user_id))
+            .into_tuple::<u64>()
+            .one(db)
+            .await?;
+        Ok(id.is_some())
+    }
 }
diff --git a/src/handlers/commands/get_account.rs b/src/handlers/commands/get_account.rs
index 64cc1a2..96e5ecb 100644
--- a/src/handlers/commands/get_account.rs
+++ b/src/handlers/commands/get_account.rs
@@ -1,7 +1,13 @@
 use crate::entity::{account, prelude::Account};
 use crate::handlers::{utils::package_handler, MainDialogue, State};
+use futures::TryStreamExt;
 use sea_orm::prelude::*;
-use teloxide::{adaptors::Throttle, prelude::*, types::ParseMode};
+use teloxide::{
+    adaptors::Throttle,
+    prelude::*,
+    types::ParseMode,
+    types::{KeyboardButton, KeyboardMarkup},
+};
 use tokio::task::spawn_blocking;
 
 async fn get_master_pass(
@@ -12,6 +18,7 @@ async fn get_master_pass(
     name: String,
     master_pass: String,
 ) -> crate::Result<()> {
+    let _ = bot.delete_message(msg.chat.id, msg.id).await;
     dialogue.exit().await?;
     let user_id = msg.from().unwrap().id.0;
     let account = Account::find()
@@ -52,8 +59,18 @@ pub async fn get_account(
     bot: Throttle<Bot>,
     msg: Message,
     dialogue: MainDialogue,
+    db: DatabaseConnection,
 ) -> crate::Result<()> {
-    bot.send_message(msg.chat.id, "Send account name").await?;
+    let account_names: Vec<Vec<KeyboardButton>> = Account::get_names(msg.from().unwrap().id.0, &db)
+        .await?
+        .map_ok(|account| KeyboardButton::new(account))
+        .try_chunks(3)
+        .try_collect()
+        .await?;
+    let markup = KeyboardMarkup::new(account_names).resize_keyboard(true);
+    bot.send_message(msg.chat.id, "Send account name")
+        .reply_markup(markup)
+        .await?;
     dialogue
         .update(State::GetAccountName(package_handler(get_account_name)))
         .await?;
diff --git a/src/handlers/commands/set_master_pass.rs b/src/handlers/commands/set_master_pass.rs
index 9ef0daf..3bc501c 100644
--- a/src/handlers/commands/set_master_pass.rs
+++ b/src/handlers/commands/set_master_pass.rs
@@ -2,7 +2,7 @@ use crate::{
     entity::{master_pass, prelude::*},
     handlers::{utils::package_handler, MainDialogue, State},
 };
-use sea_orm::{prelude::*, QuerySelect};
+use sea_orm::prelude::*;
 use teloxide::{adaptors::Throttle, prelude::*};
 use tokio::task;
 
@@ -15,19 +15,6 @@ async fn get_master_pass(
 ) -> crate::Result<()> {
     dialogue.exit().await?;
     let user_id = msg.from().unwrap().id.0;
-    let exists = MasterPass::find()
-        .select_only()
-        .column(master_pass::Column::UserId)
-        .filter(master_pass::Column::UserId.eq(user_id))
-        .into_tuple::<u64>()
-        .one(&db)
-        .await?
-        .is_some();
-    if exists {
-        bot.send_message(msg.chat.id, "Master password already exists")
-            .await?;
-        return Ok(());
-    }
     let model = task::spawn_blocking(move || {
         master_pass::ActiveModel::from_unencrypted(user_id, &master_password)
     })
@@ -41,7 +28,14 @@ pub async fn set_master_pass(
     bot: Throttle<Bot>,
     msg: Message,
     dialogue: MainDialogue,
+    db: DatabaseConnection,
 ) -> crate::Result<()> {
+    let user_id = msg.from().unwrap().id.0;
+    if MasterPass::exists(user_id, &db).await? {
+        bot.send_message(msg.chat.id, "Master password already exists")
+            .await?;
+        return Ok(());
+    }
     bot.send_message(msg.chat.id, "Send new master password")
         .await?;
     dialogue