217 lines
6.9 KiB
Python
217 lines
6.9 KiB
Python
from ast import Is
|
|
|
|
import httpx
|
|
from fastapi import APIRouter, Depends, Request, HTTPException, BackgroundTasks
|
|
from fastapi.responses import RedirectResponse
|
|
import os
|
|
import secrets
|
|
import json
|
|
from urllib.parse import quote
|
|
from dotenv import load_dotenv
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.db import SessionLocal, get_db
|
|
from app.models import User
|
|
from auth.tokens import create_jwt_token, create_refresh_token, decode_token
|
|
from app.utils.format import get_image_url
|
|
|
|
router = APIRouter()
|
|
|
|
load_dotenv()
|
|
|
|
GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
|
|
GOOGLE_CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
|
|
GOOGLE_REDIRECT_URI = os.getenv("GOOGLE_REDIRECT_URI")
|
|
|
|
GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
|
|
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
|
GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v2/userinfo"
|
|
|
|
FRONTEND_URL = os.getenv("FRONTEND_URL", "http://localhost:3000")
|
|
ENV_MODE = os.getenv("ENV_MODE")
|
|
IS_PROD = ENV_MODE == "prod"
|
|
|
|
# 在文件顶部获取环境变量
|
|
ACCESS_TOKEN_EXPIRE_MINUTES = int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "60"))
|
|
REFRESH_TOKEN_EXPIRE_DAYS = int(os.getenv("REFRESH_TOKEN_EXPIRE_DAYS", "7"))
|
|
|
|
# 统一转换为秒数（max_age 需要秒）
|
|
ACCESS_TOKEN_MAX_AGE = ACCESS_TOKEN_EXPIRE_MINUTES * 60
|
|
REFRESH_TOKEN_MAX_AGE = REFRESH_TOKEN_EXPIRE_DAYS * 24 * 60 * 60
|
|
|
|
# 定义独立的后台任务函数，用于处理头像转存到 S3 的耗时操作
|
|
async def async_upload_avatar_task(google_avatar_url: str, user_id: str):
|
|
"""
|
|
后台任务函数：下载 Google 头像并上传到 S3，然后更新用户记录。
|
|
因为主接口的 db 已经关闭，这里需要独立管理 db 的生命周期。
|
|
"""
|
|
from app.utils.s3 import upload_google_avatar_to_s3
|
|
|
|
# 如果 Google 本身就没给头像，无需转存
|
|
if not google_avatar_url:
|
|
return
|
|
|
|
try:
|
|
# 执行耗时的网络 I/O：下载并上传到 S3
|
|
s3_avatar_url = await upload_google_avatar_to_s3(google_avatar_url, user_id)
|
|
|
|
# 成功后，开启一个独立的数据库连接写回数据
|
|
db_task = SessionLocal()
|
|
try:
|
|
user = db_task.query(User).filter(User.id == user_id).first()
|
|
if user:
|
|
user.avatar = s3_avatar_url
|
|
db_task.commit()
|
|
print(f"成功在后台为用户 {user_id} 转存头像到 S3")
|
|
finally:
|
|
db_task.close() # 必须手动关闭连接
|
|
|
|
except Exception as s3_err:
|
|
# 如果 S3 依然失败，打印错误日志，但此时绝不会打扰到已经登录的用户
|
|
print(f"Background S3 Upload failed for user {user_id}: {s3_err}")
|
|
|
|
|
|
@router.get("/login/google")
|
|
def login_google():
|
|
state = secrets.token_urlsafe(16)
|
|
|
|
url = (
|
|
f"{GOOGLE_AUTH_URL}"
|
|
f"?client_id={GOOGLE_CLIENT_ID}"
|
|
f"&redirect_uri={GOOGLE_REDIRECT_URI}"
|
|
f"&response_type=code"
|
|
f"&scope=openid email profile"
|
|
f"&state={state}"
|
|
)
|
|
|
|
response = RedirectResponse(url)
|
|
response.set_cookie(
|
|
"oauth_state",
|
|
state,
|
|
httponly=True,
|
|
secure=True,
|
|
samesite="lax"
|
|
)
|
|
return response
|
|
|
|
@router.get("/auth/callback")
|
|
async def auth_callback(
|
|
request: Request,
|
|
code: str,
|
|
state: str,
|
|
background_tasks: BackgroundTasks,
|
|
db: Session = Depends(get_db)
|
|
):
|
|
cookie_state = request.cookies.get("oauth_state")
|
|
|
|
if not cookie_state or cookie_state != state:
|
|
raise HTTPException(status_code=400, detail="Invalid state")
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
token_res = await client.post(
|
|
GOOGLE_TOKEN_URL,
|
|
data={
|
|
"code": code,
|
|
"client_id": GOOGLE_CLIENT_ID,
|
|
"client_secret": GOOGLE_CLIENT_SECRET,
|
|
"redirect_uri": GOOGLE_REDIRECT_URI,
|
|
"grant_type": "authorization_code",
|
|
},
|
|
)
|
|
|
|
token_json = token_res.json()
|
|
access_token = token_json["access_token"]
|
|
|
|
user_res = await client.get(
|
|
GOOGLE_USERINFO_URL,
|
|
headers={"Authorization": f"Bearer {access_token}"},
|
|
)
|
|
|
|
user_info = user_res.json()
|
|
|
|
user = db.query(User).filter(User.email == user_info["email"]).first()
|
|
|
|
if not user:
|
|
# 提取邮箱前缀作为 handle 基础
|
|
# 比如 "john.doe@gmail.com" -> "john_doe"
|
|
base_handle = user_info["email"].split("@")[0].replace(".", "_").lower()
|
|
|
|
# 拼接一个随机短字符串，确保 handle 的唯一性
|
|
import secrets
|
|
random_suffix = secrets.token_hex(3) # 生成 6 位随机字符
|
|
suggested_handle = f"{base_handle}_{random_suffix}"
|
|
|
|
# 如果 Google 没有头像，使用 DiceBear 头像生成器作为兜底
|
|
google_avatar = user_info.get("picture")
|
|
if not google_avatar:
|
|
default_avatar_url = f"https://api.dicebear.com/7.x/lorelei/svg?seed={base_handle}"
|
|
else:
|
|
default_avatar_url = google_avatar
|
|
|
|
user = User(
|
|
google_id=user_info["id"],
|
|
email=user_info["email"],
|
|
name=user_info["name"],
|
|
avatar=default_avatar_url,
|
|
handle=suggested_handle,
|
|
)
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
|
|
# 异步/同步将头像转存到 S3
|
|
# FastAPI 会在把下方的 return 响应发送给浏览器之后，立刻在后台启动这个任务
|
|
if google_avatar:
|
|
background_tasks.add_task(
|
|
async_upload_avatar_task,
|
|
google_avatar,
|
|
str(user.id)
|
|
)
|
|
|
|
# 这里不需要再次 refresh，除非后面还要用到更新后的 user 对象
|
|
# db.refresh(user)
|
|
|
|
jwt_token = create_jwt_token({
|
|
"sub": user.id,
|
|
"email": user.email,
|
|
})
|
|
|
|
refresh_token = create_refresh_token({
|
|
"sub": user.id,
|
|
"email": user.email,
|
|
})
|
|
|
|
response = RedirectResponse(url=f"{FRONTEND_URL}")
|
|
|
|
response.set_cookie(
|
|
key="access_token",
|
|
value=jwt_token,
|
|
httponly=True,
|
|
max_age=ACCESS_TOKEN_MAX_AGE,
|
|
samesite="lax",
|
|
secure=IS_PROD,
|
|
)
|
|
|
|
response.set_cookie(
|
|
key="refresh_token",
|
|
value=refresh_token,
|
|
httponly=True,
|
|
max_age=REFRESH_TOKEN_MAX_AGE,
|
|
samesite="lax",
|
|
secure=IS_PROD,
|
|
)
|
|
|
|
response.set_cookie(
|
|
key="user_info",
|
|
value=quote(json.dumps({
|
|
"email": user.email,
|
|
"name": user.name,
|
|
"avatar": user.avatar
|
|
})),
|
|
max_age=ACCESS_TOKEN_MAX_AGE,
|
|
httponly=False
|
|
)
|
|
|
|
response.delete_cookie("oauth_state")
|
|
|
|
return response |