完善牧安云哨-后端

bbit_ai/app/config/security.py

@@ -1,13 +1,16 @@
-import jwt
-from jwt import PyJWTError
 from uuid import UUID
-from fastapi import Header, HTTPException, Depends
 
-JWT_SECRET =  "secret_jwt"
+import jwt
+from fastapi import Header, HTTPException
+from jwt import PyJWTError
+from starlette.websockets import WebSocketDisconnect
+
+JWT_SECRET = "secret_jwt"
 JWT_ALGORITHM = "HS256"
 JWT_AUDIENCE = "snowflake-ink"
 JWT_ISSUER = "https://snowflake.ink/"
 
+
 def get_user_id_from_token(token: str = Header(..., alias="Authorization")) -> UUID:
     """
     从 Authorization 头解析 token，并返回 user_id
@@ -24,7 +27,7 @@ def get_user_id_from_token(token: str = Header(..., alias="Authorization")) -> U
             JWT_SECRET,
             algorithms=[JWT_ALGORITHM],
             audience=JWT_AUDIENCE,
-            issuer=JWT_ISSUER
+            issuer=JWT_ISSUER,
         )
     except PyJWTError:
         raise HTTPException(status_code=401, detail="Token is missing or invalid")
@@ -36,4 +39,28 @@ def get_user_id_from_token(token: str = Header(..., alias="Authorization")) -> U
     if not user_id:
         raise HTTPException(status_code=401, detail="User ID not found in token")
 
-    return UUID(user_id)
+    return UUID(user_id)
+
+
+def get_user_id_from_token_from_ws(token: str) -> UUID:
+    if token.startswith("Bearer "):
+        token = token[7:]
+    try:
+        payload = jwt.decode(
+            token,
+            JWT_SECRET,
+            algorithms=[JWT_ALGORITHM],
+            audience=JWT_AUDIENCE,
+            issuer=JWT_ISSUER,
+        )
+    except PyJWTError:
+        raise WebSocketDisconnect()  # token 无效就断开
+
+    if payload.get("token_type") != "access_token":
+        raise WebSocketDisconnect()
+
+    user_id = payload.get("user_id")
+    if not user_id:
+        raise WebSocketDisconnect()
+
+    return UUID(user_id)