120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
"""
|
||
JWT 工具：解析 Java 端签发的 token、从请求中提取 token。
|
||
与 Java 端 JwtService 共用同一个 HS256 密钥（AIIMAGE_JWT_SECRET）。
|
||
"""
|
||
import os
|
||
import sys
|
||
import time
|
||
from typing import Optional, Tuple
|
||
|
||
from flask import request
|
||
|
||
try:
|
||
import jwt as pyjwt
|
||
_PYJWT_IMPORT_ERROR = None
|
||
except ImportError as _e:
|
||
pyjwt = None
|
||
_PYJWT_IMPORT_ERROR = _e
|
||
# 启动时立刻给出醒目提示，避免上线后才发现一直 401
|
||
print(
|
||
"[auth] WARNING: PyJWT 未安装，所有依赖 JWT 的接口都会返回 401。"
|
||
" 请执行 `pip install PyJWT==2.10.1` 或 `pip install -r requirements.txt`。",
|
||
file=sys.stderr,
|
||
)
|
||
|
||
|
||
_DEFAULT_SECRET = "please-change-this-secret-please-rotate-at-least-32-bytes"
|
||
COOKIE_NAME = os.getenv("AIIMAGE_AUTH_COOKIE_NAME", "aiimage_token")
|
||
|
||
# 桌面端与 Java 服务器时钟可能漂移（笔记本休眠、用户手动改时间、跨时区等），
|
||
# 给 JWT exp/nbf 校验留出容差，避免 /api/auth/sync 因为时间不同步而 401。
|
||
# 默认 5 分钟，可通过 AIIMAGE_JWT_LEEWAY_SECONDS 覆盖。
|
||
try:
|
||
_JWT_LEEWAY_SECONDS = int(os.getenv("AIIMAGE_JWT_LEEWAY_SECONDS", "300"))
|
||
except ValueError:
|
||
_JWT_LEEWAY_SECONDS = 300
|
||
|
||
|
||
def _signing_key() -> bytes:
|
||
"""与 Java JwtService.signingKey 保持一致：UTF-8 字节，不足 32 字节右侧补 0。"""
|
||
secret = os.getenv("AIIMAGE_JWT_SECRET", _DEFAULT_SECRET)
|
||
key_bytes = secret.encode("utf-8")
|
||
if len(key_bytes) < 32:
|
||
key_bytes = key_bytes + b"\x00" * (32 - len(key_bytes))
|
||
return key_bytes
|
||
|
||
|
||
def _describe_clock_skew(token: str) -> str:
|
||
"""解析 token 里的 iat/exp，与本地时钟比较，返回 ' iat=.. exp=.. now=.. skew=..s' 字符串。
|
||
仅供日志使用，不做安全决策；解析失败返回空串，不影响主流程。"""
|
||
if pyjwt is None:
|
||
return ""
|
||
try:
|
||
unverified = pyjwt.decode(token, options={"verify_signature": False, "verify_exp": False, "verify_iat": False, "verify_nbf": False})
|
||
except Exception:
|
||
return ""
|
||
iat = unverified.get("iat")
|
||
exp = unverified.get("exp")
|
||
now = int(time.time())
|
||
parts = [f"now={now}"]
|
||
if isinstance(iat, (int, float)):
|
||
parts.append(f"iat={int(iat)} skew_iat={now - int(iat)}s")
|
||
if isinstance(exp, (int, float)):
|
||
parts.append(f"exp={int(exp)} skew_exp={now - int(exp)}s")
|
||
return ", " + " ".join(parts)
|
||
|
||
|
||
def parse_token_with_reason(token: str) -> Tuple[Optional[dict], Optional[str]]:
|
||
"""解析 JWT，返回 (payload, error_reason)。
|
||
payload 命中时 error_reason 为 None；失败时 payload 为 None，error_reason 描述根因。"""
|
||
if not token:
|
||
return None, "token 为空"
|
||
if pyjwt is None:
|
||
return None, f"PyJWT 未安装（{_PYJWT_IMPORT_ERROR}）"
|
||
# 打印 token 头，便于发现 alg 不是 HS256 等情况（不验签，仅 base64 解码 header）
|
||
try:
|
||
header = pyjwt.get_unverified_header(token)
|
||
except Exception as e:
|
||
return None, f"无法解析 token header: {type(e).__name__}: {e}"
|
||
alg = header.get("alg") or "?"
|
||
try:
|
||
payload = pyjwt.decode(
|
||
token,
|
||
_signing_key(),
|
||
algorithms=["HS256", "HS384", "HS512"],
|
||
leeway=_JWT_LEEWAY_SECONDS,
|
||
)
|
||
except Exception as e:
|
||
# 常见：ExpiredSignatureError / InvalidSignatureError / DecodeError / InvalidAlgorithmError
|
||
# 把 token 内 exp/iat 与本地时钟一起打出来，定位"时间漂移"类失败更直接
|
||
skew_info = _describe_clock_skew(token)
|
||
return None, f"{type(e).__name__}: {e} (token alg={alg}{skew_info})"
|
||
sub = payload.get("sub")
|
||
try:
|
||
user_id = int(sub) if sub is not None else None
|
||
except (TypeError, ValueError):
|
||
return None, f"sub 字段非法: {sub!r}"
|
||
if user_id is None:
|
||
return None, "payload 缺少 sub 字段"
|
||
return {
|
||
"user_id": user_id,
|
||
"username": payload.get("username") or "",
|
||
"device_id": payload.get("deviceId") or "",
|
||
}, None
|
||
|
||
|
||
def parse_token(token: str) -> Optional[dict]:
|
||
"""解析 JWT，返回 {user_id, username, device_id}；失败返回 None。"""
|
||
payload, _ = parse_token_with_reason(token)
|
||
return payload
|
||
|
||
|
||
def get_token_from_request() -> str:
|
||
"""优先从 Authorization: Bearer 取，其次从 cookie 取。"""
|
||
auth = request.headers.get("Authorization", "")
|
||
if auth.startswith("Bearer "):
|
||
token = auth[len("Bearer "):].strip()
|
||
if token:
|
||
return token
|
||
return (request.cookies.get(COOKIE_NAME) or "").strip()
|