src/account_parsing.py

@@ -35,7 +35,7 @@ class _Accounts(pydantic.BaseModel):
 def _accounts_list_to_json(accounts: Iterable[DecryptedAccount]) -> str:
     result = _Accounts(
         accounts=[_Account.from_usual_account(i) for i in accounts],
-    ).json(ensure_ascii=False, indent=2)
+    ).json(ensure_ascii=False)
     return result
 
 

src/bot/__init__.py

@@ -3,13 +3,13 @@ import functools
 from sqlalchemy.future import Engine
 from telebot.async_telebot import AsyncTeleBot
 
-from . import callback_handlers, exception_handler, message_handlers
+from . import callback_handlers, message_handlers
 
-__all__ = ["callback_handlers", "exception_handler", "message_handlers"]
+__all__ = ["callback_handlers", "message_handlers"]
 
 
 def create_bot(token: str, engine: Engine) -> AsyncTeleBot:
-    bot = AsyncTeleBot(token, exception_handler=exception_handler.Handler)
+    bot = AsyncTeleBot(token)
     bot.register_message_handler(
         functools.partial(message_handlers.set_master_password, bot, engine),
         commands=["set_master_pass"],

src/bot/exception_handler.py

@@ -1,8 +0,0 @@
-import traceback
-from typing import Type
-
-
-class Handler:
-    @staticmethod
-    def handle(exc: Type[BaseException]) -> None:
-        traceback.print_exception(exc)

src/bot/message_handlers.py

@@ -1,7 +1,6 @@
 import asyncio
 import functools
 import gc
-import itertools
 from concurrent.futures import ProcessPoolExecutor
 
 import telebot
@@ -58,31 +57,21 @@ async def get_accounts(
 
 async def delete_all(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
     await base_handler(bot, mes)
-    master_pass = db.get.get_master_pass(engine, mes.from_user.id)
-    if master_pass is None:
-        await send_tmp_message(bot, mes.chat.id, "У вас нет мастер пароля")
-        return
     bot_mes = await bot.send_message(
         mes.chat.id,
         "Вы действительно хотите удалить все ваши аккаунты? Это действие "
         "нельзя отменить. "
-        "Отправьте мастер пароль для подтверждения",
-    )
-    register_state(
-        mes, functools.partial(_delete_all2, bot, engine, master_pass, bot_mes)
+        "Отправьте YES для подтверждения",
     )
+    register_state(mes, functools.partial(_delete_all2, bot, engine, bot_mes))
 
 
 async def _delete_all2(
-    bot: AsyncTeleBot,
-    engine: Engine,
-    master_pass: db.models.MasterPass,
-    prev_mes: Message,
-    mes: Message,
+    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
     await base_handler(bot, mes, prev_mes)
     text = mes.text.strip()
-    if encryption.master_pass.check_master_pass(text, master_pass):
+    if text == "YES":
         db.delete.purge_accounts(engine, mes.from_user.id)
         db.delete.delete_master_pass(engine, mes.from_user.id)
         await send_tmp_message(
@@ -95,7 +84,7 @@ async def _delete_all2(
         await send_tmp_message(
             bot,
             mes.chat.id,
-            "Вы отправили не верный мастер пароль, ничего не удалено",
+            "Вы отправили не YES, ничего не удалено",
         )
 
 
@@ -144,8 +133,7 @@ async def reset_master_pass(
 ) -> None:
     await base_handler(bot, mes)
 
-    master_pass = db.get.get_master_pass(engine, mes.from_user.id)
-    if master_pass is None:
+    if db.get.get_master_pass(engine, mes.from_user.id) is None:
         return await send_tmp_message(
             bot,
             mes.chat.id,
@@ -154,48 +142,17 @@ async def reset_master_pass(
 
     bot_mes = await bot.send_message(
         mes.chat.id,
-        "Отправьте текущий мастер пароль",
+        "Отправьте новый мастер пароль, осторожно, все текущие аккаунты "
+        "будут удалены навсегда",
     )
 
     register_state(
         mes,
-        functools.partial(
-            _reset_master_pass2,
-            bot,
-            engine,
-            master_pass,
-            bot_mes,
-        ),
+        functools.partial(_reset_master_pass2, bot, engine, bot_mes),
     )
 
 
 async def _reset_master_pass2(
-    bot: AsyncTeleBot,
-    engine: Engine,
-    master_pass: db.models.MasterPass,
-    prev_mes: Message,
-    mes: Message,
-) -> None:
-    await base_handler(bot, mes, prev_mes)
-    text = mes.text.strip()
-    if text == "/cancel":
-        await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
-
-    if not encryption.master_pass.check_master_pass(text, master_pass):
-        await send_tmp_message(bot, mes.chat.id, "Неверный мастер пароль")
-        return
-
-    bot_mes = await bot.send_message(
-        mes.chat.id,
-        "Отправьте новый мастер пароль. Осторожно, все аккаунты будут удалены",
-    )
-    register_state(
-        mes,
-        functools.partial(_reset_master_pass3, bot, engine, bot_mes),
-    )
-
-
-async def _reset_master_pass3(
     bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
     await base_handler(bot, mes, prev_mes)
@@ -470,14 +427,13 @@ async def delete_account(
 
     register_state(
         mes,
-        functools.partial(_delete_account2, bot, engine, master_pass, bot_mes),
+        functools.partial(_delete_account2, bot, engine, bot_mes),
     )
 
 
 async def _delete_account2(
     bot: AsyncTeleBot,
     engine: Engine,
-    master_pass: db.models.MasterPass,
     prev_mes: Message,
     mes: Message,
 ):
@@ -491,36 +447,27 @@ async def _delete_account2(
 
     bot_mes = await bot.send_message(
         mes.from_user.id,
-        f'Вы уверены, что хотите удалить аккаунт "{text}"?\nОтправьте мастер '
-        "пароль для подтверждения",
+        f'Вы уверены, что хотите удалить аккаунт "{text}"?\nОтправьте YES для '
+        "подтверждения",
     )
 
     register_state(
         mes,
-        functools.partial(
-            _delete_account3,
-            bot,
-            engine,
-            master_pass,
-            bot_mes,
-            text,
-        ),
+        functools.partial(_delete_account3, bot, engine, bot_mes, text),
     )
 
 
 async def _delete_account3(
     bot: AsyncTeleBot,
     engine: Engine,
-    master_pass: db.models.MasterPass,
     prev_mes: Message,
     account_name: str,
     mes: Message,
 ) -> None:
     await base_handler(bot, mes, prev_mes)
     text = mes.text.strip()
-    if not encryption.master_pass.check_master_pass(text, master_pass):
-        await send_tmp_message(bot, mes.chat.id, "Неверный пароль")
-        return
+    if text != "YES":
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
     db.delete.delete_account(engine, mes.from_user.id, account_name)
     await send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
@@ -593,6 +540,7 @@ async def _export2(
             )
             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,
@@ -692,26 +640,14 @@ async def _import3(
     # List of names of accounts, which failed to be added to the database
     # or failed the tests
     failed: list[str] = []
-    tasks: list[asyncio.Future[db.models.Account]] = []
-    loop = asyncio.get_running_loop()
-    with ProcessPoolExecutor() as pool:
     for account in accounts:
         if not check_account(account):
             failed.append(account.name)
             continue
-            function = functools.partial(
-                encryption.accounts.encrypt,
-                account,
-                text,
-            )
-            tasks.append(loop.run_in_executor(pool, function))
-        enc_accounts: list[db.models.Account] = await asyncio.gather(*tasks)
-    results = db.add.add_accounts(engine, enc_accounts)
-
-    failed_accounts = itertools.compress(
-        enc_accounts, (not result for result in results)
-    )
-    failed.extend((account.name for account in failed_accounts))
+        account = encryption.accounts.encrypt(account, text)
+        result = db.add.add_account(engine, account)
+        if not result:
+            failed.append(account.name)
 
     if failed:
         await send_deleteable_message(
@@ -720,7 +656,7 @@ async def _import3(
     else:
         await send_tmp_message(bot, mes.chat.id, "Успех")
 
-    del text, mes, accounts, function, tasks, failed_accounts
+    del text, mes, accounts
     gc.collect()
 
 
@@ -742,20 +678,10 @@ async def message_handler(bot: AsyncTeleBot, mes: Message) -> None:
         await delete_message(bot, mes)
         if mes.text.strip() == "/cancel":
             await send_tmp_message(bot, mes.chat.id, "Нет активного действия")
-            return
         await send_tmp_message(
             bot,
             mes.chat.id,
             "Вы отправили не корректное сообщение",
         )
         return
-
-    try:
     await handler(mes)
-    except Exception:
-        await send_tmp_message(
-            bot,
-            mes.chat.id,
-            "Произошла непредвиденная ошибка",
-        )
-        raise

src/db/add.py

@@ -5,42 +5,27 @@ from sqlalchemy.future import Engine
 from . import models
 
 
-def _add_model(
-    session: sqlmodel.Session, model: models.Account | models.MasterPass
-) -> bool:
-    """Adds model to the session. Returns true on success,
+def add_account(engine: Engine, account: models.Account) -> bool:
+    """Adds account to the database. Returns true on success,
     false otherwise"""
     try:
-        session.add(model)
+        with sqlmodel.Session(engine) as session:
+            session.add(account)
+            session.commit()
     except IntegrityError:
         return False
     else:
         return True
 
 
-def add_account(engine: Engine, account: models.Account) -> bool:
-    """Adds account to the database. Returns true on success,
-    false otherwise"""
-    with sqlmodel.Session(engine) as session:
-        result = _add_model(session, account)
-        session.commit()
-    return result
-
-
 def add_master_pass(engine: Engine, master_pass: models.MasterPass) -> bool:
     """Adds master password the database. Returns true on success,
     false otherwise"""
+    try:
         with sqlmodel.Session(engine) as session:
-        result = _add_model(session, master_pass)
+            session.add(master_pass)
             session.commit()
-    return result
-
-
-def add_accounts(
-    engine: Engine,
-    accounts: list[models.Account],
-) -> list[bool]:
-    with sqlmodel.Session(engine) as session:
-        result = [_add_model(session, account) for account in accounts]
-        session.commit()
-    return result
+    except IntegrityError:
+        return False
+    else:
+        return True

src/db/get.py

@@ -36,13 +36,9 @@ def get_account_names(
 
 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(
+    statement = sqlmodel.select(models.Account).where(
         models.Account.user_id == user_id,
     )
-        .order_by(models.Account.name)
-    )
     with sqlmodel.Session(engine) as session:
         result = session.exec(statement).fetchall()
     return result

src/db/models.py

@@ -4,21 +4,13 @@ import sqlmodel
 class MasterPass(sqlmodel.SQLModel, table=True):
     __tablename__ = "master_passwords"
     user_id: int = sqlmodel.Field(
-        sa_column=sqlmodel.Column(
-            sqlmodel.INT(),
-            primary_key=True,
-            autoincrement=False,
-        )
+        sa_column=sqlmodel.Column(sqlmodel.INT(), primary_key=True)
     )
     salt: bytes = sqlmodel.Field(
-        sa_column=sqlmodel.Column(sqlmodel.BINARY(64), nullable=False),
-        max_length=64,
-        min_length=64,
+        sa_column=sqlmodel.Column(sqlmodel.BINARY(64), nullable=False)
     )
     password_hash: bytes = sqlmodel.Field(
-        sa_column=sqlmodel.Column(sqlmodel.BINARY(128), nullable=False),
-        max_length=128,
-        min_length=128,
+        sa_column=sqlmodel.Column(sqlmodel.BINARY(128), nullable=False)
     )
 
 
@@ -26,17 +18,13 @@ class Account(sqlmodel.SQLModel, table=True):
     __tablename__ = "accounts"
     __table_args__ = (sqlmodel.PrimaryKeyConstraint("user_id", "name"),)
     user_id: int = sqlmodel.Field()
-    name: str = sqlmodel.Field(max_length=256)
+    name: str = sqlmodel.Field(max_length=255)
     salt: bytes = sqlmodel.Field(
-        sa_column=sqlmodel.Column(sqlmodel.BINARY(64), nullable=False),
-        max_length=64,
-        min_length=64,
+        sa_column=sqlmodel.Column(sqlmodel.BINARY(64), nullable=False)
     )
     enc_login: bytes = sqlmodel.Field(
-        sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False),
-        max_length=256,
+        sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False)
     )
     enc_password: bytes = sqlmodel.Field(
-        sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False),
-        max_length=256,
+        sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False)
     )

src/encryption/accounts.py

@@ -1,43 +1,26 @@
+import base64
 import os
-from typing import Self
 
+from cryptography.fernet import Fernet
+from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
 from ..db.models import Account
 from ..decrypted_account import DecryptedAccount
 
 
-class Cipher:
-    def __init__(self, key: bytes) -> None:
-        self._chacha = ChaCha20Poly1305(key)
-
-    @classmethod
-    def generate_cipher(cls, salt: bytes, password: bytes) -> Self:
-        """Generates cipher which uses key derived from a given password"""
+def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
+    """Generates key for fernet encryption"""
     kdf = PBKDF2HMAC(
         algorithm=hashes.SHA256(),
         length=32,
         salt=salt,
         iterations=100000,
+        backend=default_backend(),
     )
-        return cls(kdf.derive(password))
-
-    def encrypt(self, data: bytes) -> bytes:
-        nonce = os.urandom(12)
-        return nonce + self._chacha.encrypt(
-            nonce,
-            data,
-            associated_data=None,
-        )
-
-    def decrypt(self, data: bytes) -> bytes:
-        return self._chacha.decrypt(
-            nonce=data[:12],
-            data=data[12:],
-            associated_data=None,
-        )
+    key = base64.urlsafe_b64encode(kdf.derive(master_pass))
+    return key
 
 
 def encrypt(
@@ -46,10 +29,15 @@ def encrypt(
 ) -> Account:
     """Encrypts account using master password and returns Account object"""
     salt = os.urandom(64)
-    cipher = Cipher.generate_cipher(salt, master_pass.encode("utf-8"))
+    key = _generate_key(salt, master_pass.encode("utf-8"))
+    f = Fernet(key)
 
-    enc_login = cipher.encrypt(account.login.encode("utf-8"))
-    enc_password = cipher.encrypt(account.password.encode("utf-8"))
+    enc_login = base64.urlsafe_b64decode(
+        f.encrypt(account.login.encode("utf-8")),
+    )
+    enc_password = base64.urlsafe_b64decode(
+        f.encrypt(account.password.encode("utf-8")),
+    )
 
     return Account(
         user_id=account.user_id,
@@ -66,10 +54,15 @@ def decrypt(
 ) -> DecryptedAccount:
     """Decrypts account using master password and returns
     DecryptedAccount object"""
-    cipher = Cipher.generate_cipher(account.salt, master_pass.encode("utf-8"))
+    key = _generate_key(account.salt, master_pass.encode("utf-8"))
+    f = Fernet(key)
 
-    login = cipher.decrypt(account.enc_login).decode("utf-8")
-    password = cipher.decrypt(account.enc_password).decode("utf-8")
+    login = f.decrypt(
+        base64.urlsafe_b64encode(account.enc_login),
+    ).decode("utf-8")
+    password = f.decrypt(
+        base64.urlsafe_b64encode(account.enc_password),
+    ).decode("utf-8")
 
     return DecryptedAccount(
         user_id=account.user_id,