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