#!/usr/bin/env python3
"""ADVERSARIAL independent validator for the PreAuthResponse TDF payload.

Deliberately re-implemented from the documented wire rules rather than
importing heat2's decoder, so an encoder/decoder bug that cancels out in a
round-trip is still caught.  Checks:
  * every tag decodes to a legal 4-char label (chars 0x20..0x5F, no embedded
    space, trailing-space padding only)
  * fields at every nesting level are in STRICTLY ascending packed-tag order
  * varints are canonical (shortest form), no 0x40 sign bit set
  * string lengths include exactly one trailing NUL and no interior NUL
  * every struct/group is terminated by exactly one 0x00
  * the payload is consumed exactly (no trailing bytes, no overrun)
  * list/map headers use legal element type codes
"""
import struct
import sys

PROBLEMS = []
def bad(off, msg, ctx=b""):
    PROBLEMS.append((off, msg, ctx))

VALID_TYPES = {0x00: "int", 0x01: "string", 0x02: "blob", 0x03: "struct",
               0x04: "list", 0x05: "map", 0x06: "union", 0x07: "intlist",
               0x08: "objtype", 0x09: "objid", 0x0A: "float"}


def dec_tag(b, off):
    a, b1, c = b[0], b[1], b[2]
    v = [(a >> 2) & 0x3F, ((a & 3) << 4) | ((b1 >> 4) & 0xF),
         ((b1 & 0xF) << 2) | ((c >> 6) & 3), c & 0x3F]
    chars = []
    for x in v:
        chars.append(chr(x + 0x20) if x else " ")
    raw = "".join(chars)
    label = raw.rstrip()
    if not label:
        bad(off, "tag decodes to all-padding (empty label) raw=%s" % b.hex())
    if " " in label:
        bad(off, "tag %r has an interior space (padding is not trailing-only) raw=%s"
            % (raw, b.hex()))
    for ch in label:
        if not (0x20 <= ord(ch) <= 0x5F):
            bad(off, "tag %r contains non-Heat2 char %r raw=%s" % (raw, ch, b.hex()))
        if not (ch.isupper() or ch.isdigit()):
            bad(off, "tag %r char %r is not [A-Z0-9] (suspicious for a Blaze tag)"
                % (raw, ch))
    # re-encode check
    cc = [(ord(ch) - 0x20) & 0x3F for ch in raw]
    re_enc = bytes(((cc[0] << 2) | (cc[1] >> 4),
                    ((cc[1] & 0xF) << 4) | (cc[2] >> 2),
                    ((cc[2] & 3) << 6) | cc[3]))
    if re_enc != bytes(b):
        bad(off, "tag %r does not re-encode: %s != %s" % (raw, re_enc.hex(), bytes(b).hex()))
    return label, bytes(b)


def rd_varint(buf, i, what):
    start = i
    b = buf[i]; i += 1
    if b & 0x40:
        bad(start, "%s: first varint byte 0x%02x has sign bit 0x40 set" % (what, b))
    val = b & 0x3F
    nbytes = 1
    if b & 0x80:
        shift = 6
        while True:
            if i >= len(buf):
                bad(start, "%s: varint runs past end of buffer" % what)
                return val, i
            b = buf[i]; i += 1
            nbytes += 1
            val |= (b & 0x7F) << shift
            shift += 7
            if not (b & 0x80):
                if b == 0x00:
                    bad(start, "%s: non-canonical varint (trailing zero group) %s"
                        % (what, buf[start:i].hex()))
                break
    # canonical length check
    v = val
    exp = 1
    v >>= 6
    while v:
        exp += 1
        v >>= 7
    if exp != nbytes:
        bad(start, "%s: varint for %d used %d bytes, canonical is %d (%s)"
            % (what, val, nbytes, exp, buf[start:i].hex()))
    return val, i


def rd_value(buf, i, typ, path, depth):
    if typ == 0x00:
        v, i = rd_varint(buf, i, path)
        return v, i
    if typ == 0x01:
        start = i
        ln, i = rd_varint(buf, i, path + ".len")
        if ln == 0:
            bad(start, "%s: string length 0 (must be >=1 to hold the NUL)" % path)
            return "", i
        if i + ln > len(buf):
            bad(start, "%s: string length %d overruns buffer" % (path, ln))
            return "", len(buf)
        raw = buf[i:i + ln]; i += ln
        if raw[-1] != 0x00:
            bad(start, "%s: string not NUL-terminated, last byte 0x%02x (%s)"
                % (path, raw[-1], raw.hex()))
        if 0x00 in raw[:-1]:
            bad(start, "%s: string has interior NUL (%s)" % (path, raw.hex()))
        return raw[:-1].decode("utf-8", "replace"), i
    if typ == 0x02:
        ln, i = rd_varint(buf, i, path + ".len")
        return bytes(buf[i:i + ln]), i + ln
    if typ == 0x03:
        return rd_struct(buf, i, path, depth + 1, terminated=True)
    if typ == 0x04:
        et = buf[i]
        if et not in VALID_TYPES:
            bad(i, "%s: list element type 0x%02x is not a legal TDF type" % (path, et))
        i += 1
        n, i = rd_varint(buf, i, path + ".count")
        items = []
        for k in range(n):
            v, i = rd_value(buf, i, et, "%s[%d]" % (path, k), depth)
            items.append(v)
        return (VALID_TYPES.get(et), items), i
    if typ == 0x05:
        kt = buf[i]; vt = buf[i + 1]
        for nm, t in (("key", kt), ("value", vt)):
            if t not in VALID_TYPES:
                bad(i, "%s: map %s type 0x%02x is not a legal TDF type" % (path, nm, t))
        i += 2
        n, i = rd_varint(buf, i, path + ".count")
        items = []
        for k in range(n):
            kk, i = rd_value(buf, i, kt, "%s{%d}.k" % (path, k), depth)
            vv, i = rd_value(buf, i, vt, "%s{%d}.v" % (path, k), depth)
            items.append((kk, vv))
        return (VALID_TYPES.get(kt), VALID_TYPES.get(vt), items), i
    bad(i, "%s: type 0x%02x not handled by validator" % (path, typ))
    raise SystemExit("cannot continue")


def rd_struct(buf, i, path, depth, terminated):
    fields = []
    prev = None
    while True:
        if i >= len(buf):
            if terminated:
                bad(i, "%s: struct ran off the end without a 0x00 terminator" % path)
            break
        if terminated and buf[i] == 0x00:
            i += 1
            break
        if i + 4 > len(buf):
            bad(i, "%s: %d trailing bytes, too short for a tag+type header (%s)"
                % (path, len(buf) - i, buf[i:].hex()))
            break
        label, packed = dec_tag(buf[i:i + 3], i)
        typ = buf[i + 3]
        if typ not in VALID_TYPES:
            bad(i + 3, "%s.%s: type byte 0x%02x is not a legal TDF type" % (path, label, typ))
        if prev is not None and packed <= prev[1]:
            rel = "==" if packed == prev[1] else "<"
            bad(i, "%s: field %r (tag %s) is %s previous %r (tag %s) -- ORDER VIOLATION"
                % (path, label, packed.hex(), rel, prev[0], prev[1].hex()))
        prev = (label, packed)
        i += 4
        val, i = rd_value(buf, i, typ, "%s.%s" % (path, label), depth)
        fields.append((label, VALID_TYPES.get(typ), val))
    return fields, i


def show(fields, d=0):
    for lbl, tn, v in fields:
        if tn == "struct":
            print("  " * d + "%s (struct) {" % lbl)
            show(v, d + 1)
            print("  " * d + "}")
        else:
            print("  " * d + "%s (%s) = %r" % (lbl, tn, v))


def main(path):
    data = open(path, "rb").read()
    plen = struct.unpack_from(">I", data, 0)[0]
    mlen = struct.unpack_from(">H", data, 4)[0]
    comp = struct.unpack_from(">H", data, 6)[0]
    cmd = struct.unpack_from(">H", data, 8)[0]
    msgnum = (data[10] << 16) | (data[11] << 8) | data[12]
    mtype = (data[13] >> 5) & 7
    uidx = data[13] & 0x1F
    print("== %s (%d bytes) ==" % (path, len(data)))
    print("payload_len=%d meta_len=%d comp=0x%04x cmd=0x%04x msgNum=%d "
          "msgType=%d userIdx=%d opts=0x%02x rsv=0x%02x"
          % (plen, mlen, comp, cmd, msgnum, mtype, uidx, data[14], data[15]))
    if 16 + mlen + plen != len(data):
        bad(0, "frame size mismatch: 16+%d+%d=%d but file is %d"
            % (mlen, plen, 16 + mlen + plen, len(data)))
    payload = data[16 + mlen:16 + mlen + plen]
    fields, end = rd_struct(payload, 0, "", 0, terminated=False)
    if end != len(payload):
        bad(end, "payload not fully consumed: stopped at %d of %d (rest=%s)"
            % (end, len(payload), payload[end:].hex()))
    print("--- decoded ---")
    show(fields)
    print("--- result ---")
    if PROBLEMS:
        for off, msg, _ in PROBLEMS:
            lo = max(0, off - 8)
            print("BUG @0x%04x (payload): %s" % (off, msg))
            print("     bytes %s" % payload[lo:off + 16].hex(" "))
        return 1
    print("PASS: %d top-level fields, payload consumed exactly (%d bytes)"
          % (len(fields), len(payload)))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1]))
