ai-admin-server/app/controller/add_user_route.py

from datetime import timedelta
from enum import Enum

from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer
from jwt import InvalidTokenError
from pydantic import BaseModel
from sqlmodel import select, Field
from starlette import status

from app.config.env import env
from app.model.BasicModel import BasicModel
from app.utils.CrpyUtils import CryptUtils
from app.utils.db_utils import AsyncSessionDep
from app.utils.next_id import next_id


class UserValidate(str, Enum):
  Y = 'Y'
  N = 'N'


# 公共的，也是最后返回给前端的一个用户信息数据类型
class PublicUser(BasicModel):
  username: str = Field(..., description="用户名")
  email: str = Field(..., description="邮箱")
  full_name: str = Field(..., description="用户全名")
  valid: UserValidate = Field(default=UserValidate.N, description="用户账号是否已经激活")


# 注册的时候，客户端传入的用户信息，需要包含这个明文密码字段
class RegistryUser(PublicUser):
  password: str


# 对pl_user表进行增删改查时的这个model类
class UserModel(PublicUser, table=True):
  __tablename__ = "pl_user"
  hash_password: str


class Token(BaseModel):
  token: str
  token_type: str


def add_user_route(app: FastAPI):
  @app.post("/registry")
  async def _registry(registry_user: RegistryUser, session: AsyncSessionDep):

    # /*---------------------------------------检查用户名是否已经注册-------------------------------------------*/
    query = select(UserModel).where(UserModel.username == registry_user.username)
    result = await session.execute(query)
    item_cls = result.scalars().first()

    if item_cls:
      return {"result": None, "error": f"用户名：{registry_user.username} 已经存在"}

    # /*---------------------------------------检查邮箱是否已经注册-------------------------------------------*/

    query = select(UserModel).where(UserModel.email == registry_user.email)
    result = await session.execute(query)
    item_cls = result.scalars().first()

    if item_cls:
      return {"result": None, "error": f"邮箱：{registry_user.email} 已经注册"}

    # /*---------------------------------------开始注册流程-------------------------------------------*/

    hash_password = CryptUtils.get_password_hash(registry_user.password)
    user = UserModel(
      username=registry_user.username,
      email=registry_user.email,
      full_name=registry_user.full_name,
      hash_password=hash_password,
      valid=UserValidate.N,
    )
    user.id = await next_id()
    session.add(user)
    await session.commit()
    await session.refresh(user)

    public_user = PublicUser(**user.model_dump())

    active_user_token = CryptUtils.create_access_token(public_user.username, expires_delta=timedelta(days=365 * 3))
    active_url = f"{env.server_domain}:{env.server_port}/verify?token={active_user_token}"

    return {
      "result": public_user,
      "active_url": active_url
    }

  @app.get("/verify")
  async def _verify(token: str, session: AsyncSessionDep):
    username = CryptUtils.get_username_from_token(token)

    if not username:
      return {"result": None, "error": "token无效或者已经过期"}

    query = select(UserModel).where(UserModel.username == username)
    result = await session.execute(query)
    item_cls: UserModel | None = result.scalars().first()

    if not item_cls:
      return {"result": None, "error": f"用户 {username} 不存在"}

    item_cls.valid = UserValidate.Y
    session.add(item_cls)
    await session.commit()
    await session.refresh(item_cls)

    public_user = PublicUser(**item_cls.model_dump())

    return {
      "result": public_user,
      "message": f"用户 {username} 激活成功"
    }

  @app.post("/login")
  @app.post("/token")
  async def _token(session: AsyncSessionDep, form_data: OAuth2PasswordRequestForm = Depends()):
    print("login", form_data)
    user = await authenticate_user(session, form_data.username, form_data.password)

    if not user:
      raise HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="用户名或者密码不正确",
        headers={"WWW-Authenticate": "Bearer"},
      )

    token = Token(
      token=CryptUtils.create_access_token(user.username),
      token_type="Bearer",
    )

    return {
      "result": user,
      "token": token,
    }

  @app.get("/users/me")
  async def _me(current_user: PublicUser = Depends(get_current_user)):
    return current_user

  @app.post("/order")
  async def _query_order(product_name: str, current_user: PublicUser = Depends(get_current_user)):
    return [product_name]


async def authenticate_user(session: AsyncSessionDep, username: str, password: str):
  query = select(UserModel).where(UserModel.username == username)
  result = await session.execute(query)
  item_cls: UserModel | None = result.scalars().first()
  if not item_cls:
    return None

  if item_cls.valid != UserValidate.Y:
    return None

  if not CryptUtils.verify_password(password, item_cls.hash_password):
    return None

  public_user = PublicUser(**item_cls.model_dump())

  return public_user


oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

unauthorized_exception = HTTPException(
  status_code=status.HTTP_401_UNAUTHORIZED,
  detail="The token is invalid or had expired",
  headers={"WWW-Authenticate": "Bearer"},
)


async def get_current_user(session: AsyncSessionDep, token: str = Depends(oauth2_scheme)):
  try:
    username = CryptUtils.get_username_from_token(token)
    if not username:
      raise unauthorized_exception
  except InvalidTokenError:
    raise unauthorized_exception

  user_model = await get_user_by_username(username, session)
  if not user_model:
    raise unauthorized_exception

  return PublicUser(**user_model.model_dump())


async def get_user_by_username(username: str, session: AsyncSessionDep):
  query = (
    select(UserModel)
    .where(UserModel.username == username)
    .where(UserModel.valid == UserValidate.Y)
  )
  result = await session.execute(query)
  item_cls: UserModel | None = result.scalars().first()
  return item_cls