MediaCrawler/mac-agent-os-main/05_tools/09_ave/scripts/lib/kling_webhook.py

"""
lib.kling_webhook — 可灵 Webhook 回调验签（按官方规范）
=========================================================

官方文档：https://klingai.com/document-api/api/get-started/callbacks

回调请求头（3 个）：
    webhook-id         任务/事件 ID（参与签名）
    webhook-timestamp  Unix 秒（参与签名，±5 分钟防重放）
    webhook-signature  形如 "v1,<base64>" 或 "v1,<a> v1,<b>"（可能多个）

验签规则：
    signedContent = f"{webhook-id}.{webhook-timestamp}.{rawBody}"
    key           = base64Decode(secret.removeprefix("whsec_"))
    expected      = base64(HMAC-SHA256(key, signedContent))
    与每个 "v1,<sig>" 做**常量时间比较**，任一匹配即通过

安全要点：
    · secret 只从配置读取（agent-local，600 权限，不进 git），不写日志
    · 常量时间比较（hmac.compare_digest）防时序攻击
    · 时间窗校验防重放
"""
from __future__ import annotations

import base64
import hashlib
import hmac
import json
import logging
import os
import time
from pathlib import Path
from typing import Any, Optional

log = logging.getLogger("kling.webhook")

TOLERANCE_SECONDS = 5 * 60          # 官方建议 ±5 分钟


def load_webhook_secret() -> str:
    """读取 webhook secret（配置来源：环境变量 > agent-local local.yaml → kling.webhook_secret）"""
    if os.environ.get("KLING_WEBHOOK_SECRET"):
        return os.environ["KLING_WEBHOOK_SECRET"].strip()
    try:
        import yaml
        base = Path(os.environ.get("AGENT_LOCAL")
                    or (Path.home() / "workbuddy-agent-os" / "agent-local"))
        p = base / "tools" / "ave" / "config" / "local.yaml"
        if p.exists():
            d = yaml.safe_load(p.read_text(encoding="utf-8")) or {}
            return ((d.get("kling") or {}).get("webhook_secret") or "").strip()
    except Exception as e:
        log.warning("读取 webhook secret 失败: %s", e)
    return ""


def _decode_key(secret: str) -> bytes:
    """去掉 whsec_ 前缀后 Base64 解码（官方要求）"""
    raw = secret.strip()
    if raw.startswith("whsec_"):
        raw = raw[len("whsec_"):]
    # 补齐 padding（部分 secret 可能被截断/省略 =）
    pad = (-len(raw)) % 4
    return base64.b64decode(raw + "=" * pad)


def sign(secret: str, webhook_id: str, timestamp: str, raw_body: bytes) -> str:
    """按官方规范生成签名（供自测/回放用）"""
    signed = f"{webhook_id}.{timestamp}.".encode("utf-8") + raw_body
    digest = hmac.new(_decode_key(secret), signed, hashlib.sha256).digest()
    return base64.b64encode(digest).decode("utf-8")


def verify(secret: str, raw_body: bytes, webhook_id: str,
           timestamp: str, signature_header: str,
           tolerance: int = TOLERANCE_SECONDS) -> tuple[bool, str]:
    """验签（返回 (通过与否, 原因)）

    失败原因便于排查（但**不泄露签名/密钥内容**）。
    """
    if not secret:
        return False, "未配置 webhook_secret"
    for name, v in (("webhook-id", webhook_id), ("webhook-timestamp", timestamp),
                    ("webhook-signature", signature_header)):
        if not v:
            return False, f"缺少请求头 {name}"

    # ① 时间窗（防重放）
    try:
        ts = int(float(timestamp))
    except Exception:
        return False, "webhook-timestamp 格式错误"
    drift = abs(int(time.time()) - ts)
    if tolerance and drift > tolerance:
        return False, f"时间戳偏差过大（{drift}s > {tolerance}s）"

    # ② 期望签名
    try:
        expected = sign(secret, webhook_id, timestamp, raw_body)
    except Exception as e:
        return False, f"签名计算失败: {e}"

    # ③ 逐个 "v1,<sig>" 常量时间比较
    parts = [p for p in (signature_header or "").split(" ") if p]
    for item in parts:
        ver, _, sig = item.partition(",")
        if ver.strip() != "v1" or not sig:
            continue
        if hmac.compare_digest(sig.strip(), expected):
            return True, "ok"
    return False, "签名不匹配"


def parse_callback(body: dict) -> dict:
    """解析回调 body → 统一结构（供业务处理）

    官方字段：
        id / status(submitted|processing|succeeded|failed) / message /
        create_time / update_time / external_id / outputs[]
    """
    if not isinstance(body, dict):
        return {"ok": False, "reason": "body 不是 JSON 对象"}
    outputs = body.get("outputs") or []
    files: list[dict] = []
    for o in outputs:
        if not isinstance(o, dict):
            continue
        t = o.get("type")
        url = o.get("url") or o.get("mp3_url") or o.get("wav_url") or ""
        if not url:
            continue
        files.append({
            "type": t, "url": url,
            "watermark_url": o.get("watermark_url") or "",
            "duration": o.get("duration") or o.get("mp3_duration") or "",
            "id": o.get("id") or "",
        })
    status = (body.get("status") or "").lower()
    return {
        "ok": True,
        "task_id": body.get("id") or "",
        "external_id": body.get("external_id") or "",
        "status": status,                                  # submitted|processing|succeeded|failed
        "succeeded": status == "succeeded",
        "failed": status == "failed",
        "message": body.get("message") or "",
        "create_time": body.get("create_time"),
        "update_time": body.get("update_time"),
        "files": files,
        "raw_keys": sorted(body.keys()),
    }