256 lines
7.6 KiB
Python
256 lines
7.6 KiB
Python
"""认证模块"""
|
|
import asyncio
|
|
from datetime import datetime, timedelta
|
|
from typing import Optional
|
|
|
|
import redis.asyncio as redis
|
|
from redis import exceptions as redis_exceptions
|
|
from fastapi import Depends, HTTPException, status, Request
|
|
from fastapi.security import OAuth2PasswordBearer
|
|
from jose import JWTError, jwt
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.config import settings
|
|
from app.core.database import get_db
|
|
from app.core.security import verify_password
|
|
from app.models.models import SysUser
|
|
from app.crud.user import user_crud
|
|
from app.core.utils import generate_uuid
|
|
|
|
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/public/login")
|
|
|
|
REFRESH_TOKEN_EXPIRE_DAYS = 7
|
|
REFRESH_TOKEN_EXPIRE_SECONDS = int(timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS).total_seconds())
|
|
|
|
_redis_client: Optional[redis.Redis] = None
|
|
|
|
|
|
def _redis_session_enabled() -> bool:
|
|
# 只要配置了 REDIS_HOST，就启用“会话校验/refreshToken 吊销”
|
|
return bool(settings.REDIS_HOST)
|
|
|
|
|
|
def _get_redis_client() -> Optional[redis.Redis]:
|
|
global _redis_client
|
|
if not _redis_session_enabled():
|
|
return None
|
|
if _redis_client is None:
|
|
_redis_client = redis.Redis(
|
|
host=settings.REDIS_HOST,
|
|
port=settings.REDIS_PORT,
|
|
db=settings.REDIS_DB,
|
|
password=settings.REDIS_PASSWORD or None,
|
|
socket_timeout=settings.REDIS_TIMEOUT_SEC,
|
|
ssl=settings.REDIS_SSL_ENABLED,
|
|
decode_responses=True,
|
|
)
|
|
return _redis_client
|
|
|
|
|
|
def _session_key(session_id: str) -> str:
|
|
return f"auth:sessions:{session_id}"
|
|
|
|
|
|
def _refresh_key(refresh_jti: str) -> str:
|
|
return f"auth:refresh:{refresh_jti}"
|
|
|
|
|
|
async def store_session(session_id: str, user_id: str) -> None:
|
|
"""写入会话（用于强制失效）"""
|
|
r = _get_redis_client()
|
|
if r is None:
|
|
return
|
|
await r.set(_session_key(session_id), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS)
|
|
|
|
|
|
async def revoke_session(session_id: str) -> None:
|
|
"""撤销会话（登出/强制失效）"""
|
|
r = _get_redis_client()
|
|
if r is None:
|
|
return
|
|
await r.delete(_session_key(session_id))
|
|
|
|
|
|
async def extend_session(session_id: str) -> None:
|
|
"""刷新会话 TTL（用于 refreshToken 延长会话）"""
|
|
r = _get_redis_client()
|
|
if r is None:
|
|
return
|
|
await r.expire(_session_key(session_id), REFRESH_TOKEN_EXPIRE_SECONDS)
|
|
|
|
|
|
async def is_session_valid(session_id: str) -> bool:
|
|
"""检查会话是否仍然有效"""
|
|
r = _get_redis_client()
|
|
if r is None:
|
|
return True
|
|
if not session_id:
|
|
return False
|
|
try:
|
|
# 添加连接超时保护，Redis 不可用时降级为验证通过
|
|
return await asyncio.wait_for(r.exists(_session_key(session_id)), timeout=3.0) > 0
|
|
except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
|
|
# Redis 连接失败时降级：只验证 JWT，不过滤会话吊销
|
|
return True
|
|
|
|
|
|
async def store_refresh_token(refresh_jti: str, user_id: str) -> None:
|
|
"""写入 refreshToken（用于 one-time 使用/吊销）"""
|
|
r = _get_redis_client()
|
|
if r is None:
|
|
return
|
|
await r.set(_refresh_key(refresh_jti), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS)
|
|
|
|
|
|
async def is_refresh_token_valid(refresh_jti: str) -> bool:
|
|
r = _get_redis_client()
|
|
if r is None:
|
|
return True
|
|
if not refresh_jti:
|
|
return False
|
|
try:
|
|
return await asyncio.wait_for(r.exists(_refresh_key(refresh_jti)), timeout=3.0) > 0
|
|
except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
|
|
return True
|
|
|
|
|
|
async def revoke_refresh_token(refresh_jti: str) -> None:
|
|
r = _get_redis_client()
|
|
if r is None:
|
|
return
|
|
await r.delete(_refresh_key(refresh_jti))
|
|
|
|
|
|
def create_access_token(
|
|
data: dict,
|
|
session_id: str,
|
|
token_jti: str,
|
|
expires_delta: Optional[timedelta] = None,
|
|
) -> str:
|
|
"""创建JWT访问令牌"""
|
|
to_encode = data.copy()
|
|
if expires_delta:
|
|
expire = datetime.utcnow() + expires_delta
|
|
else:
|
|
expire = datetime.utcnow() + timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
|
|
to_encode.update(
|
|
{
|
|
"exp": expire,
|
|
"type": "access",
|
|
"sid": session_id,
|
|
"jti": token_jti,
|
|
}
|
|
)
|
|
encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
|
|
return encoded_jwt
|
|
|
|
|
|
def create_refresh_token(data: dict, session_id: str, refresh_jti: str) -> str:
|
|
"""创建刷新令牌"""
|
|
to_encode = data.copy()
|
|
expire = datetime.utcnow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS) # 有效期与 Redis TTL 保持一致
|
|
to_encode.update(
|
|
{
|
|
"exp": expire,
|
|
"type": "refresh",
|
|
"sid": session_id,
|
|
"jti": refresh_jti,
|
|
}
|
|
)
|
|
encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
|
|
return encoded_jwt
|
|
|
|
|
|
def decode_token(token: str) -> dict:
|
|
"""解码JWT令牌"""
|
|
try:
|
|
payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
|
|
return payload
|
|
except JWTError:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="无效的认证凭据",
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
)
|
|
|
|
|
|
async def get_current_user_by_token(
|
|
token: str,
|
|
db: Session,
|
|
) -> dict:
|
|
"""根据 token 获取当前用户（包含 Redis 会话校验）"""
|
|
payload = decode_token(token)
|
|
|
|
user_id: str = payload.get("sub")
|
|
session_id: str = payload.get("sid")
|
|
|
|
if not await is_session_valid(session_id):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Token无效或已过期",
|
|
)
|
|
|
|
if user_id is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="无效的认证凭据",
|
|
)
|
|
|
|
user = user_crud.get(db, user_id)
|
|
if not user or user.is_del == 1:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="用户不存在或已删除",
|
|
)
|
|
|
|
if not user.is_active:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="用户已被禁用",
|
|
)
|
|
|
|
return {
|
|
"user_id": user.id,
|
|
"username": user.username,
|
|
"real_name": user.real_name,
|
|
"role": user.role,
|
|
"region_id": user.region_id,
|
|
"phone": getattr(user, "phone", None),
|
|
"avatar_url": getattr(user, "avatar_url", None),
|
|
"is_active": getattr(user, "is_active", True),
|
|
}
|
|
|
|
|
|
async def get_current_user(
|
|
token: str = Depends(oauth2_scheme),
|
|
db: Session = Depends(get_db)
|
|
) -> dict:
|
|
"""获取当前用户"""
|
|
return await get_current_user_by_token(token, db)
|
|
|
|
|
|
async def get_current_user_simple(token: str = Depends(oauth2_scheme)) -> dict:
|
|
"""简化版获取当前用户（不查数据库）"""
|
|
payload = decode_token(token)
|
|
user_id: str = payload.get("sub")
|
|
if user_id is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="无效的认证凭据",
|
|
)
|
|
return {
|
|
"user_id": user_id,
|
|
"username": payload.get("username"),
|
|
"role": payload.get("role"),
|
|
"region_id": payload.get("region_id"),
|
|
"real_name": payload.get("real_name")
|
|
}
|
|
|
|
|
|
def get_client_ip(request: Request) -> str:
|
|
"""获取客户端IP"""
|
|
forwarded = request.headers.get("X-Forwarded-For")
|
|
if forwarded:
|
|
return forwarded.split(",")[0].strip()
|
|
return request.client.host if request.client else "unknown"
|