File size: 2,550 Bytes
b461de0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import time
import json
import hmac
import hashlib
import base64
from fastapi import Request, HTTPException
from core.config import ConfigManager, hash_password
TOKEN_EXPIRY = 86400 # 24 小时
def _b64url_encode(data: bytes) -> str:
return base64.urlsafe_b64encode(data).rstrip(b"=").decode()
def _b64url_decode(s: str) -> bytes:
s += "=" * (4 - len(s) % 4)
return base64.urlsafe_b64decode(s)
def create_jwt(config: ConfigManager) -> str:
secret = config.get("admin", "jwt_secret")
header = _b64url_encode(json.dumps({"alg": "HS256", "typ": "JWT"}).encode())
payload = _b64url_encode(
json.dumps({"role": "admin", "exp": int(time.time()) + TOKEN_EXPIRY}).encode()
)
signature = _b64url_encode(
hmac.new(secret.encode(), f"{header}.{payload}".encode(), hashlib.sha256).digest()
)
return f"{header}.{payload}.{signature}"
def verify_jwt(token: str, config: ConfigManager) -> dict:
try:
parts = token.split(".")
if len(parts) != 3:
raise ValueError("invalid token format")
header_b64, payload_b64, sig_b64 = parts
secret = config.get("admin", "jwt_secret")
expected_sig = _b64url_encode(
hmac.new(
secret.encode(), f"{header_b64}.{payload_b64}".encode(), hashlib.sha256
).digest()
)
if not hmac.compare_digest(sig_b64, expected_sig):
raise ValueError("invalid signature")
payload = json.loads(_b64url_decode(payload_b64))
if payload.get("exp", 0) < time.time():
raise ValueError("token expired")
return payload
except Exception as e:
raise HTTPException(status_code=401, detail=str(e))
def verify_password(password: str, config: ConfigManager) -> bool:
stored_hash = config.get("admin", "password_hash")
salt = config.get("admin", "salt")
computed, _ = hash_password(password, salt)
return hmac.compare_digest(computed, stored_hash)
def require_admin(config: ConfigManager):
"""返回一个 FastAPI 依赖函数"""
async def dependency(request: Request):
# 先试 X-Admin-Token，再试 Authorization
token = request.headers.get("X-Admin-Token", "")
if not token:
auth = request.headers.get("Authorization", "")
if auth.startswith("Bearer "):
token = auth[7:]
if not token:
raise HTTPException(status_code=401, detail="missing token")
return verify_jwt(token, config)
return dependency
|