#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""BotMSG 受众端（Server 角色）—— **本地**验证「登录令牌」，不读 bmsg 库、不回调 bmsg。

Server 是任意系统：它**不必是 bmsg 成员**。三步自助：
  ① 生成密钥对：        python3 bmsg_audience.py keygen --dir ./bmsg-keys
  ② 提交公钥(零审批)：  POST {bmsg}/audience  {"pubkey": "<base64(32B)>", "name": "<可选别名>",
                                              "owner_agent_id": "<可选: 运维者 BMSG>"}
  ③ 入站验票(~10 行)：
        from bmsg_audience import Audience
        aud = Audience(pubkey_file="./bmsg-keys/pub.key")     # 只在启动时读一次
        ok, claims = aud.verify(request.headers.get("Authorization", "").removeprefix("BotMSG "))
        if ok and claims["scope"] in (None, "read"):  # 票已证明"这是谁"(sub=成对匿名)
            ...  # 「能取什么」由 S 自己的授权表判定（票只答身份，不答权限）

票里有什么：{aud(受众指纹), sub(成对匿名), iat, exp, jti, scope?} —— **不含 B 的长期凭据**，
也不含 B 的真实身份（要真名需向 bmsg 换，留审计）。
"""
import argparse
import base64
import hashlib
import json
import os
import sys
import time
import urllib.request

DEFAULT_JWKS = os.environ.get("BMSG_JWKS_URL", "https://bmsg.zaixiangdu.com/jwks")
CACHE_PATH = os.environ.get("BMSG_JWKS_CACHE", os.path.expanduser("~/.cache/bmsg/jwks.json"))
SKEW = int(os.environ.get("BMSG_TICKET_SKEW", "60"))


def b64u(b):
    return base64.urlsafe_b64encode(b).decode().rstrip("=")


def ub64u(s):
    return base64.urlsafe_b64decode(str(s) + "=" * (-len(str(s)) % 4))


def keygen(dirpath):
    """生成 S 自己的 Ed25519 密钥对（私钥 0600，公钥可公开）。"""
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
    os.makedirs(dirpath, mode=0o700, exist_ok=True)
    priv_path = os.path.join(dirpath, "priv.key")
    pub_path = os.path.join(dirpath, "pub.key")
    if os.path.exists(priv_path):
        raise SystemExit("已存在 %s（不覆盖；要轮换请先备份再删）" % priv_path)
    pk = Ed25519PrivateKey.generate()
    raw_priv = pk.private_bytes(serialization.Encoding.Raw, serialization.PrivateFormat.Raw,
                                serialization.NoEncryption())
    raw_pub = pk.public_key().public_bytes(serialization.Encoding.Raw,
                                           serialization.PublicFormat.Raw)
    fd = os.open(priv_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, "wb") as f:
        f.write(b64u(raw_priv).encode())
    with open(pub_path, "w") as f:
        f.write(b64u(raw_pub))
    return {"priv": priv_path, "pub": pub_path, "pubkey": b64u(raw_pub),
            "aud": "k1:" + b64u(hashlib.sha256(raw_pub).digest())}


class Audience:
    """Server 侧验票器：只用 bmsg 公钥（JWKS）+ 自己的公钥指纹；**不回调 bmsg**（JWKS 有磁盘缓存）。"""

    def __init__(self, pubkey_file=None, pubkey_b64=None, jwks_url=DEFAULT_JWKS,
                 cache_path=CACHE_PATH, jwks_ttl=3600, skew=SKEW, accept_names=()):
        if pubkey_file:
            with open(pubkey_file) as f:
                pubkey_b64 = f.read().strip()
        raw = ub64u(pubkey_b64)
        if len(raw) != 32:
            raise SystemExit("pubkey 必须是 Ed25519 原始公钥的 base64(32 字节)")
        self.my_aud = "k1:" + b64u(hashlib.sha256(raw).digest())
        self.jwks_url = jwks_url
        self.cache_path = cache_path
        self.jwks_ttl = jwks_ttl
        self.skew = skew
        self.accept_names = set(accept_names)   # 允许的 aud_name 别名（可选）
        self._jwks = None
        self._jwks_at = 0
        self._seen = {}                          # jti -> exp（防重放；进程内窗口）

    # -- JWKS --
    def jwks(self, force=False):
        if self._jwks and not force and time.time() - self._jwks_at < self.jwks_ttl:
            return self._jwks
        try:
            with urllib.request.urlopen(self.jwks_url, timeout=8) as r:
                doc = json.load(r)
            self._jwks = doc
            self._jwks_at = time.time()
            try:
                os.makedirs(os.path.dirname(self.cache_path), exist_ok=True)
                with open(self.cache_path, "w") as f:
                    json.dump(doc, f)
            except Exception:
                pass
            return doc
        except Exception:
            # bmsg 不可达 ⇒ 退回磁盘缓存（**已有票不受影响**）
            try:
                with open(self.cache_path) as f:
                    self._jwks = json.load(f)
                    self._jwks_at = time.time() - self.jwks_ttl + 60
                    return self._jwks
            except Exception:
                return None

    def _key_for(self, kid):
        doc = self.jwks() or {}
        for k in doc.get("keys") or []:
            if not kid or k.get("kid") == kid:
                return k
        # kid 不匹配 ⇒ 强制刷新一次（轮换场景）
        doc = self.jwks(force=True) or {}
        for k in doc.get("keys") or []:
            if not kid or k.get("kid") == kid:
                return k
        return None

    # -- 验票 --
    def verify(self, ticket):
        """返回 (ok, claims|reason)。校验：签名 → aud 归属 → 时效 → 重放。"""
        from cryptography.exceptions import InvalidSignature
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
        t = (ticket or "").strip()
        if t.lower().startswith("botmsg "):
            t = t[7:].strip()
        parts = t.split(".")
        if len(parts) != 3:
            return False, "malformed_ticket"
        h, p, s = parts
        try:
            header = json.loads(ub64u(h))
            claims = json.loads(ub64u(p))
            sig = ub64u(s)
        except Exception:
            return False, "malformed_ticket"
        k = self._key_for(header.get("kid"))
        if not k or k.get("kty") != "OKP":
            return False, "unknown_signing_key"
        try:
            Ed25519PublicKey.from_public_bytes(ub64u(k["x"])).verify(
                sig, ("%s.%s" % (h, p)).encode("utf-8"))
        except InvalidSignature:
            return False, "bad_signature"
        except Exception:
            return False, "verify_error"
        # 受众归属：票必须就是发给**我**的
        if claims.get("aud") != self.my_aud and (claims.get("aud_name") or "") not in self.accept_names:
            return False, "audience_mismatch"
        now = int(time.time())
        if now > int(claims.get("exp", 0)) + self.skew:
            return False, "expired"
        if now + self.skew < int(claims.get("iat", 0)):
            return False, "not_yet_valid"
        jti = claims.get("jti") or ""
        for j, e in list(self._seen.items()):
            if e + self.skew < now:
                self._seen.pop(j, None)
        if jti in self._seen:
            return False, "replay"
        self._seen[jti] = int(claims.get("exp", now + 300))
        claims["_verified_at"] = now
        return True, claims


def main():
    ap = argparse.ArgumentParser(description="BotMSG 受众端（Server 角色）自助工具")
    sub = ap.add_subparsers(dest="cmd", required=True)
    p1 = sub.add_parser("keygen", help="生成 S 的 Ed25519 密钥对")
    p1.add_argument("--dir", default="./bmsg-keys")
    p2 = sub.add_parser("submit", help="向 bmsg 提交公钥（自证登记，零审批）")
    p2.add_argument("--dir", default="./bmsg-keys")
    p2.add_argument("--name", default=None, help="可选可读别名（如你的域名/服务名）")
    p2.add_argument("--owner", default=None, help="可选：运维者的 BMSG（便于通知/审计）")
    p2.add_argument("--bmsg", default=os.environ.get("BMSG_BASE", "https://bmsg.zaixiangdu.com"))
    p3 = sub.add_parser("verify", help="本地验一张票（不回调 bmsg）")
    p3.add_argument("--dir", default="./bmsg-keys")
    p3.add_argument("--ticket", required=True)
    p3.add_argument("--jwks", default=DEFAULT_JWKS)
    a = ap.parse_args()

    if a.cmd == "keygen":
        r = keygen(a.dir)
        print(json.dumps(r, ensure_ascii=False, indent=1))
        print("\n下一步：python3 bmsg_audience.py submit --dir %s --name <你的域名>" % a.dir, file=sys.stderr)
        return 0

    if a.cmd == "submit":
        with open(os.path.join(a.dir, "pub.key")) as f:
            pub = f.read().strip()
        body = {"pubkey": pub}
        if a.name:
            body["name"] = a.name
        if a.owner:
            body["owner_agent_id"] = a.owner
        req = urllib.request.Request(a.bmsg.rstrip("/") + "/audience",
                                     data=json.dumps(body).encode(), method="POST",
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=10) as r:
            print(json.dumps(json.load(r), ensure_ascii=False, indent=1))
        return 0

    if a.cmd == "verify":
        aud = Audience(pubkey_file=os.path.join(a.dir, "pub.key"), jwks_url=a.jwks)
        ok, info = aud.verify(a.ticket)
        print(json.dumps({"ok": ok, "info": info, "my_aud": aud.my_aud},
                         ensure_ascii=False, indent=1))
        return 0 if ok else 2
    return 1


if __name__ == "__main__":
    sys.exit(main())