auth.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255
  1. """认证模块"""
  2. import asyncio
  3. from datetime import datetime, timedelta
  4. from typing import Optional
  5. import redis.asyncio as redis
  6. from redis import exceptions as redis_exceptions
  7. from fastapi import Depends, HTTPException, status, Request
  8. from fastapi.security import OAuth2PasswordBearer
  9. from jose import JWTError, jwt
  10. from sqlalchemy.orm import Session
  11. from app.core.config import settings
  12. from app.core.database import get_db
  13. from app.core.security import verify_password
  14. from app.models.models import SysUser
  15. from app.crud.user import user_crud
  16. from app.core.utils import generate_uuid
  17. oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/public/login")
  18. REFRESH_TOKEN_EXPIRE_DAYS = 7
  19. REFRESH_TOKEN_EXPIRE_SECONDS = int(timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS).total_seconds())
  20. _redis_client: Optional[redis.Redis] = None
  21. def _redis_session_enabled() -> bool:
  22. # 只要配置了 REDIS_HOST,就启用“会话校验/refreshToken 吊销”
  23. return bool(settings.REDIS_HOST)
  24. def _get_redis_client() -> Optional[redis.Redis]:
  25. global _redis_client
  26. if not _redis_session_enabled():
  27. return None
  28. if _redis_client is None:
  29. _redis_client = redis.Redis(
  30. host=settings.REDIS_HOST,
  31. port=settings.REDIS_PORT,
  32. db=settings.REDIS_DB,
  33. password=settings.REDIS_PASSWORD or None,
  34. socket_timeout=settings.REDIS_TIMEOUT_SEC,
  35. ssl=settings.REDIS_SSL_ENABLED,
  36. decode_responses=True,
  37. )
  38. return _redis_client
  39. def _session_key(session_id: str) -> str:
  40. return f"auth:sessions:{session_id}"
  41. def _refresh_key(refresh_jti: str) -> str:
  42. return f"auth:refresh:{refresh_jti}"
  43. async def store_session(session_id: str, user_id: str) -> None:
  44. """写入会话(用于强制失效)"""
  45. r = _get_redis_client()
  46. if r is None:
  47. return
  48. await r.set(_session_key(session_id), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS)
  49. async def revoke_session(session_id: str) -> None:
  50. """撤销会话(登出/强制失效)"""
  51. r = _get_redis_client()
  52. if r is None:
  53. return
  54. await r.delete(_session_key(session_id))
  55. async def extend_session(session_id: str) -> None:
  56. """刷新会话 TTL(用于 refreshToken 延长会话)"""
  57. r = _get_redis_client()
  58. if r is None:
  59. return
  60. await r.expire(_session_key(session_id), REFRESH_TOKEN_EXPIRE_SECONDS)
  61. async def is_session_valid(session_id: str) -> bool:
  62. """检查会话是否仍然有效"""
  63. r = _get_redis_client()
  64. if r is None:
  65. return True
  66. if not session_id:
  67. return False
  68. try:
  69. # 添加连接超时保护,Redis 不可用时降级为验证通过
  70. return await asyncio.wait_for(r.exists(_session_key(session_id)), timeout=3.0) > 0
  71. except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
  72. # Redis 连接失败时降级:只验证 JWT,不过滤会话吊销
  73. return True
  74. async def store_refresh_token(refresh_jti: str, user_id: str) -> None:
  75. """写入 refreshToken(用于 one-time 使用/吊销)"""
  76. r = _get_redis_client()
  77. if r is None:
  78. return
  79. await r.set(_refresh_key(refresh_jti), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS)
  80. async def is_refresh_token_valid(refresh_jti: str) -> bool:
  81. r = _get_redis_client()
  82. if r is None:
  83. return True
  84. if not refresh_jti:
  85. return False
  86. try:
  87. return await asyncio.wait_for(r.exists(_refresh_key(refresh_jti)), timeout=3.0) > 0
  88. except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
  89. return True
  90. async def revoke_refresh_token(refresh_jti: str) -> None:
  91. r = _get_redis_client()
  92. if r is None:
  93. return
  94. await r.delete(_refresh_key(refresh_jti))
  95. def create_access_token(
  96. data: dict,
  97. session_id: str,
  98. token_jti: str,
  99. expires_delta: Optional[timedelta] = None,
  100. ) -> str:
  101. """创建JWT访问令牌"""
  102. to_encode = data.copy()
  103. if expires_delta:
  104. expire = datetime.utcnow() + expires_delta
  105. else:
  106. expire = datetime.utcnow() + timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
  107. to_encode.update(
  108. {
  109. "exp": expire,
  110. "type": "access",
  111. "sid": session_id,
  112. "jti": token_jti,
  113. }
  114. )
  115. encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
  116. return encoded_jwt
  117. def create_refresh_token(data: dict, session_id: str, refresh_jti: str) -> str:
  118. """创建刷新令牌"""
  119. to_encode = data.copy()
  120. expire = datetime.utcnow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS) # 有效期与 Redis TTL 保持一致
  121. to_encode.update(
  122. {
  123. "exp": expire,
  124. "type": "refresh",
  125. "sid": session_id,
  126. "jti": refresh_jti,
  127. }
  128. )
  129. encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
  130. return encoded_jwt
  131. def decode_token(token: str) -> dict:
  132. """解码JWT令牌"""
  133. try:
  134. payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
  135. return payload
  136. except JWTError:
  137. raise HTTPException(
  138. status_code=status.HTTP_401_UNAUTHORIZED,
  139. detail="无效的认证凭据",
  140. headers={"WWW-Authenticate": "Bearer"},
  141. )
  142. async def get_current_user_by_token(
  143. token: str,
  144. db: Session,
  145. ) -> dict:
  146. """根据 token 获取当前用户(包含 Redis 会话校验)"""
  147. payload = decode_token(token)
  148. user_id: str = payload.get("sub")
  149. session_id: str = payload.get("sid")
  150. if not await is_session_valid(session_id):
  151. raise HTTPException(
  152. status_code=status.HTTP_401_UNAUTHORIZED,
  153. detail="Token无效或已过期",
  154. )
  155. if user_id is None:
  156. raise HTTPException(
  157. status_code=status.HTTP_401_UNAUTHORIZED,
  158. detail="无效的认证凭据",
  159. )
  160. user = user_crud.get(db, user_id)
  161. if not user or user.is_del == 1:
  162. raise HTTPException(
  163. status_code=status.HTTP_401_UNAUTHORIZED,
  164. detail="用户不存在或已删除",
  165. )
  166. if not user.is_active:
  167. raise HTTPException(
  168. status_code=status.HTTP_401_UNAUTHORIZED,
  169. detail="用户已被禁用",
  170. )
  171. return {
  172. "user_id": user.id,
  173. "username": user.username,
  174. "real_name": user.real_name,
  175. "role": user.role,
  176. "region_id": user.region_id,
  177. "phone": getattr(user, "phone", None),
  178. "avatar_url": getattr(user, "avatar_url", None),
  179. "is_active": getattr(user, "is_active", True),
  180. }
  181. async def get_current_user(
  182. token: str = Depends(oauth2_scheme),
  183. db: Session = Depends(get_db)
  184. ) -> dict:
  185. """获取当前用户"""
  186. return await get_current_user_by_token(token, db)
  187. async def get_current_user_simple(token: str = Depends(oauth2_scheme)) -> dict:
  188. """简化版获取当前用户(不查数据库)"""
  189. payload = decode_token(token)
  190. user_id: str = payload.get("sub")
  191. if user_id is None:
  192. raise HTTPException(
  193. status_code=status.HTTP_401_UNAUTHORIZED,
  194. detail="无效的认证凭据",
  195. )
  196. return {
  197. "user_id": user_id,
  198. "username": payload.get("username"),
  199. "role": payload.get("role"),
  200. "region_id": payload.get("region_id"),
  201. "real_name": payload.get("real_name")
  202. }
  203. def get_client_ip(request: Request) -> str:
  204. """获取客户端IP"""
  205. forwarded = request.headers.get("X-Forwarded-For")
  206. if forwarded:
  207. return forwarded.split(",")[0].strip()
  208. return request.client.host if request.client else "unknown"