#!/usr/bin/env python3
"""Decode QConnect WebSocket messages from Burp capture."""
import json, sys

def decode_varint(data, pos):
    val = 0; shift = 0
    while pos < len(data):
        b = data[pos]; pos += 1
        val |= (b & 0x7F) << shift
        if not (b & 0x80): return val, pos
        shift += 7
    return val, pos

def parse_fields(data):
    fields = []; pos = 0
    while pos < len(data):
        if pos >= len(data): break
        tag, pos = decode_varint(data, pos)
        fnum = tag >> 3; wt = tag & 7
        if wt == 0:
            val, pos = decode_varint(data, pos)
            fields.append((fnum, wt, val.to_bytes(8, 'little')))
        elif wt == 1:
            fields.append((fnum, wt, data[pos:pos+8])); pos += 8
        elif wt == 2:
            ln, pos = decode_varint(data, pos)
            fields.append((fnum, wt, data[pos:pos+ln])); pos += ln
        elif wt == 5:
            fields.append((fnum, wt, data[pos:pos+4])); pos += 4
        else: break
    return fields

def get_varint(fields, num):
    for f, w, d in fields:
        if f == num and w == 0:
            return int.from_bytes(d[:8], 'little')
    return None

def get_bytes(fields, num):
    for f, w, d in fields:
        if f == num and w == 2: return d
    return None

def get_fixed32(fields, num):
    for f, w, d in fields:
        if f == num and w == 5: return int.from_bytes(d[:4], 'little')
    return None

def get_fixed64(fields, num):
    for f, w, d in fields:
        if f == num and w == 1: return int.from_bytes(d[:8], 'little')
    return None

MSG_NAMES = {
    1: "ERROR", 21: "RNDR_JOIN", 23: "STATE_UPDATED", 25: "VOLUME_CHANGED",
    28: "QUALITY_CHANGED", 29: "VOLUME_MUTED", 41: "SET_STATE", 42: "SET_VOLUME",
    43: "SET_ACTIVE", 44: "SET_QUALITY", 45: "SET_LOOP", 46: "SET_SHUFFLE",
    47: "MUTE_VOLUME", 61: "CTRL_JOIN", 66: "QUEUE_REQ", 67: "QUEUE_REQ2",
    75: "QUEUE_UPDATE", 76: "SESSION_SETUP", 77: "ACK_JOIN", 79: "TRACK_LIST",
    82: "SRV_STATE_ECHO", 83: "DEVICE_LIST", 86: "SRV_VOL_ECHO",
    87: "SRV_VOL_ECHO2", 90: "SRV_QUEUE_INFO", 91: "SRV_QUEUE_RESP",
    97: "SRV_ACK", 98: "SRV_ACK2", 99: "SRV_QUALITY_ECHO",
    100: "SRV_QUALITY_INFO", 103: "SRV_TRACKLIST_RESP"
}

def decode_renderer_state(data):
    """Decode RendererState proto"""
    f = parse_fields(data)
    ps = get_varint(f, 1)
    bs = get_varint(f, 2)
    pos_data = get_bytes(f, 3)
    dur = get_varint(f, 4)
    qi = get_varint(f, 6)
    nqi = get_varint(f, 7)

    pos_ms = None; ts = None
    if pos_data:
        pf = parse_fields(pos_data)
        ts = get_fixed64(pf, 1)
        pos_ms = get_varint(pf, 2)

    return f"playing={ps} buffer={bs} pos={pos_ms}ms dur={dur}ms qi={qi} nqi={nqi}"

def decode_set_state(data):
    """Decode SET_STATE (type 41) payload"""
    f = parse_fields(data)
    ps = get_varint(f, 1)
    pos = get_varint(f, 2)
    qv = get_bytes(f, 3)
    ct = get_bytes(f, 4)
    nt = get_bytes(f, 5)

    qv_str = ""
    if qv:
        qvf = parse_fields(qv)
        qv_str = f" qver={get_varint(qvf, 1)}.{get_varint(qvf, 2)}"

    ct_str = ""
    if ct:
        ctf = parse_fields(ct)
        ct_str = f" cur_track(qi={get_varint(ctf, 1)} tid={get_fixed32(ctf, 2)})"

    nt_str = ""
    if nt:
        ntf = parse_fields(nt)
        nt_str = f" next_track(qi={get_varint(ntf, 1)} tid={get_fixed32(ntf, 2)})"

    return f"playing_state={'None' if ps is None else ps} pos={pos}{qv_str}{ct_str}{nt_str}"

def process_message(payload_str, direction):
    data = payload_str.encode('latin-1')
    if len(data) < 2: return

    # Decode frame layer
    pos = 0
    while pos < len(data):
        if pos >= len(data): break
        ft = data[pos]; pos += 1
        flen, pos = decode_varint(data, pos)
        if pos + flen > len(data): break
        fbody = data[pos:pos+flen]; pos += flen

        if ft != 6: continue  # only data frames

        # Parse frame body fields
        ff = parse_fields(fbody)
        f7 = get_bytes(ff, 7)
        if not f7: continue

        # Inside f7, field 3 = QConnect message
        cf = parse_fields(f7)
        for cfnum, cwt, cdata in cf:
            if cfnum != 3 or cwt != 2: continue
            mf = parse_fields(cdata)
            mt = get_varint(mf, 1)
            if mt is None: continue

            name = MSG_NAMES.get(mt, f"TYPE_{mt}")
            payload = get_bytes(mf, mt)

            extra = ""
            if mt == 23 and payload:  # STATE_UPDATED
                # Unwrap field 1 (RendererState wrapper)
                sf = parse_fields(payload)
                state_data = get_bytes(sf, 1)
                if state_data:
                    extra = " " + decode_renderer_state(state_data)
                else:
                    extra = " " + decode_renderer_state(payload)
            elif mt == 41 and payload:  # SET_STATE
                extra = " " + decode_set_state(payload)
            elif mt == 82 and payload:  # SRV_STATE_ECHO
                sf = parse_fields(payload)
                state_data = get_bytes(sf, 1)
                if state_data:
                    extra = " " + decode_renderer_state(state_data)

            arrow = ">>>" if "CLIENT" in direction else "<<<"
            print(f"  {arrow} {name}({mt}){extra}")

# Read from burp JSON
import subprocess
# Just process messages passed on stdin
for line in sys.stdin:
    line = line.strip()
    if not line: continue
    try:
        msg = json.loads(line)
        d = msg.get('direction', '')
        p = msg.get('payload', '')
        process_message(p, d)
    except: pass