Extended functionality

src/entity/account.rs

@@ -1,6 +1,10 @@
 //! `SeaORM` Entity. Generated by sea-orm-codegen 0.11.2
 
-use sea_orm::entity::prelude::*;
+use chacha20poly1305::{aead::Aead, AeadCore, ChaCha20Poly1305, KeyInit};
+use pbkdf2::pbkdf2_hmac_array;
+use rand::{rngs::OsRng, RngCore};
+use sea_orm::{prelude::*, ActiveValue::Set, QuerySelect};
+use sha2::Sha256;
 
 #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
 #[sea_orm(table_name = "account")]
@@ -21,3 +25,82 @@ pub struct Model {
 pub enum Relation {}
 
 impl ActiveModelBehavior for ActiveModel {}
+
+struct Cipher {
+    chacha: ChaCha20Poly1305,
+}
+
+impl Cipher {
+    fn new(password: &[u8], salt: &[u8]) -> Self {
+        let key = pbkdf2_hmac_array::<Sha256, 32>(password, salt, 480000);
+
+        Self {
+            chacha: ChaCha20Poly1305::new(&key.into()),
+        }
+    }
+
+    pub fn encrypt(&self, value: &[u8]) -> crate::Result<Vec<u8>> {
+        let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
+        let mut result = self.chacha.encrypt(&nonce, value).unwrap();
+        result.extend(nonce);
+        Ok(result)
+    }
+
+    fn decrypt(&self, value: &[u8]) -> crate::Result<Vec<u8>> {
+        let (data, nonce) = value.split_at(value.len() - 12);
+        self.chacha
+            .decrypt(nonce.into(), data)
+            .map_err(|err| err.into())
+    }
+}
+
+impl ActiveModel {
+    pub fn from_unencrypted(
+        user_id: u64,
+        name: String,
+        login: &str,
+        password: &str,
+        master_pass: &str,
+    ) -> crate::Result<Self> {
+        let mut salt = vec![0; 64];
+        OsRng.fill_bytes(&mut salt);
+        let cipher = Cipher::new(master_pass.as_ref(), &salt);
+        let enc_login = Set(cipher.encrypt(login.as_ref())?);
+        let enc_password = Set(cipher.encrypt(password.as_ref())?);
+        Ok(Self {
+            name: Set(name),
+            user_id: Set(user_id),
+            salt: Set(salt),
+            enc_login,
+            enc_password,
+        })
+    }
+}
+
+impl Model {
+    pub fn decrypt(&self, master_pass: &str) -> crate::Result<(String, String)> {
+        let cipher = Cipher::new(master_pass.as_ref(), self.salt.as_ref());
+        let login = String::from_utf8(cipher.decrypt(self.enc_login.as_ref())?)?;
+        let password = String::from_utf8(cipher.decrypt(self.enc_password.as_ref())?)?;
+        Ok((login, password))
+    }
+}
+
+#[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)
+            .filter(Column::UserId.eq(user_id))
+            .into_values::<_, GetNamesQuery>()
+            .all(db)
+            .await
+            .map_err(|err| err.into())
+    }
+}

src/entity/master_pass.rs

@@ -1,6 +1,8 @@
 //! `SeaORM` Entity. Generated by sea-orm-codegen 0.11.2
 
-use sea_orm::entity::prelude::*;
+use rand::{rngs::OsRng, RngCore};
+use scrypt::{scrypt, Params};
+use sea_orm::{entity::prelude::*, ActiveValue::Set};
 
 #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
 #[sea_orm(table_name = "master_pass")]
@@ -17,3 +19,23 @@ pub struct Model {
 pub enum Relation {}
 
 impl ActiveModelBehavior for ActiveModel {}
+
+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,
+            128,
+        )?;
+        let mut password_hash = vec![0; 128];
+        scrypt(password.as_ref(), &salt, &params, &mut password_hash)?;
+        Ok(Self {
+            user_id: Set(user_id),
+            salt: Set(salt),
+            password_hash: Set(password_hash),
+        })
+    }
+}

src/handlers/add_account.rs

@@ -0,0 +1,18 @@
+use sea_orm::prelude::*;
+use teloxide::{adaptors::Throttle, prelude::*};
+
+use crate::entity::account;
+
+pub async fn add_account(
+    bot: Throttle<Bot>,
+    msg: Message,
+    db: DatabaseConnection,
+    (name, login, password, master_pass): (String, String, String, String),
+) -> crate::Result<()> {
+    let user_id = msg.from().unwrap().id.0;
+    let account =
+        account::ActiveModel::from_unencrypted(user_id, name, &login, &password, &master_pass)?;
+    account.insert(&db).await?;
+    bot.send_message(msg.chat.id, "Success").await?;
+    Ok(())
+}

src/handlers/default.rs

@@ -0,0 +1,6 @@
+use teloxide::{adaptors::Throttle, prelude::*};
+
+pub async fn default(bot: Throttle<Bot>, msg: Message) -> crate::Result<()> {
+    bot.send_message(msg.chat.id, "Unknown command").await?;
+    Ok(())
+}

src/handlers/get_account.rs

@@ -0,0 +1,26 @@
+use crate::entity::{account, prelude::Account};
+use sea_orm::prelude::*;
+use teloxide::{adaptors::Throttle, prelude::*, types::ParseMode};
+
+pub async fn get_account(
+    bot: Throttle<Bot>,
+    msg: Message,
+    db: DatabaseConnection,
+    (name, master_pass): (String, String),
+) -> crate::Result<()> {
+    let account = Account::find()
+        .filter(
+            account::Column::UserId
+                .eq(msg.from().unwrap().id.0)
+                .add(account::Column::Name.eq(&name)),
+        )
+        .one(&db)
+        .await?
+        .unwrap();
+    let (login, password) = account.decrypt(&master_pass)?;
+    let message = format!("Name:\n`{name}`\nLogin:\n`{login}`\nPassword:\n`{password}`");
+    bot.send_message(msg.chat.id, message)
+        .parse_mode(ParseMode::MarkdownV2)
+        .await?;
+    Ok(())
+}

src/handlers/get_accounts.rs

@@ -0,0 +1,34 @@
+use crate::entity::prelude::Account;
+use sea_orm::prelude::*;
+use teloxide::{adaptors::Throttle, prelude::*, types::ParseMode};
+
+#[derive(Clone, Copy, EnumIter, DeriveColumn, Debug)]
+enum Query {
+    AccountName,
+}
+
+pub async fn get_accounts(
+    bot: Throttle<Bot>,
+    msg: Message,
+    db: DatabaseConnection,
+) -> crate::Result<()> {
+    let user_id = msg.from().unwrap().id.0;
+    let mut account_names = Account::get_names(user_id, &db).await?.into_iter();
+    let mut result = match account_names.next() {
+        Some(name) => format!("Accounts:\n`{name}`"),
+        None => {
+            bot.send_message(msg.chat.id, "No accounts found").await?;
+            return Ok(());
+        }
+    };
+    for name in account_names {
+        result.reserve(name.len() + 3);
+        result.push_str("\n`");
+        result.push_str(&name);
+        result.push('\'')
+    }
+    bot.send_message(msg.chat.id, result)
+        .parse_mode(ParseMode::MarkdownV2)
+        .await?;
+    Ok(())
+}

src/handlers/help.rs

@@ -0,0 +1,7 @@
+use teloxide::{adaptors::Throttle, prelude::*, utils::command::BotCommands};
+
+pub async fn help(bot: Throttle<Bot>, msg: Message) -> crate::Result<()> {
+    bot.send_message(msg.chat.id, super::Command::descriptions().to_string())
+        .await?;
+    Ok(())
+}

src/handlers/mod.rs

@@ -0,0 +1,56 @@
+use sea_orm::prelude::*;
+use teloxide::{
+    adaptors::{throttle::Limits, Throttle},
+    filter_command,
+    prelude::*,
+    utils::command::BotCommands,
+};
+
+mod add_account;
+mod default;
+mod get_account;
+mod get_accounts;
+mod help;
+
+#[derive(BotCommands, Clone)]
+#[command(rename_rule = "snake_case")]
+enum Command {
+    #[command()]
+    Help,
+    #[command(parse_with = "split")]
+    AddAccount(String, String, String, String),
+    #[command(parse_with = "split")]
+    GetAccount(String, String),
+    #[command()]
+    GetAccounts,
+}
+
+pub fn get_dispatcher(
+    token: String,
+    db: DatabaseConnection,
+) -> Dispatcher<Throttle<Bot>, crate::Error, teloxide::dispatching::DefaultKey> {
+    let bot = Bot::new(token).throttle(Limits::default());
+    Dispatcher::builder(
+        bot,
+        Update::filter_message()
+            .branch(
+                filter_command::<Command, _>()
+                    .branch(dptree::case![Command::Help].endpoint(help::help))
+                    .branch(
+                        dptree::case![Command::AddAccount(name, login, password, master_pass)]
+                            .endpoint(add_account::add_account),
+                    )
+                    .branch(
+                        dptree::case![Command::GetAccount(name, master_pass)]
+                            .endpoint(get_account::get_account),
+                    )
+                    .branch(
+                        dptree::case![Command::GetAccounts].endpoint(get_accounts::get_accounts),
+                    ),
+            )
+            .branch(dptree::endpoint(default::default)),
+    )
+    .dependencies(dptree::deps![db])
+    .enable_ctrlc_handler()
+    .build()
+}

src/main.rs

@@ -1,126 +1,22 @@
 mod entity;
+mod handlers;
 
-use entity::{account, prelude::Account};
-
-use anyhow::Result;
-use chacha20poly1305::{
-    aead::{rand_core::RngCore, Aead, AeadCore, OsRng},
-    ChaCha20Poly1305, KeyInit,
-};
+use anyhow::{Error, Result};
 use dotenv::dotenv;
-use pbkdf2::pbkdf2_hmac_array;
-use sea_orm::{
-    ActiveModelTrait, ActiveValue::Set, ColumnTrait, ConnectOptions, Database, EntityTrait,
-    QueryFilter,
-};
-use sha2::Sha256;
+use handlers::get_dispatcher;
+use migration::{Migrator, MigratorTrait};
+use sea_orm::Database;
 use std::env;
-use teloxide::{prelude::*, utils::command::BotCommands};
-
-#[derive(BotCommands, Clone)]
-#[command(rename_rule = "lowercase")]
-enum Command {
-    #[command()]
-    Help,
-    #[command(parse_with = "split")]
-    AddAccount(String, String, String, String),
-    #[command(parse_with = "split")]
-    GetAccount(String, String),
-}
-
-use Command::*;
-
-struct Cipher {
-    chacha: ChaCha20Poly1305,
-}
-
-impl Cipher {
-    fn new(password: &[u8], salt: &[u8]) -> Self {
-        let key = pbkdf2_hmac_array::<Sha256, 32>(password, salt, 480000);
-
-        Self {
-            chacha: ChaCha20Poly1305::new(&key.into()),
-        }
-    }
-
-    fn encrypt(&self, value: &[u8]) -> Result<Vec<u8>> {
-        let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
-        let mut result = self.chacha.encrypt(&nonce, value).unwrap();
-        result.extend(nonce);
-        Ok(result)
-    }
-
-    fn decrypt(&self, value: &[u8]) -> Result<Vec<u8>> {
-        let (data, nonce) = value.split_at(value.len() - 12);
-        assert!(nonce.len() == 12);
-        self.chacha
-            .decrypt(nonce.into(), data)
-            .map_err(|err| err.into())
-    }
-}
-
-async fn answer_command(db: ConnectOptions, bot: Bot, msg: Message, cmd: Command) -> Result<()> {
-    let user_id = msg.from().unwrap().id.0;
-    let db = Database::connect(db).await?;
-    match cmd {
-        Help => {
-            bot.send_message(msg.chat.id, Command::descriptions().to_string())
-                .await?;
-        }
-        AddAccount(name, login, password, master_pass) => {
-            let mut salt = vec![0; 64];
-            OsRng.fill_bytes(&mut salt);
-            let cipher = Cipher::new(master_pass.as_ref(), &salt);
-            let enc_login = Set(cipher.encrypt(login.as_ref())?);
-            let enc_password = Set(cipher.encrypt(password.as_ref())?);
-            let account = account::ActiveModel {
-                name: Set(name),
-                user_id: Set(user_id),
-                salt: Set(salt),
-                enc_login,
-                enc_password,
-            };
-            account.insert(&db).await?;
-            bot.send_message(msg.chat.id, "Success").await?;
-        }
-        GetAccount(name, master_pass) => {
-            let account = Account::find()
-                .filter(
-                    account::Column::UserId
-                        .eq(user_id)
-                        .add(account::Column::Name.eq(&name)),
-                )
-                .one(&db)
-                .await?
-                .unwrap();
-            let cipher = Cipher::new(master_pass.as_ref(), &account.salt);
-            let login = String::from_utf8(cipher.decrypt(&account.enc_login)?)?;
-            let password = String::from_utf8(cipher.decrypt(&account.enc_password)?)?;
-            let message = format!("Account `{name}`\nLogin: `{login}`\nPassword: `{password}`");
-            bot.send_message(msg.chat.id, message).await?;
-        }
-    }
-    Ok(())
-}
 
 #[tokio::main]
-async fn main() {
+async fn main() -> Result<()> {
     let _ = dotenv();
     pretty_env_logger::init();
 
     let token = env::var("TOKEN").unwrap();
     let database_url = env::var("DATABASE_URL").unwrap();
-    let bot = Bot::new(token);
-    let db = ConnectOptions::new(database_url);
-    Dispatcher::builder(
-        bot,
-        Update::filter_message()
-            .filter_command::<Command>()
-            .endpoint(answer_command),
-    )
-    .dependencies(dptree::deps![db])
-    .enable_ctrlc_handler()
-    .build()
-    .dispatch()
-    .await;
+    let db = Database::connect(database_url).await?;
+    Migrator::up(&db, None).await?;
+    get_dispatcher(token, db).dispatch().await;
+    Ok(())
 }