From 042ca9312eaa7d4b0a7307c49beed729fe74c772 Mon Sep 17 00:00:00 2001
From: StNicolay <stnicolay2007@mail.ru>
Date: Fri, 4 Nov 2022 01:15:28 +0300
Subject: [PATCH] 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
---
 src/bot/handlers.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/bot/handlers.py b/src/bot/handlers.py
index c3b17a0..697924f 100644
--- a/src/bot/handlers.py
+++ b/src/bot/handlers.py
@@ -83,8 +83,8 @@ def _set_master_pass2(
     if text == "/cancel":
         return send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
-    hash_, salt = cryptography.master_pass.encrypt_master_pass(text)
-    database.add.add_master_pass(engine, mes.from_user.id, salt, hash_)
+    hash_pass, master_salt = cryptography.master_pass.encrypt_master_pass(text)
+    database.add.add_master_pass(engine, mes.from_user.id, master_salt, hash_pass)
     send_tmp_message(bot, mes.chat.id, "Успех")
     del mes, text
     gc.collect()
@@ -242,6 +242,11 @@ def _add_account5(
 def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
     base_handler(bot, mes)
     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(
         mes, functools.partial(_get_account2, bot, engine, bot_mes)
     )
@@ -272,11 +277,7 @@ def _get_account3(
     if text == "/cancel":
         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
+    master_salt, hash_pass = database.get.get_master_pass(engine, mes.from_user.id)
 
     if cryptography.master_pass.encrypt_master_pass(text, master_salt) != hash_pass:
         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:
     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(
         mes.chat.id, "Отправьте название аккаунта, который вы хотите удалить"
     )
@@ -407,14 +413,16 @@ def _import2(
             return send_tmp_message(bot, mes.chat.id, "Успешная отмена")
 
     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)
     downloaded_file = bot.download_file(file_info.file_path)
     try:
         accounts = json_to_accounts(downloaded_file.decode("utf-8"))
     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.register_next_step_handler(