auth.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238
  1. """认证模块"""
  2. from datetime import datetime, timedelta
  3. from typing import Optional
  4. from redis import exceptions as redis_exceptions
  5. from fastapi import Depends, HTTPException, status, Request
  6. from fastapi.security import OAuth2PasswordBearer
  7. from jose import JWTError, jwt
  8. from app.core.context import current_user_dept_id
  9. from sqlalchemy.orm import Session
  10. from app.core.config import settings
  11. from app.core.database import get_db
  12. from app.core.redis_client import get_sync_redis_client
  13. from app.crud.user import user_crud
  14. from app.services.web_auth_roles import resolve_web_role_for_user
  15. oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/public/login")
  16. REFRESH_TOKEN_EXPIRE_DAYS = 7
  17. REFRESH_TOKEN_EXPIRE_SECONDS = int(timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS).total_seconds())
  18. def _session_key(session_id: str) -> str:
  19. return f"auth:sessions:{session_id}"
  20. def _refresh_key(refresh_jti: str) -> str:
  21. return f"auth:refresh:{refresh_jti}"
  22. def _redis_call_safe(fn):
  23. """同步 Redis 调用;失败时返回 None 供调用方降级。"""
  24. r = get_sync_redis_client()
  25. if r is None:
  26. return None
  27. try:
  28. return fn(r)
  29. except (redis_exceptions.TimeoutError, redis_exceptions.ConnectionError, redis_exceptions.RedisError):
  30. return None
  31. def store_session(session_id: str, user_id: str) -> None:
  32. """写入会话(用于强制失效)"""
  33. _redis_call_safe(lambda r: r.set(_session_key(session_id), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS))
  34. def revoke_session(session_id: str) -> None:
  35. """撤销会话(登出/强制失效)"""
  36. _redis_call_safe(lambda r: r.delete(_session_key(session_id)))
  37. def extend_session(session_id: str) -> None:
  38. """刷新会话 TTL(已登录用户每次 API 校验通过后滑动续期)"""
  39. if not session_id:
  40. return
  41. _redis_call_safe(lambda r: r.expire(_session_key(session_id), REFRESH_TOKEN_EXPIRE_SECONDS))
  42. def is_session_valid(session_id: str) -> bool:
  43. """检查会话是否仍然有效"""
  44. if get_sync_redis_client() is None:
  45. return True
  46. if not session_id:
  47. return False
  48. exists = _redis_call_safe(lambda r: r.exists(_session_key(session_id)))
  49. if exists is None:
  50. return True
  51. return exists > 0
  52. def store_refresh_token(refresh_jti: str, user_id: str) -> None:
  53. """写入 refreshToken(用于 one-time 使用/吊销)"""
  54. _redis_call_safe(
  55. lambda r: r.set(_refresh_key(refresh_jti), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS)
  56. )
  57. def is_refresh_token_valid(refresh_jti: str) -> bool:
  58. if get_sync_redis_client() is None:
  59. return True
  60. if not refresh_jti:
  61. return False
  62. exists = _redis_call_safe(lambda r: r.exists(_refresh_key(refresh_jti)))
  63. if exists is None:
  64. return True
  65. return exists > 0
  66. def revoke_refresh_token(refresh_jti: str) -> None:
  67. _redis_call_safe(lambda r: r.delete(_refresh_key(refresh_jti)))
  68. def create_access_token(
  69. data: dict,
  70. session_id: str,
  71. token_jti: str,
  72. expires_delta: Optional[timedelta] = None,
  73. ) -> str:
  74. """创建JWT访问令牌"""
  75. to_encode = data.copy()
  76. if expires_delta:
  77. expire = datetime.utcnow() + expires_delta
  78. else:
  79. expire = datetime.utcnow() + timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
  80. to_encode.update(
  81. {
  82. "exp": expire,
  83. "type": "access",
  84. "sid": session_id,
  85. "jti": token_jti,
  86. }
  87. )
  88. encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
  89. return encoded_jwt
  90. def create_refresh_token(data: dict, session_id: str, refresh_jti: str) -> str:
  91. """创建刷新令牌"""
  92. to_encode = data.copy()
  93. expire = datetime.utcnow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
  94. to_encode.update(
  95. {
  96. "exp": expire,
  97. "type": "refresh",
  98. "sid": session_id,
  99. "jti": refresh_jti,
  100. }
  101. )
  102. encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
  103. return encoded_jwt
  104. def decode_token(token: str) -> dict:
  105. """解码JWT令牌"""
  106. try:
  107. payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
  108. return payload
  109. except JWTError:
  110. raise HTTPException(
  111. status_code=status.HTTP_401_UNAUTHORIZED,
  112. detail="无效的认证凭据",
  113. headers={"WWW-Authenticate": "Bearer"},
  114. )
  115. def get_current_user_by_token(token: str, db: Session) -> dict:
  116. """根据 token 获取当前用户(包含 Redis 会话校验;会话滑动续期)"""
  117. payload = decode_token(token)
  118. user_id: str = payload.get("sub")
  119. session_id: str = payload.get("sid")
  120. if not is_session_valid(session_id):
  121. raise HTTPException(
  122. status_code=status.HTTP_401_UNAUTHORIZED,
  123. detail="Token无效或已过期",
  124. )
  125. extend_session(session_id)
  126. if user_id is None:
  127. raise HTTPException(
  128. status_code=status.HTTP_401_UNAUTHORIZED,
  129. detail="无效的认证凭据",
  130. )
  131. user = user_crud.get(db, user_id)
  132. if not user or getattr(user, "del_flag", 0) != 0:
  133. raise HTTPException(
  134. status_code=status.HTTP_401_UNAUTHORIZED,
  135. detail="用户不存在或已删除",
  136. )
  137. if getattr(user, "status", "1") != "1":
  138. raise HTTPException(
  139. status_code=status.HTTP_401_UNAUTHORIZED,
  140. detail="用户已被禁用",
  141. )
  142. web_role = resolve_web_role_for_user(db, user.id)
  143. phone = (getattr(user, "phone_no", None) or "").strip() or None
  144. email = (getattr(user, "email", None) or "").strip() or None
  145. avatar = getattr(user, "avatar", None)
  146. current_user_dept_id.set(user.dept_id or "")
  147. tenant_no = (getattr(user, "tenant_no", None) or "").strip() or settings.DEFAULT_TENANT_NO
  148. return {
  149. "user_id": user.id,
  150. "username": user.user_no,
  151. "real_name": user.nick_name,
  152. "role": web_role,
  153. "region_id": user.dept_id,
  154. "tenant_no": tenant_no,
  155. "phone": phone,
  156. "email": email,
  157. "avatar_url": avatar,
  158. "is_active": user.status == "1",
  159. }
  160. def get_current_user(
  161. token: str = Depends(oauth2_scheme),
  162. db: Session = Depends(get_db),
  163. ) -> dict:
  164. """获取当前用户(sync 依赖,配合 sync 路由在线程池执行,避免阻塞事件循环)。"""
  165. return get_current_user_by_token(token, db)
  166. def get_current_user_simple(token: str = Depends(oauth2_scheme)) -> dict:
  167. """简化版获取当前用户(不查数据库)"""
  168. payload = decode_token(token)
  169. user_id: str = payload.get("sub")
  170. if user_id is None:
  171. raise HTTPException(
  172. status_code=status.HTTP_401_UNAUTHORIZED,
  173. detail="无效的认证凭据",
  174. )
  175. current_user_dept_id.set(payload.get("region_id") or "")
  176. return {
  177. "user_id": user_id,
  178. "username": payload.get("username"),
  179. "role": payload.get("role"),
  180. "region_id": payload.get("region_id"),
  181. "real_name": payload.get("real_name"),
  182. }
  183. def get_client_ip(request: Request) -> str:
  184. """获取客户端IP"""
  185. forwarded = request.headers.get("X-Forwarded-For")
  186. if forwarded:
  187. return forwarded.split(",")[0].strip()
  188. return request.client.host if request.client else "unknown"