Fixed sending a final message in /import

Removed checking if sleep_time is 0
Moved sleeping outside of try block

src/__init__.py

@@ -7,7 +7,7 @@ from . import (
     account_checks,
     account_parsing,
     bot,
-    database,
+    db,
     decrypted_account,
     encryption,
     generate_password,
@@ -19,19 +19,19 @@ __all__ = [
     "bot",
     "decrypted_account",
     "encryption",
-    "database",
+    "db",
     "generate_password",
 ]
 
 
 def main() -> None:
     load_dotenv("./.env")
-    engine = database.prepare.get_engine(
+    engine = db.prepare.get_engine(
         host=os.getenv("DB_HOST"),
         user=os.getenv("DB_USER"),
         passwd=os.getenv("DB_PASS"),
         db=os.getenv("DB_NAME"),
     )
-    database.prepare.prepare(engine)
+    db.prepare.prepare(engine)
     bot_ = bot.create_bot(os.getenv("TG_TOKEN"), engine)
     asyncio.run(bot_.infinity_polling())

src/bot/helper_functions.py

@@ -18,7 +18,7 @@ states: dict[tuple[int, int], Handler] = {}
 
 def register_state(
     message: Message,
-    handler: Callable[[Message], Any],
+    handler: Handler,
 ) -> None:
     states[(message.chat.id, message.from_user.id)] = handler
 
@@ -40,9 +40,8 @@ async def delete_message(
     *,
     sleep_time: int = 0,
 ) -> bool:
+    await asyncio.sleep(sleep_time)
     try:
-        if sleep_time != 0:
-            await asyncio.sleep(sleep_time)
         await bot.delete_message(mes.chat.id, mes.id)
     except telebot.apihelper.ApiException:
         return False
@@ -86,10 +85,3 @@ async def send_deleteable_message(
         parse_mode="MarkdownV2",
         reply_markup=markup,
     )
-
-
-def escape(text: str) -> str:
-    escaped_chars = "*_~|`[("
-    for char in escaped_chars:
-        text = text.replace(char, rf"\\{char}")
-    return text

src/bot/message_handlers.py

@@ -1,11 +1,13 @@
+import asyncio
 import functools
 import gc
+from concurrent.futures import ProcessPoolExecutor
 
 import telebot
 from sqlalchemy.future import Engine
 from telebot.async_telebot import AsyncTeleBot
 
-from .. import database, encryption, generate_password
+from .. import db, encryption, generate_password
 from ..account_checks import (
     check_account,
     check_account_name,
@@ -18,7 +20,6 @@ from . import markups
 from .helper_functions import (
     base_handler,
     delete_message,
-    escape,
     get_state,
     register_state,
     send_deleteable_message,
@@ -34,7 +35,7 @@ async def get_accounts(
     mes: Message,
 ) -> None:
     await base_handler(bot, mes)
-    accounts = database.get.get_accounts(
+    accounts = db.get.get_account_names(
         engine,
         mes.from_user.id,
         to_sort=True,
@@ -71,8 +72,8 @@ async def _delete_all2(
     await base_handler(bot, mes, prev_mes)
     text = mes.text.strip()
     if text == "YES":
-        database.delete.purge_accounts(engine, mes.from_user.id)
-        database.delete.delete_master_pass(engine, mes.from_user.id)
+        db.delete.purge_accounts(engine, mes.from_user.id)
+        db.delete.delete_master_pass(engine, mes.from_user.id)
         await send_tmp_message(
             bot,
             mes.chat.id,
@@ -93,7 +94,7 @@ async def set_master_password(
     mes: Message,
 ) -> None:
     await base_handler(bot, mes, None)
-    if database.get.get_master_pass(engine, mes.from_user.id) is not None:
+    if db.get.get_master_pass(engine, mes.from_user.id) is not None:
         return await send_tmp_message(
             bot,
             mes.chat.id,
@@ -118,7 +119,7 @@ async def _set_master_pass2(
         mes.from_user.id,
         text,
     )
-    database.add.add_master_pass(engine, master_password)
+    db.add.add_master_pass(engine, master_password)
 
     await send_tmp_message(bot, mes.chat.id, "Успех")
     del mes, text
@@ -132,7 +133,7 @@ async def reset_master_pass(
 ) -> None:
     await base_handler(bot, mes)
 
-    if database.get.get_master_pass(engine, mes.from_user.id) is None:
+    if db.get.get_master_pass(engine, mes.from_user.id) is None:
         return await send_tmp_message(
             bot,
             mes.chat.id,
@@ -163,8 +164,8 @@ async def _reset_master_pass2(
         mes.from_user.id,
         text,
     )
-    database.delete.purge_accounts(engine, mes.from_user.id)
-    database.change.change_master_pass(engine, master_password)
+    db.delete.purge_accounts(engine, mes.from_user.id)
+    db.change.change_master_pass(engine, master_password)
 
     await send_tmp_message(
         bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
@@ -176,7 +177,7 @@ async def _reset_master_pass2(
 async def add_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
     await base_handler(bot, mes)
 
-    master_password_from_db = database.get.get_master_pass(
+    master_password_from_db = db.get.get_master_pass(
         engine,
         mes.from_user.id,
     )
@@ -208,7 +209,7 @@ async def _add_account2(
             mes.chat.id,
             "Не корректное название аккаунта",
         )
-    if text in database.get.get_accounts(engine, mes.from_user.id):
+    if text in db.get.get_account_names(engine, mes.from_user.id):
         return await send_tmp_message(
             bot, mes.chat.id, "Аккаунт с таким именем уже существует"
         )
@@ -289,7 +290,7 @@ async def _add_account5(
     if text == "/cancel":
         return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    master_password = database.get.get_master_pass(engine, mes.from_user.id)
+    master_password = db.get.get_master_pass(engine, mes.from_user.id)
     if not encryption.master_pass.check_master_pass(text, master_password):
         return await send_tmp_message(
             bot,
@@ -305,12 +306,9 @@ async def _add_account5(
         password=data["passwd"],
     )
 
-    encrypted_account = encryption.other_accounts.encrypt(
-        account,
-        text,
-    )
+    encrypted_account = encryption.accounts.encrypt(account, text)
 
-    result = database.add.add_account(
+    result = db.add.add_account(
         engine,
         encrypted_account,
     )
@@ -329,11 +327,11 @@ async def _add_account5(
 async def get_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
     await base_handler(bot, mes)
 
-    master_pass = database.get.get_master_pass(engine, mes.from_user.id)
+    master_pass = db.get.get_master_pass(engine, mes.from_user.id)
     if master_pass is None:
         return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
 
-    accounts = database.get.get_accounts(
+    accounts = db.get.get_account_names(
         engine,
         mes.from_user.id,
         to_sort=True,
@@ -356,7 +354,7 @@ async def _get_account2(
     if text == "/cancel":
         return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    if text not in database.get.get_accounts(engine, mes.from_user.id):
+    if text not in db.get.get_account_names(engine, mes.from_user.id):
         return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
 
     bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
@@ -378,7 +376,7 @@ async def _get_account3(
     if text == "/cancel":
         return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    master_password = database.get.get_master_pass(
+    master_password = db.get.get_master_pass(
         engine,
         mes.from_user.id,
     )
@@ -390,17 +388,14 @@ async def _get_account3(
             "Не подходит мастер пароль",
         )
 
-    account = database.get.get_account_info(engine, mes.from_user.id, name)
-    account = encryption.other_accounts.decrypt(
-        account,
-        text,
-    )
+    account = db.get.get_account_info(engine, mes.from_user.id, name)
+    account = encryption.accounts.decrypt(account, text)
     await send_deleteable_message(
         bot,
         mes.chat.id,
-        f"Название:\n{escape(account.name)}\n"
+        f"Название:\n`{account.name}`\n"
         f"Логин:\n`{account.login}`\nПароль:\n`{account.password}`\nНажмите "
-        "на логин или пароль, чтобы скопировать",
+        "на название, логин или пароль, чтобы скопировать",
     )
 
     del text, mes
@@ -414,11 +409,11 @@ async def delete_account(
 ) -> None:
     await base_handler(bot, mes)
 
-    master_pass = database.get.get_master_pass(engine, mes.from_user.id)
+    master_pass = db.get.get_master_pass(engine, mes.from_user.id)
     if master_pass is None:
         return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
 
-    accounts = database.get.get_accounts(
+    accounts = db.get.get_account_names(
         engine,
         mes.from_user.id,
         to_sort=True,
@@ -447,7 +442,7 @@ async def _delete_account2(
     if text == "/cancel":
         return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    if text not in database.get.get_accounts(engine, mes.from_user.id):
+    if text not in db.get.get_account_names(engine, mes.from_user.id):
         return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
 
     bot_mes = await bot.send_message(
@@ -474,7 +469,7 @@ async def _delete_account3(
     if text != "YES":
         return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    database.delete.delete_account(engine, mes.from_user.id, account_name)
+    db.delete.delete_account(engine, mes.from_user.id, account_name)
     await send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
 
 
@@ -498,7 +493,7 @@ async def help_command(bot: AsyncTeleBot, mes: Message) -> None:
 
 async def export(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
     await base_handler(bot, mes)
-    master_password_from_db = database.get.get_master_pass(
+    master_password_from_db = db.get.get_master_pass(
         engine,
         mes.from_user.id,
     )
@@ -506,7 +501,7 @@ async def export(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
     if master_password_from_db is None:
         return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
 
-    if not database.get.get_accounts(engine, mes.from_user.id):
+    if not db.get.get_account_names(engine, mes.from_user.id):
         return await send_tmp_message(bot, mes.chat.id, "Нет аккаунтов")
 
     bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
@@ -522,7 +517,7 @@ async def _export2(
     if text == "/cancel":
         return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    master_password = database.get.get_master_pass(
+    master_password = db.get.get_master_pass(
         engine,
         mes.from_user.id,
     )
@@ -533,8 +528,19 @@ async def _export2(
             "Не подходит мастер пароль",
         )
 
-    accounts = database.get.get_all_accounts(engine, mes.from_user.id)
-    accounts = encryption.other_accounts.decrypt_multiple(accounts, text)
+    accounts = db.get.get_accounts(engine, mes.from_user.id)
+    with ProcessPoolExecutor() as pool:
+        loop = asyncio.get_running_loop()
+        tasks = []
+        for account in accounts:
+            function = functools.partial(
+                encryption.accounts.decrypt,
+                account,
+                text,
+            )
+            tasks.append(loop.run_in_executor(pool, function))
+        accounts = await asyncio.gather(*tasks)
+        accounts.sort(key=lambda account: account.name)
     json_io = accounts_to_json(accounts)
     await bot.send_document(
         mes.chat.id,
@@ -552,7 +558,7 @@ async def import_accounts(
     mes: Message,
 ) -> None:
     await base_handler(bot, mes)
-    master_password_from_db = database.get.get_master_pass(
+    master_password_from_db = db.get.get_master_pass(
         engine,
         mes.from_user.id,
     )
@@ -620,7 +626,7 @@ async def _import3(
     if text == "/cancel":
         return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    master_password = database.get.get_master_pass(
+    master_password = db.get.get_master_pass(
         engine,
         mes.from_user.id,
     )
@@ -638,20 +644,18 @@ async def _import3(
         if not check_account(account):
             failed.append(account.name)
             continue
-        account = encryption.other_accounts.encrypt(
-            account,
-            text,
-        )
-        result = database.add.add_account(engine, account)
+        account = encryption.accounts.encrypt(account, text)
+        result = db.add.add_account(engine, account)
         if not result:
             failed.append(account.name)
 
     if failed:
-        mes_text = "Не удалось добавить:\n" + "\n".join(failed)
+        await send_deleteable_message(
+            bot, mes.chat.id, "Не удалось добавить:\n" + "\n".join(failed)
+        )
     else:
-        mes_text = "Успех"
+        await send_tmp_message(bot, mes.chat.id, "Успех")
 
-    await send_tmp_message(bot, mes.chat.id, mes_text, 10)
     del text, mes, accounts
     gc.collect()
 

src/db/get.py

@@ -17,7 +17,7 @@ def get_master_pass(
     return result
 
 
-def get_accounts(
+def get_account_names(
     engine: Engine,
     user_id: int,
     *,
@@ -34,14 +34,10 @@ def get_accounts(
     return result
 
 
-def get_all_accounts(engine: Engine, user_id: int) -> list[models.Account]:
+def get_accounts(engine: Engine, user_id: int) -> list[models.Account]:
     """Returns a list of accounts of a user"""
-    statement = (
-        sqlmodel.select(models.Account)
-        .where(
-            models.Account.user_id == user_id,
-        )
-        .order_by(models.Account.name)
+    statement = sqlmodel.select(models.Account).where(
+        models.Account.user_id == user_id,
     )
     with sqlmodel.Session(engine) as session:
         result = session.exec(statement).fetchall()

src/encryption/__init__.py

@@ -1,3 +1,3 @@
-from . import master_pass, other_accounts
+from . import accounts, master_pass
 
-__all__ = ["master_pass", "other_accounts"]
+__all__ = ["master_pass", "accounts"]

src/encryption/accounts.py

@@ -1,13 +1,12 @@
 import base64
 import os
-from typing import Iterable, Iterator
 
 from cryptography.fernet import Fernet
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
-from ..database.models import Account
+from ..db.models import Account
 from ..decrypted_account import DecryptedAccount
 
 
@@ -71,12 +70,3 @@ def decrypt(
         login=login,
         password=password,
     )
-
-
-def decrypt_multiple(
-    accounts: Iterable[Account], master_pass: str
-) -> Iterator[DecryptedAccount]:
-    """Decrypts an iterable of accounts using master_pass and
-    returns an Iterator of decrypted accounts"""
-    for account in accounts:
-        yield decrypt(account, master_pass)

src/encryption/master_pass.py

@@ -3,7 +3,7 @@ import os
 from cryptography.exceptions import InvalidKey
 from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
 
-from ..database.models import MasterPass
+from ..db.models import MasterPass
 
 MEMORY_USAGE = 2**14