diff --git a/src/bot/__init__.py b/src/bot/__init__.py index b40789c..4139f92 100644 --- a/src/bot/__init__.py +++ b/src/bot/__init__.py @@ -22,7 +22,7 @@ def create_bot(token: str, engine: mariadb.Connection) -> telebot.TeleBot: functools.partial(handlers.get_accounts, bot, engine), commands=["get_accounts"] ) bot.register_message_handler( - functools.partial(handlers.add_record, bot, engine), commands=["add_account"] + functools.partial(handlers.add_account, bot, engine), commands=["add_account"] ) bot.register_message_handler( functools.partial(handlers.delete_all, bot, engine), commands=["delete_all"] diff --git a/src/bot/handlers.py b/src/bot/handlers.py index 59a62eb..f857d0b 100644 --- a/src/bot/handlers.py +++ b/src/bot/handlers.py @@ -1,5 +1,5 @@ +import functools import gc -import shlex import time import telebot @@ -7,6 +7,8 @@ from sqlalchemy.future import Engine from .. import cryptography, database +Message = telebot.types.Message + def _send_tmp_message( bot: telebot.TeleBot, chat_id: telebot.types.Message, text: str, timeout: int = 5 @@ -16,41 +18,6 @@ def _send_tmp_message( bot.delete_message(chat_id, bot_mes.id) -def add_record( - bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message -) -> None: - data = shlex.split(mes.text) - bot.delete_message(mes.chat.id, mes.id) - if len(data) != 5: - return _send_tmp_message(bot, mes.chat.id, "Неправильное кол-во аргументов") - - master_password = data[4] - master_password_from_db = database.get.get_master_pass(engine, mes.from_user.id) - if master_password_from_db is None: - return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") - - master_salt, hash_pass = master_password_from_db - if ( - cryptography.master_pass.encrypt_master_pass(master_password, master_salt) - != hash_pass - ): - return _send_tmp_message(bot, mes.chat.id, "Не подходит главный пароль") - - enc_login, enc_pass, salt = cryptography.other_accounts.encrypt_account_info( - data[2], data[3], master_password.encode("utf-8") - ) - result = database.add.add_account( - engine, mes.from_user.id, data[1], salt, enc_login, enc_pass - ) - _send_tmp_message( - bot, - mes.chat.id, - "Успех" if result else "Ошибка, вероятно аккаунт с таким же именем существует", - ) - del data, master_password, mes - gc.collect() - - def get_accounts( bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message ) -> None: @@ -61,65 +28,10 @@ def get_accounts( bot, mes.chat.id, "Ваши аккаунты:\n" + "\n".join(accounts) if accounts else "У вас нет аккаунтов", + timeout=30, ) -def set_master_password( - bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message -) -> None: - data = shlex.split(mes.text) - bot.delete_message(mes.chat.id, mes.id) - - if database.get.get_master_pass(engine, mes.from_user.id) is not None: - return _send_tmp_message(bot, mes.chat.id, "У вас уже установлен мастер пароль") - - if len(data) != 2: - return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов") - - pass_, salt = cryptography.master_pass.encrypt_master_pass(data[1]) - result = database.add.add_master_pass(engine, mes.from_user.id, salt, pass_) - _send_tmp_message( - bot, - mes.chat.id, - "Успех" if result else "Ошибка, вероятно аккаунт с таким же именем существует", - ) - - del data, mes - gc.collect() - - -def get_account( - bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message -) -> None: - data = shlex.split(mes.text) - bot.delete_message(mes.chat.id, mes.id) - - if len(data) != 3: - return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов") - - master_pass = database.get.get_master_pass(engine, mes.from_user.id) - if master_pass is None: - return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") - master_salt, hash_pass = master_pass - - if data[1] not in database.get.get_accounts(engine, mes.from_user.id): - return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта") - - if cryptography.master_pass.encrypt_master_pass(data[2], master_salt) != hash_pass: - return _send_tmp_message(bot, mes.chat.id, "Не подходит главный пароль") - - salt, enc_login, enc_pass = database.get.get_account_info( - engine, mes.from_user.id, data[1] - ) - login, passwd = cryptography.other_accounts.decrypt_account_info( - enc_login, enc_pass, data[2].encode("utf-8"), salt - ) - _send_tmp_message(bot, mes.chat.id, f"Логин:\n{login}\nПароль:\n{passwd}", 30) - - del data, mes, passwd, login - gc.collect() - - def delete_all( bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message ) -> None: @@ -128,50 +40,236 @@ def delete_all( bot.send_message(mes.chat.id, "Все ваши данные удалены из базы данных") -def reset_master_pass( - bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message -) -> None: +def _base(bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None) -> None: bot.delete_message(mes.chat.id, mes.id) - data = shlex.split(mes.text) - if len(data) != 2: - return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов") + if prev_mes is not None: + bot.delete_message(prev_mes.chat.id, prev_mes.id) + + +def set_master_password(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: + _base(bot, mes, None) + if database.get.get_master_pass(engine, mes.from_user.id) is not None: + return _send_tmp_message(bot, mes.chat.id, "Мастер пароль уже существует") + bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль") + bot.register_next_step_handler( + mes, functools.partial(_set_master_pass2, bot, engine, bot_mes) + ) + + +def _set_master_pass2( + bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message +) -> None: + _base(bot, mes, prev_mes) + hash_, salt = cryptography.master_pass.encrypt_master_pass(mes.text.strip()) + database.add.add_master_pass(engine, mes.from_user.id, salt, hash_) + _send_tmp_message(bot, mes.chat.id, "Успех") + del mes + gc.collect() + + +def reset_master_pass(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: + _base(bot, mes) + if database.get.get_master_pass(engine, mes.from_user.id) is None: + return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан") + bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль") + bot.register_next_step_handler( + mes, functools.partial(_reset_master_pass2, bot, engine, bot_mes) + ) + + +def _reset_master_pass2( + bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message +) -> None: + _base(bot, mes, prev_mes) + if database.get.get_master_pass(engine, mes.from_user.id) is None: + return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан") + hash_, salt = cryptography.master_pass.encrypt_master_pass(mes.text.strip()) + database.delete.purge_accounts(engine, mes.from_user.id) + database.change.change_master_pass(engine, mes.from_user.id, salt, hash_) + _send_tmp_message( + bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён" + ) + del mes + gc.collect() + + +def add_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: + _base(bot, mes) master_password_from_db = database.get.get_master_pass(engine, mes.from_user.id) if master_password_from_db is None: return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") - master_password = data[1] - enc_pass, salt = cryptography.master_pass.encrypt_master_pass(master_password) - database.delete.purge_accounts(engine, mes.from_user.id) - database.change.change_master_pass(engine, mes.from_user.id, salt, enc_pass) - _send_tmp_message( - bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён" + bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта") + + bot.register_next_step_handler( + mes, functools.partial(_add_account2, bot, engine, bot_mes) ) -def delete_account( - bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message +def _add_account2( + bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message ) -> None: - data = shlex.split(mes.text) - bot.delete_message(mes.chat.id, mes.id) + _base(bot, mes, prev_mes) - if len(data) != 2: - return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов") + if ( + database.get.get_account_info(engine, mes.from_user.id, mes.text.strip()) + is not None + ): + _send_tmp_message(bot, mes.chat.id, "Аккаунт с таким именем уже существует") - if data[1] not in database.get.get_accounts(engine, mes.from_user.id): + bot_mes = bot.send_message(mes.chat.id, "Отправьте логин") + data = {"name": mes.text.strip()} + + bot.register_next_step_handler( + mes, functools.partial(_add_account3, bot, engine, bot_mes, data) + ) + + +def _add_account3( + bot: telebot.TeleBot, + engine: Engine, + prev_mes: Message, + data: dict[str, str], + mes: Message, +) -> None: + _base(bot, mes, prev_mes) + data["login"] = mes.text.strip() + + bot_mes = bot.send_message(mes.chat.id, "Отправьте пароль от аккаунта") + + bot.register_next_step_handler( + mes, functools.partial(_add_account4, bot, engine, bot_mes, data) + ) + + +def _add_account4( + bot: telebot.TeleBot, + engine: Engine, + prev_mes: Message, + data: dict[str, str], + mes: Message, +) -> None: + _base(bot, mes, prev_mes) + data["passwd"] = mes.text.strip() + + bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль") + + bot.register_next_step_handler( + mes, functools.partial(_add_account5, bot, engine, bot_mes, data) + ) + + +def _add_account5( + bot: telebot.TeleBot, + engine: Engine, + prev_mes: Message, + data: dict[str, str], + mes: Message, +) -> None: + _base(bot, mes, prev_mes) + text = mes.text.strip() + + salt, hash_ = database.get.get_master_pass(engine, mes.from_user.id) + if cryptography.master_pass.encrypt_master_pass(text, salt) != hash_: + return _send_tmp_message(bot, mes.chat.id, "Не подходит главный пароль") + + name, login, passwd = data["name"], data["login"], data["passwd"] + + enc_login, enc_pass, salt = cryptography.other_accounts.encrypt_account_info( + login, passwd, text.encode("utf-8") + ) + + database.add.add_account(engine, mes.from_user.id, name, salt, enc_login, enc_pass) + + _send_tmp_message(bot, mes.chat.id, "Успех") + + del data, name, login, passwd, enc_login + + gc.collect() + + +def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: + _base(bot, mes) + bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта") + bot.register_next_step_handler( + mes, functools.partial(_get_account2, bot, engine, bot_mes) + ) + + +def _get_account2( + bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message +) -> None: + _base(bot, mes, prev_mes) + text = mes.text.strip() + + if text not in database.get.get_accounts(engine, mes.from_user.id): return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта") - database.delete.delete_account(engine, mes.from_user.id, data[1]) + bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль") + bot.register_next_step_handler( + mes, functools.partial(_get_account3, bot, engine, bot_mes, text) + ) + + +def _get_account3( + bot: telebot.TeleBot, engine: Engine, prev_mes: Message, name: str, mes: Message +) -> None: + _base(bot, mes, prev_mes) + text = mes.text.strip() + + master_pass = database.get.get_master_pass(engine, mes.from_user.id) + if master_pass is None: + return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") + + master_salt, hash_pass = master_pass + + if cryptography.master_pass.encrypt_master_pass(text, master_salt) != hash_pass: + return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль") + + salt, enc_login, enc_pass = database.get.get_account_info( + engine, mes.from_user.id, name + ) + login, passwd = cryptography.other_accounts.decrypt_account_info( + enc_login, enc_pass, text.encode("utf-8"), salt + ) + _send_tmp_message(bot, mes.chat.id, f"Логин:\n{login}\nПароль:\n{passwd}", 30) + + del text, mes, passwd, login + gc.collect() + + +def delete_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: + _base(bot, mes) + bot_mes = bot.send_message( + mes.chat.id, "Отправьте название аккаунта, который вы хотите удалить" + ) + + bot.register_next_step_handler( + mes, functools.partial(_delete_account2, bot, engine, bot_mes) + ) + + +def _delete_account2( + bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message +) -> None: + _base(bot, mes, prev_mes) + text = mes.text.strip() + + if text not in database.get.get_accounts(engine, mes.from_user.id): + return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта") + + database.delete.delete_account(engine, mes.from_user.id, text) _send_tmp_message(bot, mes.chat.id, "Аккаунт удалён") def help(bot: telebot.TeleBot, mes: telebot.types.Message) -> None: message = """Команды: -/set_master_pass {Мастер пароль} - установить мастер пароль -/add_account {Название} {Логин} {Пароль} {Мастер пароль} - создать аккаунт +/set_master_pass - установить мастер пароль +/add_account - создать аккаунт /get_accounts - получить список аккаунтов -/get_account {Название} {Мастер пароль} - получить логин и пароль аккаунта -/delete_account {Название} - удалить аккаунт -/delete_all - удалить все аккаунты и мастер пароль -/reset_master_pass {Новый мастер пароль} - удалить все аккаунты и изменить мастер пароль""" +/get_account - получить логин и пароль аккаунта +/delete_account - удалить аккаунт +/delete_all - удалить все аккаунты и мастер пароль (Спать хочу -> подтверждения нету) +/reset_master_pass- удалить все аккаунты и изменить мастер пароль""" bot.send_message(mes.chat.id, message)