281 lines
8.5 KiB
Python
281 lines
8.5 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 app.core.context import current_user_dept_id
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.core.config import settings
|
||
from app.core.database import get_db
|
||
from app.crud.user import user_crud
|
||
from app.core.utils import generate_uuid
|
||
from app.services.web_auth_roles import resolve_web_role_for_user
|
||
|
||
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 续期；已登录用户每次 API 校验通过后滑动续期）"""
|
||
if not session_id:
|
||
return
|
||
r = _get_redis_client()
|
||
if r is None:
|
||
return
|
||
try:
|
||
await asyncio.wait_for(
|
||
r.expire(_session_key(session_id), REFRESH_TOKEN_EXPIRE_SECONDS),
|
||
timeout=3.0,
|
||
)
|
||
except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
|
||
pass
|
||
|
||
|
||
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无效或已过期",
|
||
)
|
||
|
||
# 方案 C：有操作时延长 Redis 会话，连续 7 天无请求则会话过期
|
||
await extend_session(session_id)
|
||
|
||
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 getattr(user, "del_flag", 0) != 0:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="用户不存在或已删除",
|
||
)
|
||
|
||
if getattr(user, "status", "1") != "1":
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="用户已被禁用",
|
||
)
|
||
|
||
web_role = resolve_web_role_for_user(db, user.id)
|
||
|
||
phone = (getattr(user, "phone_no", None) or "").strip() or None
|
||
email = (getattr(user, "email", None) or "").strip() or None
|
||
avatar = getattr(user, "avatar", None)
|
||
|
||
current_user_dept_id.set(user.dept_id or "")
|
||
|
||
tenant_no = (getattr(user, "tenant_no", None) or "").strip() or settings.DEFAULT_TENANT_NO
|
||
|
||
return {
|
||
"user_id": user.id,
|
||
"username": user.user_no,
|
||
"real_name": user.nick_name,
|
||
"role": web_role,
|
||
"region_id": user.dept_id,
|
||
"tenant_no": tenant_no,
|
||
"phone": phone,
|
||
"email": email,
|
||
"avatar_url": avatar,
|
||
"is_active": user.status == "1",
|
||
}
|
||
|
||
|
||
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="无效的认证凭据",
|
||
)
|
||
current_user_dept_id.set(payload.get("region_id") or "")
|
||
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"
|