app/models.py - master - durovcode/backend

app/models.py
188 строк4 KB
quyxishi
implemented Gemini support
22 сен 2024, 02:05
22 сен 2024, 02:051ff8dc3
import enum
import secrets
from base64 import b64decode
from hashlib import sha512
from typing import Optional, Self

from argon2 import PasswordHasher as Argon2Hash
from pydantic import BaseModel, EmailStr, Field
from sqlalchemy import Column, Enum, Integer, String, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm.decl_api import DeclarativeMeta


# FastAPI : Models


# * /account


class RoleEnum(str, enum.Enum):
    student: str = 'student'
    lecturer: str = 'lecturer'


class StatusModel(BaseModel):
    ok: bool
    detail: Optional[str] = None


class ErrorModel(StatusModel):
    ok: bool = False
    detail: str


class UserAuthorizationForm(BaseModel, str_strip_whitespace=True):
    email: EmailStr = Field(min_length=4, max_length=64)
    passwd: str = Field(min_length=8, max_length=32)


class UserVerificationResponse(StatusModel):
    token: Optional[str]


class UserAuthorizationResponse(UserVerificationResponse):
    otp: bool


class UserVerificationForm(BaseModel, str_strip_whitespace=True):
    otp: str = Field(min_length=6, max_length=6, pattern=r'^\d+$')


class UserCreationForm(UserAuthorizationForm):
    name: str = Field(
        min_length=4,
        max_length=32,
        pattern='^[А-яё]+ [А-яё]+$',
        examples=['Иван Иванов'],
    )
    role: RoleEnum = Field()


class UserCreationResponse(StatusModel):
    pass


class UserDeletionResponse(StatusModel):
    pass


class UserMeResponse(StatusModel):
    name: str
    profile_photo: Optional[str]
    role: RoleEnum


class UserAssignOtpResponse(UserVerificationResponse):
    key: str
    uri: str


class UserAssignOtpVerifyForm(UserVerificationForm):
    pass


class UserAssignOtpVerifyResponse(StatusModel):
    pass


class UserAssignAvatarResponse(StatusModel):
    profile_photo: str


# * /ai


class AiQueryResponse(StatusModel):
    original: str
    report_docx: str


# SQLAlchemy : Models


AlchemyBase: DeclarativeMeta = declarative_base()


class Users(AlchemyBase):
    __tablename__: str = 'users'

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    email = Column(String, unique=True)
    name = Column(String)
    passwd_hash = Column(String)
    salt = Column(String)
    otp_key = Column(String, nullable=True)
    profile_photo = Column(String, nullable=True)
    role = Column(Enum(RoleEnum))

    @classmethod
    async def authorize(
        cls, session: AsyncSession, email: str, passwd: str
    ) -> Optional[Self]:
        user: Optional[Users] = await session.scalar(
            statement=select(Users).where(Users.email == email).limit(limit=1)
        )

        if user is None:
            return False

        argon2 = Argon2Hash()

        passwd_hash: str = argon2.hash(
            password=sha512(string=passwd.encode()).hexdigest(),
            salt=b64decode(s=user.salt),
        )

        return user if secrets.compare_digest(user.passwd_hash, passwd_hash) else None

    @classmethod
    async def select(cls, session: AsyncSession, id: int) -> Optional[Self]:
        return await session.scalar(
            statement=select(cls).where(cls.id == id).order_by(cls.id)
        )

    @classmethod
    async def create(
        cls,
        session: AsyncSession,
        email: str,
        name: str,
        passwd_hash: str,
        salt: str,
        role: RoleEnum,
    ) -> Optional[Self]:
        user_entry = Users(
            email=email, name=name, passwd_hash=passwd_hash, salt=salt, role=role
        )

        session.add(instance=user_entry)
        await session.flush()

        return await cls.select(session=session, id=user_entry.id)

    @classmethod
    async def is_exists(cls, session: AsyncSession, email: str) -> bool:
        return (
            await session.scalar(statement=select(cls).where(cls.email == email))
        ) is not None

    @classmethod
    async def delete(cls, session: AsyncSession, user: Self) -> None:
        await session.delete(instance=user)
        await session.flush()

    async def update(self, session: AsyncSession, **kwargs: str) -> None:
        for k, v in kwargs.items():
            self.__setattr__(k, v)

        await session.flush()


# GigaChat : Models


# class GigaAccessTokenResponse(BaseModel):
#     access_token: str
#     expires_at: int