Added ability to send message, which can be deleted by pressing a button

2022-12-29 15:25:20 +03:00
parent 0026e3321a
commit d4c50432d7
5 changed files with 230 additions and 155 deletions
@@ -3,61 +3,62 @@ import functools
 from sqlalchemy.future import Engine
 from telebot.async_telebot import AsyncTeleBot

-from . import handlers
+from . import message_handlers, callback_handlers

-__all__ = ["handlers"]
+__all__ = ["message_handlers"]


 def create_bot(token: str, engine: Engine) -> AsyncTeleBot:
    bot = AsyncTeleBot(token)
    bot.register_message_handler(
-        functools.partial(handlers.set_master_password, bot, engine),
+        functools.partial(message_handlers.set_master_password, bot, engine),
        commands=["set_master_pass"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.get_account, bot, engine),
+        functools.partial(message_handlers.get_account, bot, engine),
        commands=["get_account"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.get_accounts, bot, engine),
+        functools.partial(message_handlers.get_accounts, bot, engine),
        commands=["get_accounts"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.add_account, bot, engine),
+        functools.partial(message_handlers.add_account, bot, engine),
        commands=["add_account"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.delete_all, bot, engine),
+        functools.partial(message_handlers.delete_all, bot, engine),
        commands=["delete_all"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.reset_master_pass, bot, engine),
+        functools.partial(message_handlers.reset_master_pass, bot, engine),
        commands=["reset_master_pass"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.delete_account, bot, engine),
+        functools.partial(message_handlers.delete_account, bot, engine),
        commands=["delete_account"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.help_command, bot),
+        functools.partial(message_handlers.help_command, bot),
        commands=["help", "start"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.cancel, bot), commands=["cancel"]
+        functools.partial(message_handlers.export, bot, engine),
+        commands=["export"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.export, bot, engine), commands=["export"]
-    )
-    bot.register_message_handler(
-        functools.partial(handlers.import_accounts, bot, engine),
+        functools.partial(message_handlers.import_accounts, bot, engine),
        commands=["import"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.gen_password, bot),
+        functools.partial(message_handlers.gen_password, bot),
        commands=["gen_password"],
    )
    bot.register_message_handler(
-        functools.partial(handlers.message_handler, bot),
+        functools.partial(message_handlers.cancel, bot), commands=["cancel"]
+    )
+    bot.register_message_handler(
+        functools.partial(message_handlers.message_handler, bot),
        content_types=[
            "text",
            "audio",
@@ -71,4 +72,9 @@ def create_bot(token: str, engine: Engine) -> AsyncTeleBot:
            "contact",
        ],
    )
+
+    bot.register_callback_query_handler(
+        functools.partial(callback_handlers.delete_message, bot),
+        lambda call: call.data == "DELETE",
+    )
    return bot
@@ -0,0 +1,8 @@
+from telebot.async_telebot import AsyncTeleBot
+from telebot.types import CallbackQuery as Callback
+
+from . import helper_functions
+
+
+async def delete_message(bot: AsyncTeleBot, call: Callback) -> None:
+    await helper_functions.delete_message(bot, call.message)
@@ -0,0 +1,85 @@
+import asyncio
+from typing import Any, Awaitable, Callable
+
+import telebot
+from telebot.async_telebot import AsyncTeleBot
+from . import markups
+
+Message = telebot.types.Message
+Handler = Callable[[Message], Awaitable[Any]]
+states: dict[tuple[int, int], Handler] = {}
+
+
+def register_state(
+    message: Message,
+    handler: Callable[[Message], Any],
+) -> None:
+    states[(message.chat.id, message.from_user.id)] = handler
+
+
+def reset_state(message: Message) -> None:
+    try:
+        del states[(message.chat.id, message.from_user.id)]
+    except KeyError:
+        pass
+
+
+def get_state(message: Message) -> Handler | None:
+    return states.get((message.chat.id, message.from_user.id))
+
+
+async def delete_message(
+    bot: AsyncTeleBot,
+    mes: Message,
+    *,
+    sleep_time: int = 0,
+) -> bool:
+    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
+    else:
+        return True
+
+
+async def send_tmp_message(
+    bot: AsyncTeleBot,
+    chat_id: telebot.types.Message,
+    text: str,
+    sleep_time: int = 5,
+) -> None:
+    bot_mes = await bot.send_message(chat_id, text, parse_mode="MarkdownV2")
+    await delete_message(bot, bot_mes, sleep_time=sleep_time)
+
+
+async def base_handler(
+    bot: AsyncTeleBot, mes: Message, prev_mes: Message | None = None
+) -> None:
+    reset_state(mes)
+    await delete_message(bot, mes)
+    if prev_mes is not None:
+        await delete_message(bot, prev_mes)
+
+
+async def send_deleteable_message(
+    bot: AsyncTeleBot,
+    chat_id: int,
+    text: str,
+) -> None:
+    """Sends a message with a delete button"""
+    markup = markups.deletion_markup()
+    await bot.send_message(
+        chat_id,
+        text,
+        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
@@ -0,0 +1,8 @@
+from telebot.types import InlineKeyboardButton, InlineKeyboardMarkup
+
+
+def deletion_markup() -> InlineKeyboardMarkup:
+    markup = InlineKeyboardMarkup(row_width=5)
+    button = InlineKeyboardButton("Удалить сообщение", callback_data="DELETE")
+    markup.add(button)
+    return markup
@@ -1,7 +1,5 @@
-import asyncio
 import functools
 import gc
-from typing import Any, Awaitable, Callable

 import telebot
 from sqlalchemy.future import Engine
@@ -16,51 +14,18 @@ from ..account_checks import (
 )
 from ..account_parsing import accounts_to_json, json_to_accounts
 from ..decrypted_account import DecryptedAccount
+from . import markups
+from .helper_functions import (
+    base_handler,
+    delete_message,
+    escape,
+    get_state,
+    register_state,
+    send_deleteable_message,
+    send_tmp_message,
+)

 Message = telebot.types.Message
-states: dict[tuple[int, int], Callable[[Message], Awaitable[Any]]] = {}
-
-
-def _register_state(
-    message: Message,
-    handler: Callable[[Message], Any],
-) -> None:
-    states[(message.chat.id, message.from_user.id)] = handler
-
-
-async def _delete_message(
-    bot: AsyncTeleBot,
-    mes: Message,
-    *,
-    sleep_time: int = 0,
-) -> bool:
-    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
-    else:
-        return True
-
-
-async def _send_tmp_message(
-    bot: AsyncTeleBot,
-    chat_id: telebot.types.Message,
-    text: str,
-    sleep_time: int = 5,
-) -> None:
-    bot_mes = await bot.send_message(chat_id, text, parse_mode="MarkdownV2")
-    await _delete_message(bot, bot_mes, sleep_time=sleep_time)
-
-
-async def _base_handler(
-    bot: AsyncTeleBot, mes: Message, prev_mes: Message | None = None
-) -> None:
-    states.pop((mes.chat.id, mes.from_user.id), None)
-    await _delete_message(bot, mes)
-    if prev_mes is not None:
-        await _delete_message(bot, prev_mes)


 async def get_accounts(
@@ -68,55 +33,54 @@ async def get_accounts(
    engine: Engine,
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)
    accounts = database.get.get_accounts(
        engine,
        mes.from_user.id,
        to_sort=True,
    )
    if not accounts:
-        return await _send_tmp_message(bot, mes.chat.id, "У вас нет аккаунтов")
+        return await send_tmp_message(bot, mes.chat.id, "У вас нет аккаунтов")

    # Make accounts copyable and escape special chars
    accounts = [f"`{account}`" for account in accounts]
-    await _send_tmp_message(
+    await send_deleteable_message(
        bot,
        mes.chat.id,
        "Ваши аккаунты:\n"
        + "\n".join(accounts)
        + "\nНажмите на название, чтобы скопировать\n"
        f"Всего {len(accounts)}",
-        30,
    )


 async def delete_all(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)
    bot_mes = await bot.send_message(
        mes.chat.id,
        "Вы действительно хотите удалить все ваши аккаунты? Это действие "
        "нельзя отменить. "
        "Отправьте YES для подтверждения",
    )
-    _register_state(mes, functools.partial(_delete_all2, bot, engine, bot_mes))
+    register_state(mes, functools.partial(_delete_all2, bot, engine, bot_mes))


 async def _delete_all2(
    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    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)
-        await _send_tmp_message(
+        await send_tmp_message(
            bot,
            mes.chat.id,
            "Всё успешно удалено",
            sleep_time=10,
        )
    else:
-        await _send_tmp_message(
+        await send_tmp_message(
            bot,
            mes.chat.id,
            "Вы отправили не YES, ничего не удалено",
@@ -128,15 +92,15 @@ async def set_master_password(
    engine: Engine,
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes, None)
+    await base_handler(bot, mes, None)
    if database.get.get_master_pass(engine, mes.from_user.id) is not None:
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Мастер пароль уже существует",
        )
    bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
-    _register_state(
+    register_state(
        mes,
        functools.partial(_set_master_pass2, bot, engine, bot_mes),
    )
@@ -145,10 +109,10 @@ async def set_master_password(
 async def _set_master_pass2(
    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    master_password = encryption.master_pass.encrypt_master_pass(
        mes.from_user.id,
@@ -156,7 +120,7 @@ async def _set_master_pass2(
    )
    database.add.add_master_pass(engine, master_password)

-    await _send_tmp_message(bot, mes.chat.id, "Успех")
+    await send_tmp_message(bot, mes.chat.id, "Успех")
    del mes, text
    gc.collect()

@@ -166,10 +130,10 @@ async def reset_master_pass(
    engine: Engine,
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)

    if database.get.get_master_pass(engine, mes.from_user.id) is None:
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Мастер пароль не задан",
@@ -181,7 +145,7 @@ async def reset_master_pass(
        "будут удалены навсегда",
    )

-    _register_state(
+    register_state(
        mes,
        functools.partial(_reset_master_pass2, bot, engine, bot_mes),
    )
@@ -190,10 +154,10 @@ async def reset_master_pass(
 async def _reset_master_pass2(
    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    master_password = encryption.master_pass.encrypt_master_pass(
        mes.from_user.id,
@@ -202,7 +166,7 @@ async def _reset_master_pass2(
    database.delete.purge_accounts(engine, mes.from_user.id)
    database.change.change_master_pass(engine, master_password)

-    await _send_tmp_message(
+    await send_tmp_message(
        bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
    )
    del mes, text
@@ -210,21 +174,21 @@ async def _reset_master_pass2(


 async def add_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)

    master_password_from_db = database.get.get_master_pass(
        engine,
        mes.from_user.id,
    )
    if master_password_from_db is None:
-        return await _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
+        return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")

    bot_mes = await bot.send_message(
        mes.chat.id,
        "Отправьте название аккаунта",
    )

-    _register_state(
+    register_state(
        mes,
        functools.partial(_add_account2, bot, engine, bot_mes),
    )
@@ -233,26 +197,26 @@ async def add_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
 async def _add_account2(
    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    if not check_account_name(text):
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Не корректное название аккаунта",
        )
    if text in database.get.get_accounts(engine, mes.from_user.id):
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot, mes.chat.id, "Аккаунт с таким именем уже существует"
        )

    bot_mes = await bot.send_message(mes.chat.id, "Отправьте логин")
    data = {"name": text}

-    _register_state(
+    register_state(
        mes,
        functools.partial(_add_account3, bot, engine, bot_mes, data),
    )
@@ -265,12 +229,12 @@ async def _add_account3(
    data: dict[str, str],
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
    if not check_login(text):
-        return await _send_tmp_message(bot, mes.chat.id, "Не корректный логин")
+        return await send_tmp_message(bot, mes.chat.id, "Не корректный логин")

    data["login"] = text

@@ -279,7 +243,7 @@ async def _add_account3(
        "Отправьте пароль от аккаунта",
    )

-    _register_state(
+    register_state(
        mes,
        functools.partial(_add_account4, bot, engine, bot_mes, data),
    )
@@ -292,12 +256,12 @@ async def _add_account4(
    data: dict[str, str],
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
    if not check_password(text):
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Не корректный пароль",
@@ -307,7 +271,7 @@ async def _add_account4(

    bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")

-    _register_state(
+    register_state(
        mes,
        functools.partial(_add_account5, bot, engine, bot_mes, data),
    )
@@ -320,14 +284,14 @@ async def _add_account5(
    data: dict[str, str],
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    master_password = database.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(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Не подходит главный пароль",
@@ -351,7 +315,7 @@ async def _add_account5(
        encrypted_account,
    )

-    await _send_tmp_message(
+    await send_tmp_message(
        bot,
        mes.chat.id,
        "Успех" if result else "Произошла не предвиденная ошибка",
@@ -363,17 +327,17 @@ async def _add_account5(


 async def get_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)

    master_pass = database.get.get_master_pass(engine, mes.from_user.id)
    if master_pass is None:
-        return await _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
+        return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")

    bot_mes = await bot.send_message(
        mes.chat.id,
        "Отправьте название аккаунта",
    )
-    _register_state(
+    register_state(
        mes,
        functools.partial(_get_account2, bot, engine, bot_mes),
    )
@@ -382,16 +346,16 @@ async def get_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
 async def _get_account2(
    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    if text not in database.get.get_accounts(engine, mes.from_user.id):
-        return await _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
+        return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")

    bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
-    _register_state(
+    register_state(
        mes,
        functools.partial(_get_account3, bot, engine, bot_mes, text),
    )
@@ -404,10 +368,10 @@ async def _get_account3(
    name: str,
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    master_password = database.get.get_master_pass(
        engine,
@@ -415,7 +379,7 @@ async def _get_account3(
    )

    if not encryption.master_pass.check_master_pass(text, master_password):
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Не подходит мастер пароль",
@@ -426,12 +390,12 @@ async def _get_account3(
        account,
        text,
    )
-    await _send_tmp_message(
+    await send_deleteable_message(
        bot,
        mes.chat.id,
+        f"Название:\n{escape(account.name)}\n"
        f"Логин:\n`{account.login}`\nПароль:\n`{account.password}`\nНажмите "
        "на логин или пароль, чтобы скопировать",
-        30,
    )

    del text, mes
@@ -443,17 +407,17 @@ async def delete_account(
    engine: Engine,
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)

    master_pass = database.get.get_master_pass(engine, mes.from_user.id)
    if master_pass is None:
-        return await _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
+        return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")

    bot_mes = await bot.send_message(
        mes.chat.id, "Отправьте название аккаунта, который вы хотите удалить"
    )

-    _register_state(
+    register_state(
        mes,
        functools.partial(_delete_account2, bot, engine, bot_mes),
    )
@@ -465,13 +429,13 @@ async def _delete_account2(
    prev_mes: Message,
    mes: Message,
 ):
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    if text not in database.get.get_accounts(engine, mes.from_user.id):
-        return await _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
+        return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")

    bot_mes = await bot.send_message(
        mes.from_user.id,
@@ -479,8 +443,9 @@ async def _delete_account2(
        "подтверждения",
    )

-    _register_state(
-        mes, functools.partial(_delete_account3, bot, engine, bot_mes, text)
+    register_state(
+        mes,
+        functools.partial(_delete_account3, bot, engine, bot_mes, text),
    )


@@ -491,13 +456,13 @@ async def _delete_account3(
    account_name: str,
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text != "YES":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    database.delete.delete_account(engine, mes.from_user.id, account_name)
-    await _send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
+    await send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")


 async def help_command(bot: AsyncTeleBot, mes: Message) -> None:
@@ -519,42 +484,42 @@ async def help_command(bot: AsyncTeleBot, mes: Message) -> None:


 async def cancel(bot: AsyncTeleBot, mes: Message) -> None:
-    await _base_handler(bot, mes)
-    await _send_tmp_message(bot, mes.chat.id, "Нет активного действия")
+    await base_handler(bot, mes)
+    await send_tmp_message(bot, mes.chat.id, "Нет активного действия")


 async def export(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)
    master_password_from_db = database.get.get_master_pass(
        engine,
        mes.from_user.id,
    )

    if master_password_from_db is None:
-        return await _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
+        return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")

    if not database.get.get_accounts(engine, mes.from_user.id):
-        return await _send_tmp_message(bot, mes.chat.id, "Нет аккаунтов")
+        return await send_tmp_message(bot, mes.chat.id, "Нет аккаунтов")

    bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")

-    _register_state(mes, functools.partial(_export2, bot, engine, bot_mes))
+    register_state(mes, functools.partial(_export2, bot, engine, bot_mes))


 async def _export2(
    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    master_password = database.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(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Не подходит мастер пароль",
@@ -563,11 +528,14 @@ async def _export2(
    accounts = database.get.get_all_accounts(engine, mes.from_user.id)
    accounts = encryption.other_accounts.decrypt_multiple(accounts, text)
    json_io = accounts_to_json(accounts)
-    bot_mes = await bot.send_document(mes.chat.id, json_io)
+    await bot.send_document(
+        mes.chat.id,
+        json_io,
+        reply_markup=markups.deletion_markup(),
+    )

    del text, accounts, json_io
    gc.collect()
-    await _delete_message(bot, bot_mes, sleep_time=30)


 async def import_accounts(
@@ -575,37 +543,37 @@ async def import_accounts(
    engine: Engine,
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)
    master_password_from_db = database.get.get_master_pass(
        engine,
        mes.from_user.id,
    )

    if master_password_from_db is None:
-        return await _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
+        return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")

    bot_mes = await bot.send_message(mes.chat.id, "Отправьте json файл")

-    _register_state(mes, functools.partial(_import2, bot, engine, bot_mes))
+    register_state(mes, functools.partial(_import2, bot, engine, bot_mes))


 async def _import2(
    bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    if mes.text is not None:
        text = mes.text.strip()
        if text == "/cancel":
-            return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+            return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    if mes.document is None:
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Вы должны отправить документ",
        )
    if mes.document.file_size > 102_400:  # If file size is bigger than 100 MB
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Файл слишком большой",
@@ -619,14 +587,14 @@ async def _import2(
            mes.from_user.id,
        )
    except Exception:
-        return await _send_tmp_message(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Ошибка во время работы с файлом",
        )

    bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
-    _register_state(
+    register_state(
        mes,
        functools.partial(_import3, bot, engine, bot_mes, accounts),
    )
@@ -639,17 +607,17 @@ async def _import3(
    accounts: list[DecryptedAccount],
    mes: Message,
 ) -> None:
-    await _base_handler(bot, mes, prev_mes)
+    await base_handler(bot, mes, prev_mes)
    text = mes.text.strip()
    if text == "/cancel":
-        return await _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+        return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")

    master_password = database.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(
+        return await send_tmp_message(
            bot,
            mes.chat.id,
            "Не подходит мастер пароль",
@@ -675,13 +643,13 @@ async def _import3(
    else:
        mes_text = "Успех"

-    await _send_tmp_message(bot, mes.chat.id, mes_text, 10)
+    await send_tmp_message(bot, mes.chat.id, mes_text, 10)
    del text, mes, accounts
    gc.collect()


 async def gen_password(bot: AsyncTeleBot, mes: Message) -> None:
-    await _base_handler(bot, mes)
+    await base_handler(bot, mes)
    # Generate 10 passwords and put 'em in the backticks
    passwords = (f"`{generate_password.gen_password()}`" for _ in range(10))
    text = (
@@ -689,14 +657,14 @@ async def gen_password(bot: AsyncTeleBot, mes: Message) -> None:
        + "\n".join(passwords)
        + "\nНажмите на пароль, чтобы его скопировать"
    )
-    await _send_tmp_message(bot, mes.chat.id, text, 15)
+    await send_tmp_message(bot, mes.chat.id, text, 15)


 async def message_handler(bot: AsyncTeleBot, mes: Message) -> None:
-    handler = states.get((mes.chat.id, mes.from_user.id))
+    handler = get_state(mes)
    if handler is None:
-        await _delete_message(bot, mes)
-        await _send_tmp_message(
+        await delete_message(bot, mes)
+        await send_tmp_message(
            bot,
            mes.chat.id,
            "Вы отправили не корректное сообщение",