Cleaned up code in handlers

Renamed variables in _set_master_pass2 for consistency
Added a few missing returns in guard clauses
Added file size limit for importing account to 100 MB
This commit is contained in:
StNicolay 2022-11-04 01:15:28 +03:00
parent b10b1fe2b3
commit 042ca9312e

View File

@ -83,8 +83,8 @@ def _set_master_pass2(
if text == "/cancel": if text == "/cancel":
return send_tmp_message(bot, mes.chat.id, "Успешная отмена") return send_tmp_message(bot, mes.chat.id, "Успешная отмена")
hash_, salt = cryptography.master_pass.encrypt_master_pass(text) hash_pass, master_salt = cryptography.master_pass.encrypt_master_pass(text)
database.add.add_master_pass(engine, mes.from_user.id, salt, hash_) database.add.add_master_pass(engine, mes.from_user.id, master_salt, hash_pass)
send_tmp_message(bot, mes.chat.id, "Успех") send_tmp_message(bot, mes.chat.id, "Успех")
del mes, text del mes, text
gc.collect() gc.collect()
@ -242,6 +242,11 @@ def _add_account5(
def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
base_handler(bot, mes) base_handler(bot, mes)
bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта") bot_mes = bot.send_message(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, "Нет мастер пароля")
bot.register_next_step_handler( bot.register_next_step_handler(
mes, functools.partial(_get_account2, bot, engine, bot_mes) mes, functools.partial(_get_account2, bot, engine, bot_mes)
) )
@ -272,11 +277,7 @@ def _get_account3(
if text == "/cancel": if text == "/cancel":
return send_tmp_message(bot, mes.chat.id, "Успешная отмена") return send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_pass = database.get.get_master_pass(engine, mes.from_user.id) master_salt, hash_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: if cryptography.master_pass.encrypt_master_pass(text, master_salt) != hash_pass:
return send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль") return send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
@ -300,6 +301,11 @@ def _get_account3(
def delete_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: def delete_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
base_handler(bot, mes) base_handler(bot, mes)
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, "Нет мастер пароля")
bot_mes = bot.send_message( bot_mes = bot.send_message(
mes.chat.id, "Отправьте название аккаунта, который вы хотите удалить" mes.chat.id, "Отправьте название аккаунта, который вы хотите удалить"
) )
@ -407,14 +413,16 @@ def _import2(
return send_tmp_message(bot, mes.chat.id, "Успешная отмена") return send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if mes.document is None: if mes.document is None:
send_tmp_message(bot, mes.chat.id, "Вы должны отправить документ") return send_tmp_message(bot, mes.chat.id, "Вы должны отправить документ")
if mes.document.file_size > 102_400: # If file size is bigger that 100 MB
return send_tmp_message(bot, mes.chat.id, "Файл слишком большой")
file_info = bot.get_file(mes.document.file_id) file_info = bot.get_file(mes.document.file_id)
downloaded_file = bot.download_file(file_info.file_path) downloaded_file = bot.download_file(file_info.file_path)
try: try:
accounts = json_to_accounts(downloaded_file.decode("utf-8")) accounts = json_to_accounts(downloaded_file.decode("utf-8"))
except Exception: except Exception:
send_tmp_message(bot, mes.chat.id, "Ошибка во время работы с файлом") return send_tmp_message(bot, mes.chat.id, "Ошибка во время работы с файлом")
bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль") bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль")
bot.register_next_step_handler( bot.register_next_step_handler(