#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""tsa_anchor.py — RFC 3161 可信时间戳（TSA）for WDY records — 纯标准库。

把 WDY 记录的电子证据哈希（文件 SHA-256）提交至独立时间戳机构（TSA），
取得 RFC 3161 时间戳令牌（.tsr）；令牌带机构签名，可离线独立核验
（openssl ts -verify），不依赖 WDY 与链方系统存活。

内置提供方（公共端点，2026-09-12 实测签发成功）：
  digicert  http://timestamp.digicert.com   （默认）
  sectigo   http://timestamp.sectigo.com
  freetsa   https://freetsa.org/tsr

用法
----
  # 盖章（生成 .tsr + 回写 manifest.anchors.tsa + 发布到验证页目录）
  python tsa_anchor.py --stamp 原件.pdf --record WDY-2026-0012-P2XK4P
  # 解析打印
  python tsa_anchor.py --info tsa.tsr
  # 核验（摘要比对 + openssl 验签；--ca 指机构根证书）
  python tsa_anchor.py --verify tsa.tsr -f 原件.pdf --ca DigiCertTrustedRootG4.crt.pem
  # 自检
  python tsa_anchor.py --self-test

退出码：0 成功；1 失败；2 参数错误；3 机构拒签。
"""

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

import wdy_anchor_common as C

PROVIDERS = {
    "digicert": {
        "name": "DigiCert RFC 3161 TSA",
        "endpoint": "http://timestamp.digicert.com",
        "ca_hint": "https://cacerts.digicert.com/DigiCertTrustedRootG4.crt.pem",
    },
    "sectigo": {
        "name": "Sectigo RFC 3161 TSA",
        "endpoint": "http://timestamp.sectigo.com",
        "ca_hint": "https://crt.sectigo.com/SectigoPublicTimeStampingRootR46.pem",
    },
    "freetsa": {
        "name": "FreeTSA",
        "endpoint": "https://freetsa.org/tsr",
        "ca_hint": "https://freetsa.org/files/cacert.pem",
    },
}

OID_SHA256 = "2.16.840.1.101.3.4.2.1"
OID_SIGNED_DATA = "1.2.840.113549.1.7.2"
OID_TST_INFO = "1.2.840.113549.1.9.16.1.4"

STATUS_TEXT = {0: "granted", 1: "grantedWithMods", 2: "rejection", 3: "waiting",
               4: "revocationWarning", 5: "revocationNotification"}


# ─────────────────────────── 最小 DER 编解码 ───────────────────────────
def _der_len(n: int) -> bytes:
    if n < 0x80:
        return bytes([n])
    b = n.to_bytes((n.bit_length() + 7) // 8, "big")
    return bytes([0x80 | len(b)]) + b


def _tlv(tag: int, content: bytes) -> bytes:
    return bytes([tag]) + _der_len(len(content)) + content


def _der_uint(n: int) -> bytes:
    if n == 0:
        return _tlv(0x02, b"\x00")
    b = n.to_bytes((n.bit_length() + 7) // 8, "big")
    if b[0] & 0x80:
        b = b"\x00" + b
    return _tlv(0x02, b)


def _der_oid(dotted: str) -> bytes:
    parts = [int(x) for x in dotted.split(".")]
    body = bytes([40 * parts[0] + parts[1]])
    for p in parts[2:]:
        if p < 0x80:
            body += bytes([p])
        else:
            enc = []
            while p:
                enc.insert(0, p & 0x7F)
                p >>= 7
            for i in range(len(enc) - 1):
                enc[i] |= 0x80
            body += bytes(enc)
    return _tlv(0x06, body)


def _der_null() -> bytes:
    return b"\x05\x00"


def _der_bool(v: bool) -> bytes:
    return b"\x01\x01" + (b"\xff" if v else b"\x00")


def _der_octet(b: bytes) -> bytes:
    return _tlv(0x04, b)


def _der_seq(*items: bytes) -> bytes:
    return _tlv(0x30, b"".join(items))


def _read_tlv(buf, off):
    """读取一个 TLV，返回 (tag, value_bytes, next_off)。"""
    if off + 2 > len(buf):
        raise ValueError("DER 截断")
    tag = buf[off]
    off += 1
    if tag & 0x1F == 0x1F:
        raise ValueError(f"不支持多字节 tag 0x{tag:02x}")
    l = buf[off]
    off += 1
    if l & 0x80:
        n = l & 0x7F
        if n == 0 or off + n > len(buf):
            raise ValueError("DER 长度域非法")
        l = int.from_bytes(buf[off:off + n], "big")
        off += n
    if off + l > len(buf):
        raise ValueError("DER 值截断")
    return tag, buf[off:off + l], off + l


def _oid_dotted(body: bytes) -> str:
    if not body:
        return ""
    first = body[0]
    parts = [str(first // 40), str(first % 40)]
    val = 0
    for b in body[1:]:
        val = (val << 7) | (b & 0x7F)
        if not (b & 0x80):
            parts.append(str(val))
            val = 0
    return ".".join(parts)


# ─────────────────────────── 请求构造 / 响应解析 ───────────────────────────
def build_tsq(digest: bytes, cert_req=True, nonce=None):
    """构造 RFC 3161 TimeStampReq（DER）。返回 (tsq_bytes, nonce)。"""
    if nonce is None:
        nonce = int.from_bytes(os.urandom(8), "big") & 0x7FFFFFFFFFFFFFFF
    alg = _der_seq(_der_oid(OID_SHA256), _der_null())
    imprint = _der_seq(alg, _der_octet(digest))
    items = [_der_uint(1), imprint]
    if cert_req:
        items.append(_der_bool(True))
    items.append(_der_uint(nonce))
    return _der_seq(*items), nonce


def _parse_gentime(raw: bytes):
    s = raw.decode("ascii", "replace").strip()
    s2 = s
    if s2.endswith("Z"):
        s2 = s2[:-1]
    frac = ""
    if "." in s2 or "," in s2:
        s2, frac = re.split(r"[.,]", s2, maxsplit=1)
    dt = datetime.strptime(s2, "%Y%m%d%H%M%S").replace(tzinfo=timezone.utc)
    return (dt.strftime("%Y-%m-%dT%H:%M:%SZ"),
            dt.strftime("%Y-%m-%d %H:%M:%S") + " UTC")


def parse_tsr(data: bytes):
    """解析 TimeStampResp：返回 dict（status/serial/gen_time/imprint/policy…）。"""
    tag, body, end = _read_tlv(data, 0)
    if tag != 0x30:
        raise ValueError("TimeStampResp 外框不是 SEQUENCE")
    if end != len(data):
        raise ValueError("TimeStampResp 尾部多余数据")
    t1, status_info, o1 = _read_tlv(body, 0)
    if t1 != 0x30:
        raise ValueError("缺少 PKIStatusInfo")
    t2, status_val, _ = _read_tlv(status_info, 0)
    if t2 != 0x02:
        raise ValueError("缺少 PKIStatus 整数")
    status = int.from_bytes(status_val, "big")
    info = {"status": status, "status_text": STATUS_TEXT.get(status, f"unknown({status})"),
            "has_token": False}
    if status not in (0, 1):
        return info
    if o1 >= len(body):
        return info
    t3, _cinfo, _ = _read_tlv(body, o1)
    if t3 != 0x30:
        raise ValueError("缺少 timeStampToken（ContentInfo）")
    cinfo = _cinfo
    _t, oid_body, off = _read_tlv(cinfo, 0)
    if _oid_dotted(oid_body) != OID_SIGNED_DATA:
        raise ValueError("ContentInfo 不是 signedData")
    _t, explicit, off = _read_tlv(cinfo, off)   # [0] EXPLICIT
    _t, sd, _ = _read_tlv(explicit, 0)          # SignedData
    _, _, s1 = _read_tlv(sd, 0)                 # version
    _, _, s2 = _read_tlv(sd, s1)                # digestAlgorithms
    t_ec, ec, _ = _read_tlv(sd, s2)             # encapContentInfo
    if t_ec != 0x30:
        raise ValueError("缺少 encapContentInfo")
    _, _, e1 = _read_tlv(ec, 0)                 # eContentType
    t_e0, e0, _ = _read_tlv(ec, e1)             # [0] EXPLICIT
    t_oct, tst, _ = _read_tlv(e0, 0)            # OCTET STRING (TSTInfo)
    if t_oct != 0x04:
        raise ValueError("eContent 不是 OCTET STRING")
    t_ti, ti, _ = _read_tlv(tst, 0)
    if t_ti != 0x30:
        raise ValueError("TSTInfo 不是 SEQUENCE")
    _, _, off = _read_tlv(ti, 0)                            # version
    _, pol_body, off = _read_tlv(ti, off)                   # policy
    _, mi, off = _read_tlv(ti, off)                         # messageImprint
    _, _, m1 = _read_tlv(mi, 0)                             # alg
    t_mh, mh, _ = _read_tlv(mi, m1)                         # digest
    _, serial_body, off = _read_tlv(ti, off)                # serialNumber
    t_gt, gt, off = _read_tlv(ti, off)                      # genTime
    if t_gt not in (0x17, 0x18):
        raise ValueError("缺少 GeneralizedTime")
    iso, disp = _parse_gentime(gt)
    info.update({
        "has_token": True,
        "serial_hex": serial_body.hex().upper(),
        "serial_int": int.from_bytes(serial_body, "big"),
        "gen_time_iso": iso, "gen_time_disp": disp, "gen_time_raw": gt.decode("ascii", "replace"),
        "policy": _oid_dotted(pol_body),
        "imprint_hex": mh.hex(),
    })
    return info


# ─────────────────────────── 命令实现 ───────────────────────────
def tsa_stamp(file_path, provider_key="digicert", timeout=60):
    prov = PROVIDERS[provider_key]
    src = Path(file_path)
    digest = hashlib.sha256(src.read_bytes()).digest()
    tsq, nonce = build_tsq(digest)
    req = urllib.request.Request(prov["endpoint"], data=tsq, headers={
        "Content-Type": "application/timestamp-query",
        "Accept": "application/timestamp-reply",
        "User-Agent": "wdy-tsa-anchor/1.0"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        tsr = r.read(1_000_000)
    info = parse_tsr(tsr)
    return digest, tsr, info, prov


def cmd_stamp(a):
    manifest, key = a.manifest, a.key
    if a.record and not manifest:
        p = C.WORK / f"{a.record}.manifest.json"
        manifest = str(p) if p.exists() else None
    if not key:
        p = C.WORK / "wdy-2026-rot1.key"
        key = str(p) if p.exists() else None
    digest, tsr, info, prov = tsa_stamp(a.stamp, a.provider, timeout=a.timeout)
    print(f"机构：{prov['name']}（{prov['endpoint']}）")
    print(f"状态：{info.get('status_text')}")
    if not info.get("has_token"):
        print("机构未签发令牌。")
        return 3
    print(f"序列号：{info['serial_hex']}")
    print(f"签发时间：{info['gen_time_disp']}")
    print(f"消息摘要：{info['imprint_hex']}")
    if info["imprint_hex"] != digest.hex():
        print("⚠ 返回摘要与提交摘要不一致！")
        return 1
    out = Path(a.out) if a.out else (C.INTL / a.record / "tsa.tsr" if a.record
                                     else Path(a.stamp).with_suffix(".tsr"))
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_bytes(tsr)
    sha = hashlib.sha256(tsr).hexdigest()
    print(f"令牌已写入：{out}（{len(tsr)} 字节）")
    if a.record and manifest and key:
        anchor = {
            "type": "rfc3161", "provider": prov["name"], "endpoint": prov["endpoint"],
            "ca_hint": prov["ca_hint"],
            "serial": info["serial_hex"], "timestamp": info["gen_time_iso"],
            "gen_time": info["gen_time_disp"], "policy": info["policy"],
            "digest_algorithm": "sha256", "message_imprint": info["imprint_hex"],
            "tsr_file": out.name, "tsr_sha256": sha, "stamped_at": C.now_iso(),
        }
        C.merge_manifest_anchor(manifest, key, "tsa", anchor)
        print(f"manifest 已回写 anchors.tsa 并重签（{Path(manifest).name}）")
        pub = C.publish_to_verify(a.record, {out: "tsa.tsr", Path(manifest): "manifest.json"})
        if pub:
            print("已发布到验证页目录：", ", ".join(pub))
    C.append_report({"tool": "tsa_anchor", "action": "stamp", "record": a.record,
                     "provider": a.provider, "serial": info["serial_hex"],
                     "gen_time": info["gen_time_iso"], "tsr": str(out), "tsr_sha256": sha})
    return 0


def cmd_info(path):
    info = parse_tsr(Path(path).read_bytes())
    for k in ("status_text", "serial_hex", "gen_time_disp", "policy",
              "imprint_hex", "gen_time_raw"):
        if k in info:
            print(f"{k:>14}: {info[k]}")
    if not info.get("has_token"):
        print("(无令牌：非 granted)")
    return 0


def _openssl_verify(tsr, data_file, ca=None, timeout=60):
    exe = shutil.which("openssl")
    if not exe:
        return None, "未找到 openssl（--openssl 指定路径，或在 Git Bash 环境运行）"
    cmd = [exe, "ts", "-verify", "-in", str(tsr), "-data", str(data_file)]
    if ca:
        cmd += ["-CAfile", str(ca)]
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8",
                           errors="replace", timeout=timeout)
    except Exception as e:
        return None, f"openssl 执行失败：{e}"
    return p.returncode, ((p.stdout or "") + (p.stderr or "")).strip()


def cmd_verify(a):
    tsr = Path(a.verify)
    info = parse_tsr(tsr.read_bytes())
    ok = True
    print(f"状态：{info.get('status_text')}｜序列号：{info.get('serial_hex', '-')}"
          f"｜签发时间：{info.get('gen_time_disp', '-')}")
    if a.file:
        h = hashlib.sha256(Path(a.file).read_bytes()).hexdigest()
        match = (h == info.get("imprint_hex"))
        ok &= match
        print(f"原件摘要比对：{'一致 ✓' if match else '不一致 ✗'}")
    rc, out = _openssl_verify(tsr, a.file, a.ca) if a.file else (None, "未提供原件（-f），跳过 openssl 验签")
    if rc is None:
        print(f"openssl：{out}")
    else:
        print(f"openssl ts -verify：{'PASS ✓' if rc == 0 else 'FAIL ✗'}")
        print(out)
        ok &= (rc == 0)
        if rc != 0 and not a.ca:
            prov = None
            for p in PROVIDERS.values():
                if p["name"] in (info.get("provider") or ""):
                    prov = p
            print("提示：验签需要机构证书链（-CAfile），如 "
                  + (prov["ca_hint"] if prov else "机构根证书"))
    C.append_report({"tool": "tsa_anchor", "action": "verify", "tsr": str(tsr),
                     "result": "ok" if ok else "check-failed", "openssl_rc": rc})
    return 0 if ok else 1


def cmd_self_test():
    ok = True
    vec = _der_oid("1.2.840.113549.1.1.1").hex()
    ok &= vec == "06092a864886f70d010101"
    print(f"  DER OID 向量          {'PASS' if vec == '06092a864886f70d010101' else 'FAIL'}")
    tsq, nonce = build_tsq(bytes(range(32)), nonce=12345)
    ok &= tsq[0] == 0x30
    contains = bytes.fromhex("608648016503040201") in tsq and bytes(range(32)) in tsq
    ok &= contains
    print(f"  TSQ 结构（OID+摘要）  {'PASS' if contains else 'FAIL'}")
    t, v, o = _read_tlv(tsq, 0)
    ok &= (t == 0x30 and o == len(tsq))
    print(f"  TSQ TLV 往返          {'PASS' if t == 0x30 and o == len(tsq) else 'FAIL'}")
    iso, disp = _parse_gentime(b"20260912000000Z")
    ok &= iso == "2026-09-12T00:00:00Z"
    print(f"  GeneralizedTime 解析  {'PASS' if iso == '2026-09-12T00:00:00Z' else 'FAIL'}")
    print("自检", "4/4 PASS" if ok else "存在 FAIL")
    return 0 if ok else 1


def main():
    ap = argparse.ArgumentParser(description="RFC 3161 时间戳锚定（WDY）")
    ap.add_argument("--stamp", metavar="FILE", default=None)
    ap.add_argument("--info", metavar="FILE.tsr", default=None)
    ap.add_argument("--verify", metavar="FILE.tsr", default=None)
    ap.add_argument("--self-test", action="store_true")
    ap.add_argument("--provider", choices=sorted(PROVIDERS), default="digicert")
    ap.add_argument("--record", default=None)
    ap.add_argument("--manifest", default=None)
    ap.add_argument("--key", default=None)
    ap.add_argument("--out", default=None)
    ap.add_argument("-f", "--file", default=None, help="--verify 用的原件")
    ap.add_argument("--ca", default=None, help="--verify 用 CAfile（机构证书链 PEM）")
    ap.add_argument("--timeout", type=int, default=60)
    a = ap.parse_args()
    if a.self_test:
        sys.exit(cmd_self_test())
    if a.stamp:
        sys.exit(cmd_stamp(a))
    if a.info:
        sys.exit(cmd_info(a.info))
    if a.verify:
        sys.exit(cmd_verify(a))
    ap.print_help()


if __name__ == "__main__":
    main()