switched mariadb connector to pymysql

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

Dockerfile

@@ -13,9 +13,6 @@ RUN adduser -u 1000 --disabled-password --gecos "" appuser && chown -R appuser /
 
 # Install deps
 RUN apt update && apt full-upgrade -y
-RUN apt install curl gcc g++ -y
-RUN curl -sS https://downloads.mariadb.com/MariaDB/mariadb_repo_setup | bash
-RUN apt install libmariadb3 libmariadb-dev -y
 
 # Install pip requirements
 COPY requirements.txt .

requirements.txt

@@ -1,6 +1,6 @@
 bcrypt
 cryptography
-mariadb
+pymysql
 python-dotenv
 pyTelegramBotAPI
 sqlmodel

src/bot/__init__.py

@@ -1,6 +1,6 @@
 import functools
 
-import mariadb
+from sqlalchemy.future import Engine
 import telebot
 
 from . import handlers, utils
@@ -8,7 +8,7 @@ from . import handlers, utils
 __all__ = ["handlers", "utils"]
 
 
-def create_bot(token: str, engine: mariadb.Connection) -> telebot.TeleBot:
+def create_bot(token: str, engine: Engine) -> telebot.TeleBot:
     bot = telebot.TeleBot(token)
     bot.register_message_handler(
         functools.partial(handlers.set_master_password, bot, engine),

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(

src/database/prepare.py

@@ -5,9 +5,7 @@ from . import models
 
 
 def get_engine(host: str, user: str, passwd: str, db: str) -> Engine:
-    engine = sqlmodel.create_engine(
-        f"mariadb+mariadbconnector://{user}:{passwd}@{host}/{db}"
-    )
+    engine = sqlmodel.create_engine(f"mariadb+pymysql://{user}:{passwd}@{host}/{db}")
     return engine