#!/usr/bin/env python3
"""Grab one full framebuffer from a no-auth VNC server and save as PNG.

Why VNC and not QMP screendump: with virtio-vga-gl the guest renders
into a GL scanout and screendump fails with "no surface"; QEMU's VNC
server reads the GL framebuffer back, so this always works. Pair the VM
with `-display egl-headless -vnc 127.0.0.1:17`.

Stdlib only (manual RFB handshake + raw-encoding framebuffer + PNG
writer) so it runs anywhere python3 exists.

Usage: vncshot.py <out.png> [host] [port]   # default 127.0.0.1:5917
"""
import socket, struct, sys, zlib

def recv_exact(s, n):
    buf = b""
    while len(buf) < n:
        chunk = s.recv(n - len(buf))
        if not chunk:
            raise EOFError("server closed")
        buf += chunk
    return buf

def png_write(path, w, h, rgb):
    def chunk(tag, data):
        c = struct.pack(">I", len(data)) + tag + data
        return c + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)
    raw = b"".join(b"\x00" + rgb[y * w * 3:(y + 1) * w * 3] for y in range(h))
    out = (b"\x89PNG\r\n\x1a\n"
           + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
           + chunk(b"IDAT", zlib.compress(raw, 6))
           + chunk(b"IEND", b""))
    with open(path, "wb") as f:
        f.write(out)

def main():
    out = sys.argv[1]
    host = sys.argv[2] if len(sys.argv) > 2 else "127.0.0.1"
    port = int(sys.argv[3]) if len(sys.argv) > 3 else 5917
    s = socket.create_connection((host, port), timeout=30)
    s.settimeout(30)

    ver = recv_exact(s, 12)              # "RFB 003.008\n"
    s.sendall(b"RFB 003.008\n")
    ntypes = recv_exact(s, 1)[0]
    types = recv_exact(s, ntypes)
    if 1 not in types:
        sys.exit(f"no 'None' security type offered: {list(types)}")
    s.sendall(b"\x01")
    res = struct.unpack(">I", recv_exact(s, 4))[0]
    if res != 0:
        sys.exit("security handshake failed")
    s.sendall(b"\x01")                   # ClientInit: shared
    w, h = struct.unpack(">HH", recv_exact(s, 4))
    recv_exact(s, 16)                    # server pixel format (ignored)
    namelen = struct.unpack(">I", recv_exact(s, 4))[0]
    recv_exact(s, namelen)

    # SetPixelFormat: 32bpp truecolor, shifts r=16 g=8 b=0 (little-endian BGRX)
    pf = struct.pack(">BBBBHHHBBBxxx", 32, 24, 0, 1, 255, 255, 255, 16, 8, 0)
    s.sendall(b"\x00\x00\x00\x00" + pf)
    s.sendall(struct.pack(">BxH i", 2, 1, 0))      # SetEncodings: Raw only
    s.sendall(struct.pack(">BBHHHH", 3, 0, 0, 0, w, h))  # full update request

    fb = bytearray(w * h * 4)
    got = False
    while not got:
        mtype = recv_exact(s, 1)[0]
        if mtype == 0:                    # FramebufferUpdate
            recv_exact(s, 1)
            nrects = struct.unpack(">H", recv_exact(s, 2))[0]
            for _ in range(nrects):
                x, y, rw, rh = struct.unpack(">HHHH", recv_exact(s, 8))
                enc = struct.unpack(">i", recv_exact(s, 4))[0]
                if enc == 0:
                    data = recv_exact(s, rw * rh * 4)
                    for row in range(rh):
                        off = ((y + row) * w + x) * 4
                        fb[off:off + rw * 4] = data[row * rw * 4:(row + 1) * rw * 4]
                elif enc == -224:        # LastRect
                    break
                else:
                    sys.exit(f"unsupported encoding {enc}")
            got = True
        elif mtype == 2:                  # Bell
            pass
        elif mtype == 3:                  # ServerCutText
            recv_exact(s, 3)
            ln = struct.unpack(">I", recv_exact(s, 4))[0]
            recv_exact(s, ln)
        else:
            sys.exit(f"unexpected server message {mtype}")

    rgb = bytearray(w * h * 3)
    for i in range(w * h):
        rgb[i * 3 + 0] = fb[i * 4 + 2]
        rgb[i * 3 + 1] = fb[i * 4 + 1]
        rgb[i * 3 + 2] = fb[i * 4 + 0]
    png_write(out, w, h, bytes(rgb))
    print(f"saved {out} ({w}x{h})")

if __name__ == "__main__":
    main()