Added checks of the master password for /add_account and /get_account

src/entity/account.rs

@@ -86,19 +86,14 @@ impl Model {
     }
 }
 
-#[derive(Copy, Clone, EnumIter, DeriveColumn, Debug)]
-enum GetNamesQuery {
-    AccountName,
-}
-
 impl Entity {
     /// Gets a list of account names of a user
     pub async fn get_names(user_id: u64, db: &DatabaseConnection) -> crate::Result<Vec<String>> {
         Self::find()
             .select_only()
-            .column_as(Column::Name, GetNamesQuery::AccountName)
+            .column(Column::Name)
             .filter(Column::UserId.eq(user_id))
-            .into_values::<_, GetNamesQuery>()
+            .into_tuple()
             .all(db)
             .await
             .map_err(|err| err.into())

src/entity/master_pass.rs

@@ -20,22 +20,38 @@ pub enum Relation {}
 
 impl ActiveModelBehavior for ActiveModel {}
 
+fn hash_password(password: &[u8], salt: &[u8]) -> crate::Result<Vec<u8>> {
+    let params = Params::new(14, Params::RECOMMENDED_R, Params::RECOMMENDED_P, 64)?;
+    let mut password_hash = vec![0; 64];
+    scrypt(password.as_ref(), &salt, &params, &mut password_hash)?;
+    Ok(password_hash)
+}
+
 impl ActiveModel {
     pub fn from_unencrypted(user_id: u64, password: &str) -> crate::Result<Self> {
         let mut salt = vec![0; 64];
         OsRng.fill_bytes(&mut salt);
-        let params = Params::new(
-            Params::RECOMMENDED_LOG_N,
-            Params::RECOMMENDED_R,
-            Params::RECOMMENDED_P,
-            64,
-        )?;
-        let mut password_hash = vec![0; 64];
-        scrypt(password.as_ref(), &salt, &params, &mut password_hash)?;
+        let password_hash = Set(hash_password(password.as_ref(), &salt)?);
         Ok(Self {
             user_id: Set(user_id),
             salt: Set(salt),
-            password_hash: Set(password_hash),
+            password_hash,
         })
     }
 }
+
+impl Entity {
+    pub async fn verify_master_pass(
+        user_id: u64,
+        master_pass: &str,
+        db: &DatabaseConnection,
+    ) -> crate::Result<Option<bool>> {
+        let model = match Self::find_by_id(user_id).one(db).await {
+            Ok(Some(model)) => model,
+            Ok(None) => return Ok(None),
+            Err(err) => return Err(err.into()),
+        };
+        let password_hash = hash_password(master_pass.as_ref(), &model.salt)?;
+        Ok(Some(password_hash == model.password_hash))
+    }
+}

src/handlers/add_account.rs

@@ -1,8 +1,7 @@
+use crate::{entity::account, utils::handle_master_password_check};
 use sea_orm::prelude::*;
 use teloxide::{adaptors::Throttle, prelude::*};
 
-use crate::entity::account;
-
 pub async fn add_account(
     bot: Throttle<Bot>,
     msg: Message,
@@ -10,6 +9,9 @@ pub async fn add_account(
     (name, login, password, master_pass): (String, String, String, String),
 ) -> crate::Result<()> {
     let user_id = msg.from().unwrap().id.0;
+    if handle_master_password_check(&bot, &db, msg.chat.id, user_id, &master_pass).await? {
+        return Ok(());
+    };
     let account =
         account::ActiveModel::from_unencrypted(user_id, name, &login, &password, &master_pass)?;
     account.insert(&db).await?;

src/handlers/get_account.rs

@@ -1,4 +1,7 @@
-use crate::entity::{account, prelude::Account};
+use crate::{
+    entity::{account, prelude::Account},
+    utils::handle_master_password_check,
+};
 use sea_orm::prelude::*;
 use teloxide::{adaptors::Throttle, prelude::*, types::ParseMode};
 
@@ -8,10 +11,14 @@ pub async fn get_account(
     db: DatabaseConnection,
     (name, master_pass): (String, String),
 ) -> crate::Result<()> {
+    let user_id = msg.from().unwrap().id.0;
+    if handle_master_password_check(&bot, &db, msg.chat.id, user_id, &master_pass).await? {
+        return Ok(());
+    };
     let account = Account::find()
         .filter(
             account::Column::UserId
-                .eq(msg.from().unwrap().id.0)
+                .eq(user_id)
                 .add(account::Column::Name.eq(&name)),
         )
         .one(&db)

src/handlers/get_accounts.rs

@@ -25,7 +25,7 @@ pub async fn get_accounts(
         result.reserve(name.len() + 3);
         result.push_str("\n`");
         result.push_str(&name);
-        result.push('\'')
+        result.push('`')
     }
     bot.send_message(msg.chat.id, result)
         .parse_mode(ParseMode::MarkdownV2)

src/handlers/set_master_pass.rs

@@ -10,10 +10,11 @@ pub async fn set_master_pass(
     master_pass: String,
 ) -> crate::Result<()> {
     let user_id = msg.from().unwrap().id.0;
-    println!("User id: {user_id}");
     let exists = MasterPass::find()
+        .select_only()
+        .column(master_pass::Column::UserId)
         .filter(master_pass::Column::UserId.eq(user_id))
-        .limit(1)
+        .into_tuple::<u64>()
         .one(&db)
         .await?
         .is_some();

src/main.rs

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

src/utils.rs

@@ -0,0 +1,34 @@
+use crate::entity::prelude::*;
+use sea_orm::DatabaseConnection;
+use teloxide::{adaptors::Throttle, prelude::*};
+
+/// Handles checking the master password
+///
+/// # Returns
+///
+/// Returns Ok(true) if the master password wasn't right or if it wasn't set
+///
+/// # Errors
+///
+/// Returns an error if there was an error getting or hashing the master password
+#[inline]
+pub async fn handle_master_password_check(
+    bot: &Throttle<Bot>,
+    db: &DatabaseConnection,
+    chat_id: ChatId,
+    user_id: u64,
+    master_pass: &str,
+) -> crate::Result<bool> {
+    match MasterPass::verify_master_pass(user_id, &master_pass, db).await {
+        Ok(Some(true)) => Ok(false),
+        Ok(Some(false)) => {
+            bot.send_message(chat_id, "Wrong master password").await?;
+            Ok(true)
+        }
+        Ok(None) => {
+            bot.send_message(chat_id, "No master password set").await?;
+            Ok(true)
+        }
+        Err(err) => Err(err.into()),
+    }
+}