From 244c884d0a7c54c4a37de18c1c2c8ff15a506ff7 Mon Sep 17 00:00:00 2001 From: zhangqian <zhangqian@123.com> Date: 星期二, 15 十月 2024 00:11:31 +0800 Subject: [PATCH] 智能体列表接口,智能体会话记录接口 --- app/api/__init__.py | 67 +++++++++++++++++++++++++++++++++ 1 files changed, 66 insertions(+), 1 deletions(-) diff --git a/app/api/__init__.py b/app/api/__init__.py index bcd5c2a..8bd4579 100644 --- a/app/api/__init__.py +++ b/app/api/__init__.py @@ -1,10 +1,75 @@ -from fastapi import FastAPI +import jwt +from fastapi import FastAPI, Depends, HTTPException +from fastapi.security import OAuth2PasswordBearer +from passlib.context import CryptContext from pydantic import BaseModel +from starlette import status +from starlette.websockets import WebSocket, WebSocketDisconnect + +from app.models.user_model import UserModel +from app.service.auth import SECRET_KEY, ALGORITHM app = FastAPI() + +pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") +oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") class Response(BaseModel): code: int = 200 msg: str = "" data: dict = {} + + +class ResponseList(BaseModel): + code: int = 200 + msg: str = "" + data: list[dict] = [] + + +def get_current_user(token: str = Depends(oauth2_scheme)): + try: + payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM]) + username: str = payload.get("sub") + if username is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="鏃犳硶楠岃瘉鍑瘉", + headers={"WWW-Authenticate": "Bearer"}, + ) + user = UserModel(username=username, id=payload.get("user_id")) + if user.id == 0: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="鐢ㄦ埛涓嶅瓨鍦�", + headers={"WWW-Authenticate": "Bearer"}, + ) + return user + except jwt.PyJWTError: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="浠ょ墝鏃犳晥鎴栧凡杩囨湡", + headers={"WWW-Authenticate": "Bearer"}, + ) + + +async def get_current_user_websocket(websocket: WebSocket): + token = websocket.query_params.get('token') + if token is None: + await websocket.close(code=1008) + raise WebSocketDisconnect(code=status.WS_1008_POLICY_VIOLATION) + try: + payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM]) + username: str = payload.get("sub") + if username is None: + await websocket.close(code=1008) + raise WebSocketDisconnect(code=status.WS_1008_POLICY_VIOLATION) + user = UserModel(username=username, id=payload.get("user_id")) + if user is None: + await websocket.close(code=1008) + raise WebSocketDisconnect(code=status.WS_1008_POLICY_VIOLATION) + return user + except jwt.PyJWTError as e: + print(e) + await websocket.close(code=1008) + raise WebSocketDisconnect(code=status.WS_1008_POLICY_VIOLATION) \ No newline at end of file -- Gitblit v1.8.0