live-hub-py/tools/taf_dump.py

"""通用 Taf/JCE 递归解析 dump 工具(调试用)。"""
from __future__ import annotations

import base64
import json
import sys

sys.path.insert(0, ".")

from core.huya.taf_protocol import TafInputStream, TafType  # noqa: E402

STR = (TafType.STRING1, TafType.STRING4)
INTS = (TafType.INT8, TafType.INT16, TafType.INT32, TafType.INT64)


def parse_struct(i: TafInputStream, depth: int = 0):
    """解析到 STRUCT_END 或数据结束, 返回 dict。"""
    out = {}
    pad = "  " * depth
    while True:
        try:
            tag, dtype = i.peek_head()
        except Exception:
            break
        if dtype == TafType.STRUCT_END:
            i.read_head()
            break
        key = f"tag{tag}"
        if dtype == TafType.STRUCT_BEGIN:
            i.read_head()
            print(f"{pad}{key}: {{")
            out[key] = parse_struct(i, depth + 1)
            print(f"{pad}}}")
        elif dtype in STR:
            val = i.read_string(tag)
            show = val if len(val) <= 300 else val[:300] + f"...({len(val)})"
            print(f"{pad}{key}: STR {show!r}")
            out[key] = val
        elif dtype in INTS:
            val = {TafType.INT8: i.read_int8, TafType.INT16: i.read_int16,
                   TafType.INT32: i.read_int32, TafType.INT64: i.read_int64}[dtype](tag)
            print(f"{pad}{key}: INT {val}")
            out[key] = val
        elif dtype == TafType.ZERO:
            i.read_head()
            print(f"{pad}{key}: 0")
            out[key] = 0
        elif dtype == TafType.MAP:
            n = i.read_int32(0)
            print(f"{pad}{key}: MAP[{n}]")
            m = {}
            for _ in range(n):
                kt, kdt = i.peek_head()
                if kdt in STR:
                    kv = i.read_string(kt)
                elif kdt in INTS:
                    kv = str({TafType.INT8: i.read_int8, TafType.INT16: i.read_int16,
                              TafType.INT32: i.read_int32, TafType.INT64: i.read_int64}[kdt](kt))
                else:
                    break
                vt, vdt = i.peek_head()
                if vdt == TafType.STRUCT_BEGIN:
                    i.read_head()
                    print(f"{pad}  [{kv}]: {{")
                    m[str(kv)] = parse_struct(i, depth + 2)
                    print(f"{pad}  }}")
                elif vdt == TafType.SIMPLE_LIST:
                    sv = i.read_bytes(vt)
                    m[str(kv)] = sv.hex()
                    print(f"{pad}  [{kv}] BYTES[{len(sv)}] {sv[:40].hex()}")
                elif vdt in STR:
                    sv = i.read_string(vt); m[str(kv)] = sv
                    print(f"{pad}  [{kv}] STR {sv[:120]!r}")
                else:
                    i.read_head(); i.skip_field(vdt)
            out[key] = m
        elif dtype == TafType.LIST:
            n = i.read_int32(0)
            print(f"{pad}{key}: LIST[{n}]")
            items = []
            for _ in range(n):
                it, idt = i.read_head()
                if idt in STR:
                    items.append(i.read_string(it))
                elif idt in INTS:
                    items.append({TafType.INT8: i.read_int8, TafType.INT16: i.read_int16,
                                  TafType.INT32: i.read_int32, TafType.INT64: i.read_int64}[idt](it))
                else:
                    i.skip_field(idt)
            out[key] = items
        elif dtype == TafType.SIMPLE_LIST:
            i.read_head()  # 元素类型 INT8
            n = i.read_int32(0)
            raw = (getattr(i, 'stream', None) or getattr(i, 'buf')).read(n)
            print(f"{pad}{key}: BYTES[{n}] {raw[:60].hex()}")
            out[key] = raw.hex()
        else:
            print(f"{pad}{key}: ?type={hex(dtype)}")
            try:
                i.skip_field(dtype)
            except Exception:
                break
    return out


def dump_packet(data: bytes, label: str):
    print("=" * 70)
    print(label, f"({len(data)} bytes)")
    print("=" * 70)
    from core.huya.wup_protocol import WupResponse
    r = WupResponse()
    try:
        r.decode(data)
        print(f"iVersion={r.iVersion} servant={r.sServantName!r} func={r.sFuncName!r}")
        for k, v in r.newdata.items():
            print(f"--- newdata[{k!r}] len={len(v)}")
            if not k or len(v) < 8:
                print("   ", v.hex())
                continue
            try:
                parse_struct(TafInputStream(v))
            except Exception as e:
                print("   parse err:", e, "| hex:", v.hex()[:200])
    except Exception as e:
        # 可能没有长度前缀, 直接当 body
        print("decode with prefix failed:", e, "-> try raw")
        parse_struct(TafInputStream(data))


def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("path")
    ap.add_argument("--b64", action="store_true", help="输入是 base64")
    args = ap.parse_args()
    if args.path == "-":
        raw = sys.stdin.buffer.read()
    else:
        raw = open(args.path, "rb").read()
    if args.b64:
        data = base64.b64decode(raw)
    else:
        data = raw
    dump_packet(data, args.path)


if __name__ == "__main__":
    main()