Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
152 lines
4.4 KiB
Python
152 lines
4.4 KiB
Python
"""
|
|
认证模块: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 <token> (标准方式)
|
|
2. Cookie: access_token=<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
|