""" 认证模块:JWT Token + 密码哈希 部署到外网时的安全认证层 """ import hashlib import os import secrets from datetime import datetime, timedelta, timezone from typing import Optional from fastapi import Depends, HTTPException, status, Request, Response from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from jose import JWTError, jwt from pydantic import BaseModel # ==================== 配置 ==================== # JWT 密钥:优先从环境变量读取,否则生成随机密钥(每次重启失效) SECRET_KEY = os.environ.get("AUTH_SECRET_KEY", secrets.token_urlsafe(32)) ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_HOURS = int(os.environ.get("AUTH_TOKEN_EXPIRE_HOURS", "24")) # 默认用户(可通过环境变量覆盖) DEFAULT_USERNAME = os.environ.get("AUTH_USERNAME", "admin") DEFAULT_PASSWORD = os.environ.get("AUTH_PASSWORD", "pswd4admin") # Bearer token 提取器 security = HTTPBearer(auto_error=False) # ==================== 密码哈希 (SHA-256 + salt) ==================== def _hash_password(password: str, salt: str | None = None) -> str: """SHA-256 加盐哈希""" if salt is None: salt = secrets.token_hex(16) hashed = hashlib.sha256(f"{salt}:{password}".encode()).hexdigest() return f"{salt}${hashed}" def _verify_password(plain_password: str, stored_hash: str) -> bool: """验证密码""" if "$" not in stored_hash: return False salt, _ = stored_hash.split("$", 1) return _hash_password(plain_password, salt) == stored_hash # ==================== 模型 ==================== class LoginRequest(BaseModel): username: str password: str class TokenResponse(BaseModel): access_token: str token_type: str = "bearer" expires_in: int # 秒 class UserInfo(BaseModel): username: str # ==================== 用户存储(简单内存版) ==================== # 启动时对默认密码做哈希 _users_db: dict[str, str] = { DEFAULT_USERNAME: _hash_password(DEFAULT_PASSWORD), } def authenticate_user(username: str, password: str) -> Optional[str]: """验证用户,成功返回用户名,失败返回 None""" hashed = _users_db.get(username) if not hashed: return None if not _verify_password(password, hashed): return None return username # ==================== JWT 签发/验证 ==================== def create_access_token(username: str) -> tuple[str, int]: """创建 JWT token,返回 (token, expires_in_seconds)""" expires_delta = timedelta(hours=ACCESS_TOKEN_EXPIRE_HOURS) expire = datetime.now(timezone.utc) + expires_delta payload = { "sub": username, "exp": expire, "iat": datetime.now(timezone.utc), } token = jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM) return token, int(expires_delta.total_seconds()) def verify_token(token: str) -> Optional[str]: """验证 JWT token,成功返回用户名,失败返回 None""" try: payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM]) username: str = payload.get("sub") if username is None: return None # 检查用户是否仍然存在 if username not in _users_db: return None return username except JWTError: return None # ==================== FastAPI 依赖 ==================== async def get_current_user( request: Request, credentials: Optional[HTTPAuthorizationCredentials] = Depends(security), ) -> str: """ 从请求中提取并验证 JWT token 支持两种方式: 1. Authorization: Bearer (标准方式) 2. Cookie: access_token= (浏览器便捷方式) """ token = None # 方式1: Authorization header if credentials and credentials.credentials: token = credentials.credentials # 方式2: Cookie fallback if not token: token = request.cookies.get("access_token") if not token: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="未登录,请先登录", headers={"WWW-Authenticate": "Bearer"}, ) username = verify_token(token) if not username: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="登录已过期,请重新登录", headers={"WWW-Authenticate": "Bearer"}, ) return username