PassManager/src/database/add.py

import sqlmodel
from sqlalchemy.exc import IntegrityError
from sqlalchemy.future import Engine

from . import models


def add_account(
    engine: Engine,
    user_id: int,
    name: str,
    salt: bytes,
    enc_login: bytes,
    enc_password: bytes,
) -> bool:
    """Adds account to the database. Returns true on success,
    false otherwise"""
    account = models.Account(
        user_id=user_id,
        name=name,
        salt=salt,
        enc_login=enc_login,
        enc_password=enc_password,
    )
    try:
        with sqlmodel.Session(engine) as session:
            session.add(account)
            session.commit()
    except IntegrityError:
        return False
    else:
        return True


def add_master_pass(
    engine: Engine,
    user_id: int,
    salt: bytes,
    password_hash: bytes,
) -> bool:
    """Adds master password the database. Returns true on success,
    false otherwise"""
    master_pass = models.MasterPass(
        user_id=user_id,
        salt=salt,
        password_hash=password_hash,
    )
    try:
        with sqlmodel.Session(engine) as session:
            session.add(master_pass)
            session.commit()
    except IntegrityError:
        return False
    else:
        return True