Fixed a typo in README and changed one of the phrase to be more correct

.dockerignore

@@ -0,0 +1,26 @@
+**/__pycache__
+**/.venv
+**/.classpath
+**/.dockerignore
+**/.env
+**/.git
+**/.gitignore
+**/.project
+**/.settings
+**/.toolstarget
+**/.vs
+**/.vscode
+**/*.*proj.user
+**/*.dbmdl
+**/*.jfm
+**/bin
+**/charts
+**/docker-compose*
+**/compose*
+**/Dockerfile*
+**/node_modules
+**/npm-debug.log
+**/obj
+**/secrets.dev.yaml
+**/values.dev.yaml
+README.md

.gitignore

@@ -160,3 +160,5 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
+# Database data
+data/

Dockerfile

@@ -0,0 +1,28 @@
+FROM python:3.11-slim
+
+# Keeps Python from generating .pyc files in the container
+ENV PYTHONDONTWRITEBYTECODE=1
+
+# Turns off buffering for easier container logging
+ENV PYTHONUNBUFFERED=1
+
+WORKDIR /app
+
+# Creates new user
+RUN adduser -u 1000 --disabled-password --gecos "" appuser && chown -R appuser /app
+
+# Install deps
+RUN apt update && apt full-upgrade -y
+RUN apt install curl gcc -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 .
+RUN python -m pip install -r requirements.txt
+
+COPY . /app
+
+USER appuser
+
+CMD ["python", "main.py"]

README.md

@@ -1,18 +1,35 @@
 # PassManager
 
-### Цель
+## Описание
 
-**Создать бота для сохранения, созания паролей в зашифрованой базе данных.**
+Менеджер паролей для телеграм, где все пароли надёжно зашифрованы
 
-### Задачи
+## Документация
 
-* Создать телеграм бота
+### База данных
-* Научить его работать с БД (MariaDB)
-* Зашифровать БД
-* Сделать общедоступным
 
-### Инструменты
+В качестве базы данных данный бот использует mariadb, к которой бот должен иметь доступ всегда
 
-* Язык программирования Python
+### Команды и их синтаксис
-* PyTelegramBotAPI
+
-* MariaDB
+- /set_master_pass - установить мастер пароль
+- /add_account - создать аккаунт
+- /get_accounts - получить список аккаунтов
+- /get_account - получить логин и пароль аккаунта
+- /delete_account - удалить аккаунт
+- /delete_all - удалить все аккаунты и мастер пароль
+- /reset_master_pass- удалить все аккаунты и изменить мастер пароль
+- /cancel - отмена текущего действия
+- /help - помощь
+
+### Настройка
+
+Настройка производится через переменные среды. Их можно прописать в файле .env, если не хотите задавать их каждый раз в случае, если вы не используете Docker
+
+#### Переменные среды
+
+- DB_HOST - ip адрес базы данных
+- DB_USER - пользователь базы данных
+- DB_PASS - пароль пользователя базы данных
+- DB_NAME - название базы данных
+- TG_TOKEN - токен телеграма, его следует указать в файле .env в формате TG_TOKEN={токен}

compose.yaml

@@ -0,0 +1,33 @@
+version: '3.4'
+networks:
+  password_manager: {}
+
+services:
+  passmanager:
+    build:
+      context: .
+      dockerfile: ./Dockerfile
+    restart: always
+    environment:
+      DB_HOST: db
+      DB_USER: manager
+      DB_PASS: passwd123!
+      DB_NAME: passmanager
+      TG_TOKEN: ${TG_TOKEN}
+    depends_on:
+      - db
+    networks:
+      - password_manager
+
+  db:
+    image: jc21/mariadb-aria
+    restart: always
+    environment:
+      MYSQL_ROOT_PASSWORD: example123!
+      MYSQL_DATABASE: passmanager
+      MYSQL_USER: manager
+      MYSQL_PASSWORD: passwd123!
+    volumes:
+      - ./data:/var/lib/mysql
+    networks:
+      - password_manager

requirements.txt

@@ -3,3 +3,4 @@ cryptography
 mariadb
 python-dotenv
 pyTelegramBotAPI
+sqlmodel

src/__init__.py

@@ -1,20 +1,22 @@
 import os
 
-import mariadb
 from dotenv import load_dotenv
+from sqlalchemy.future import Engine
 
 from . import bot, cryptography, database
 
-
 __all__ = ["bot", "cryptography", "database"]
+engine: Engine
 
 
 def main() -> None:
     load_dotenv("./.env")
-    con = mariadb.connect(
+    engine = database.prepare.get_engine(
-        os.getenv("DB_HOST"),
+        host=os.getenv("DB_HOST"),
-        os.getenv("DB_USER"),
+        user=os.getenv("DB_USER"),
-        os.getenv("DB_PASS"),
+        passwd=os.getenv("DB_PASS"),
-        os.getenv("DB_NAME"),
+        db=os.getenv("DB_NAME"),
-    )
+    )  # type: ignore
-    database.prepare(con)
+    database.prepare.prepare(engine)
+    bot_ = bot.create_bot(os.getenv("TG_TOKEN"), engine)  # type: ignore
+    bot_.infinity_polling()

src/bot/__init__.py

@@ -0,0 +1,43 @@
+import functools
+
+import mariadb
+import telebot
+
+from . import handlers
+
+__all__ = ["handlers"]
+
+
+def create_bot(token: str, engine: mariadb.Connection) -> telebot.TeleBot:
+    bot = telebot.TeleBot(token)
+    bot.register_message_handler(
+        functools.partial(handlers.set_master_password, bot, engine),
+        commands=["set_master_pass"],
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.get_account, bot, engine), commands=["get_account"]
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.get_accounts, bot, engine), commands=["get_accounts"]
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.add_account, bot, engine), commands=["add_account"]
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.delete_all, bot, engine), commands=["delete_all"]
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.reset_master_pass, bot, engine),
+        commands=["reset_master_pass"],
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.delete_account, bot, engine),
+        commands=["delete_account"],
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.help, bot), commands=["help", "start"]
+    )
+    bot.register_message_handler(
+        functools.partial(handlers.cancel, bot), commands=["cancel"]
+    )
+    return bot

src/bot/handlers.py

@@ -0,0 +1,331 @@
+import functools
+import gc
+import time
+
+import telebot
+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
+) -> None:
+    bot_mes = bot.send_message(chat_id, text, parse_mode="MarkdownV2")
+    time.sleep(timeout)
+    bot.delete_message(chat_id, bot_mes.id)
+
+
+def get_accounts(
+    bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
+) -> None:
+    accounts = database.get.get_accounts(engine, mes.from_user.id)
+    bot.delete_message(mes.chat.id, mes.id)
+
+    return _send_tmp_message(
+        bot,
+        mes.chat.id,
+        "Ваши аккаунты:\n" + "\n".join(accounts) if accounts else "У вас нет аккаунтов",
+        timeout=30,
+    )
+
+
+def _base(bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None) -> None:
+    bot.delete_message(mes.chat.id, mes.id)
+    if prev_mes is not None:
+        bot.delete_message(prev_mes.chat.id, prev_mes.id)
+
+
+def delete_all(
+    bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
+) -> None:
+    _base(bot, mes)
+    bot_mes = bot.send_message(
+        mes.chat.id,
+        "Вы действительно хотите удалить все ваши аккаунты? Это действие нельзя отменить. Отправьте YES для подтверждения",
+    )
+    bot.register_next_step_handler(
+        mes, functools.partial(_delete_all, bot, engine, bot_mes)
+    )
+
+
+def _delete_all(
+    bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
+) -> None:
+    _base(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)
+        _send_tmp_message(bot, mes.chat.id, "Всё успешно удалено", timeout=10)
+    else:
+        _send_tmp_message(bot, mes.chat.id, "Вы отправили не YES, ничего не удалено")
+
+
+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)
+    text = mes.text.strip()
+    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_)
+    _send_tmp_message(bot, mes.chat.id, "Успех")
+    del mes, text
+    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)
+    text = mes.text.strip()
+    if text == "/cancel":
+        return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+
+    hash_, salt = cryptography.master_pass.encrypt_master_pass(text)
+    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, text
+    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, "Нет мастер пароля")
+
+    bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта")
+
+    bot.register_next_step_handler(
+        mes, functools.partial(_add_account2, bot, engine, bot_mes)
+    )
+
+
+def _add_account2(
+    bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
+) -> None:
+    _base(bot, mes, prev_mes)
+    text = mes.text.strip()
+    if text == "/cancel":
+        return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+
+    if text in database.get.get_accounts(engine, mes.from_user.id):
+        return _send_tmp_message(
+            bot, mes.chat.id, "Аккаунт с таким именем уже существует"
+        )
+
+    bot_mes = bot.send_message(mes.chat.id, "Отправьте логин")
+    data = {"name": text}
+
+    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)
+    text = mes.text.strip()
+    if text == "/cancel":
+        return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+
+    data["login"] = text
+
+    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)
+    text = mes.text.strip()
+    if text == "/cancel":
+        return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+
+    data["passwd"] = text
+
+    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()
+    if text == "/cancel":
+        return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+
+    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 == "/cancel":
+        return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+
+    if text not in database.get.get_accounts(engine, mes.from_user.id):
+        return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
+
+    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()
+    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
+
+    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}`\nНажмите на логин или пароль, чтобы скопировать",
+        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 == "/cancel":
+        return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
+
+    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 - создать аккаунт
+/get_accounts - получить список аккаунтов
+/get_account - получить логин и пароль аккаунта
+/delete_account - удалить аккаунт
+/delete_all - удалить все аккаунты и мастер пароль
+/reset_master_pass - удалить все аккаунты и изменить мастер пароль
+/cancel - отмена текущего действия
+/help - помощь"""
+    bot.send_message(mes.chat.id, message)
+
+
+def cancel(bot: telebot.TeleBot, mes: Message) -> None:
+    _send_tmp_message(bot, mes.chat.id, "Нет активного действия")

src/cryptography/master_pass.py

@@ -1,14 +1,26 @@
+from typing import overload
+
 import bcrypt
 
-print("Hi")
+
+@overload
+def encrypt_master_pass(passwd: str, salt: bytes) -> bytes:
+    ...
 
 
+@overload
 def encrypt_master_pass(passwd: str) -> tuple[bytes, bytes]:
+    ...
+
+
+def encrypt_master_pass(
+    passwd: str, salt: bytes | None = None
+) -> tuple[bytes, bytes] | bytes:
     """Hashes master password and return tuple of hashed password and salt"""
+    if salt is None:
         salt = bcrypt.gensalt()
+        gened_salt = True
+    else:
+        gened_salt = False
     hashed = bcrypt.hashpw(passwd.encode("utf-8"), salt)
-    return (hashed, salt)
+    return (hashed, salt) if gened_salt else hashed
-
-
-def encrypt_master_pass_known_salt(passwd: str, salt: bytes) -> bytes:
-    return bcrypt.hashpw(passwd.encode("utf-8"), salt)

src/cryptography/other_accounts.py

@@ -8,7 +8,7 @@ from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
 
-def _generate_key(salt: bytes, master_pass_hash: bytes) -> bytes:
+def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
     kdf = PBKDF2HMAC(
         algorithm=hashes.SHA256(),
         length=32,
@@ -16,17 +16,17 @@ def _generate_key(salt: bytes, master_pass_hash: bytes) -> bytes:
         iterations=100000,
         backend=default_backend(),
     )
-    key = base64.urlsafe_b64encode(kdf.derive(master_pass_hash))
+    key = base64.urlsafe_b64encode(kdf.derive(master_pass))
     return key
 
 
 def encrypt_account_info(
-    login: str, passwd: str, master_pass_hash: bytes
+    login: str, passwd: str, master_pass: bytes
 ) -> tuple[bytes, bytes, bytes]:
     """Encrypts login and password of a user using hash of their master password as a key.
     Returns a tuple of encrypted login password and salt"""
     salt = bcrypt.gensalt()
-    key = _generate_key(salt, master_pass_hash)
+    key = _generate_key(salt, master_pass)
     f = Fernet(key)
     enc_login = f.encrypt(login.encode("utf-8"))
     enc_passwd = f.encrypt(passwd.encode("utf-8"))
@@ -34,9 +34,9 @@ def encrypt_account_info(
 
 
 def decrypt_account_info(
-    enc_login: bytes, enc_pass: bytes, master_pass_hash: bytes, salt: bytes
+    enc_login: bytes, enc_pass: bytes, master_pass: bytes, salt: bytes
 ) -> tuple[str, str]:
-    key = _generate_key(salt, master_pass_hash)
+    key = _generate_key(salt, master_pass)
     f = Fernet(key)
     login_bytes = f.decrypt(enc_login)
     pass_bytes = f.decrypt(enc_pass)

src/database/__init__.py

@@ -1,3 +1,3 @@
-from . import add, delete, get, prepare
+from . import add, delete, get, models, prepare, change
 
-__all__ = ["add", "delete", "get", "prepare"]
+__all__ = ["add", "delete", "get", "models", "prepare", "change"]

src/database/add.py

@@ -1,42 +1,40 @@
-import traceback
+import sqlmodel
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.future import Engine
 
-import mariadb
+from . import models
-
-
-def add_master_pass(
-    id: int, hashed_passwd: bytes, salt: bytes, con: mariadb.Connection
-) -> bool:
-    cursor = con.cursor()
-    try:
-        cursor.execute(
-            "INSERT INTO master_pass (user_id, salt, passwd) VALUES (?, ?, ?)",
-            [id, hashed_passwd, salt],
-        )
-        cursor.close()
-    except Exception:
-        traceback.print_exc()
-        return False
-    else:
-        return True
 
 
 def add_account(
-    id: int,
+    engine: Engine,
-    acc_name: str,
+    user_id: int,
+    name: str,
     salt: bytes,
     enc_login: bytes,
-    enc_passwd: bytes,
+    enc_pass: bytes,
-    con: mariadb.Connection,
 ) -> bool:
-    cursor = con.cursor()
+    """Adds account to db. Returns true, if on success"""
-    try:
+    account = models.Account(
-        cursor.execute(
+        user_id=user_id, name=name, salt=salt, enc_login=enc_login, enc_pass=enc_pass
-            "INSERT INTO accounts (user_id, acc_name, salt, enc_login, enc_pass) VALUES (?, ?, ?, ?, ?, ?)",
-            [id, acc_name, salt, enc_login, enc_passwd],
     )
-        cursor.close()
+    try:
-    except Exception:
+        with sqlmodel.Session(engine) as session:
-        traceback.print_exc()
+            session.add(account)
+            session.commit()
+    except IntegrityError:
+        return False
+    else:
+        return True
+
+
+def add_master_pass(engine: Engine, user_id: int, salt: bytes, passwd: bytes) -> bool:
+    """Adds master password to db. Returns true, if on success"""
+    master_pass = models.MasterPass(user_id=user_id, salt=salt, passwd=passwd)
+    try:
+        with sqlmodel.Session(engine) as session:
+            session.add(master_pass)
+            session.commit()
+    except IntegrityError:
         return False
     else:
         return True

src/database/change.py

@@ -0,0 +1,17 @@
+import sqlmodel
+from sqlalchemy.future import Engine
+
+from . import models
+
+
+def change_master_pass(
+    engine: Engine, user_id: int, salt: bytes, passwd: bytes
+) -> None:
+    statement = sqlmodel.update(
+        models.MasterPass,
+        models.MasterPass.user_id == user_id,
+        {"salt": salt, "passwd": passwd},
+    )
+    with sqlmodel.Session(engine) as session:
+        session.exec(statement)
+        session.commit()

src/database/delete.py

@@ -1,15 +1,29 @@
-import mariadb
+import sqlmodel
+from sqlalchemy.future import Engine
+
+from . import models
 
 
-def delete_master_pass(con: mariadb.Connection, user_id: int) -> None:
+def purge_accounts(engine: Engine, user_id: int) -> None:
-    cursor = con.cursor()
+    statement = sqlmodel.delete(models.Account).where(models.Account.user_id == user_id)
-    cursor.execute("DELETE FROM master_pass WHERE user_id=?", [user_id])
+    with sqlmodel.Session(engine) as session:
-    cursor.close()
+        session.exec(statement)
+        session.commit()
 
 
-def delete_account(con: mariadb.Connection, user_id: int, account: str):
+def delete_master_pass(engine: Engine, user_id: int) -> None:
-    cursor = con.cursor()
+    statement = sqlmodel.delete(models.MasterPass).where(
-    cursor.execute(
+        models.MasterPass.user_id == user_id
-        "DELETE FROM accounts WHERE user_id = ? AND acc_name = ?", [user_id, account]
     )
-    cursor.close()
+    with sqlmodel.Session(engine) as session:
+        session.exec(statement)
+        session.commit()
+
+
+def delete_account(engine: Engine, user_id: int, name: str) -> None:
+    statement = sqlmodel.delete(models.Account).where(
+        models.Account.user_id == user_id, models.Account.name == name
+    )
+    with sqlmodel.Session(engine) as session:
+        session.exec(statement)
+        session.commit()

src/database/get.py

@@ -1,31 +1,38 @@
-import mariadb
+import sqlmodel
+from sqlalchemy.future import Engine
+
+from . import models
 
 
-def get_master_pass(con: mariadb.Connection, id: int) -> tuple[bytes, bytes]:
+def get_master_pass(engine: Engine, user_id: int) -> tuple[bytes, bytes] | None:
-    """Returns tuple of salt and hashed master password"""
+    """Gets master pass. Returns tuple of salt and password"""
-    cursor = con.cursor()
+    statement = sqlmodel.select(models.MasterPass).where(
-    cursor.execute("SELECT salt, passwd FROM master_pass IF user_id = ?", [id])
+        models.MasterPass.user_id == user_id
-    result = cursor.fetchone()
+    )
-    cursor.close()
+    with sqlmodel.Session(engine) as session:
-    return result
+        result = session.exec(statement).first()
+    if result is None:
+        return
+    return (result.salt, result.passwd)
 
 
-def get_accounts(con: mariadb.Connection, id: int) -> list[str]:
+def get_accounts(engine: Engine, user_id: int) -> list[str]:
-    """Returns list of user accounts"""
+    """Gets list of account names"""
-    cursor = con.cursor()
+    statement = sqlmodel.select(models.Account).where(models.Account.user_id == user_id)
-    cursor.execute("SELECT acc_name FROM accounts IF user_id = ?", [id])
+    with sqlmodel.Session(engine) as session:
-    return [i[0] for i in cursor.fetchall()]
+        result = session.exec(statement)
+        return [account.name for account in result]
 
 
 def get_account_info(
-    id: int, name: str, con: mariadb.Connection
+    engine: Engine, user_id: int, name: str
 ) -> tuple[bytes, bytes, bytes]:
-    """Returns tuple of salt, login and password"""
+    """Gets account info. Returns tuple of salt, login and password"""
-    cursor = con.cursor()
+    statement = sqlmodel.select(models.Account).where(
-    cursor.execute(
+        models.Account.user_id == user_id and models.Account.name == name
-        """SELECT salt, enc_login, enc_pass IF user_id = ? AND acc_name = ?""",
-        [id, name],
     )
-    result = cursor.fetchone()
+    with sqlmodel.Session(engine) as session:
-    cursor.close()
+        result = session.exec(statement).first()
-    return result
+        if result is None:
+            return
+        return (result.salt, result.enc_login, result.enc_pass)

src/database/models.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+import sqlmodel
+
+
+class MasterPass(sqlmodel.SQLModel, table=True):
+    __tablename__ = "master_passwords"
+    id: Optional[int] = sqlmodel.Field(primary_key=True)
+    user_id: int = sqlmodel.Field(nullable=False, index=True, unique=True)
+    salt: bytes = sqlmodel.Field(
+        sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
+    )
+    passwd: bytes = sqlmodel.Field(
+        sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
+    )
+
+
+class Account(sqlmodel.SQLModel, table=True):
+    __tablename__ = "accounts"
+    __table_args__ = (sqlmodel.UniqueConstraint("user_id", "name"),)
+    id: Optional[int] = sqlmodel.Field(primary_key=True)
+    user_id: int = sqlmodel.Field(nullable=False, index=True)
+    name: str = sqlmodel.Field(nullable=False, index=True, max_length=255)
+    salt: bytes = sqlmodel.Field(
+        sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
+    )
+    enc_login: bytes = sqlmodel.Field(
+        sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
+    )
+    enc_pass: bytes = sqlmodel.Field(
+        sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
+    )

src/database/prepare.py

@@ -1,35 +1,18 @@
-import mariadb
+import sqlmodel
+from sqlalchemy.future import Engine
+
+from . import models
 
 
-def _create_tables(con: mariadb.Connection) -> None:
+def get_engine(host: str, user: str, passwd: str, db: str) -> Engine:
-    cursor = con.cursor()
+    engine = sqlmodel.create_engine(
-    cursor.execute(
+        f"mariadb+mariadbconnector://{user}:{passwd}@{host}/{db}"
-        """CREATE TABLE IF NOT EXISTS master_pass (user_id INT,
-                        salt BINARY(64),
-                        passwd BINARY(64),
-                        PRIMARY KEY(user_id)
-                        )"""
     )
-    cursor.execute(
+    return engine
-        """CREATE TABLE IF NOT EXISTS accounts(user_id INT,
+
-                        acc_name VARCHAR(255),
+
-                        salt BINARY(64),
+def prepare(engine: Engine) -> None:
-                        enc_login BINARY(64),
+    sqlmodel.SQLModel.metadata.create_all(
-                        enc_pass BINARY(64),
+        engine,
-                        UNIQUE(acc_name, user_id)
+        # [models.Account, models.MasterPass]
-                        )"""
     )
-    cursor.close()
-
-
-def _create_index(con: mariadb.Connection) -> None:
-    cursor = con.cursor()
-    cursor.execute(
-        """CREATE INDEX IF NOT EXISTS user_id_to_acc on accounts(user_id)
-    """
-    )
-
-
-def prepare(con: mariadb.Connection) -> None:
-    _create_tables(con)
-    _create_index(con)