CCE 2018 Quals - VNC

from PIL import Image
import zlib
import struct

def db(d):
    return d[1:], struct.unpack('>B', d[:1])[0]
def dw(d):
    return d[2:], struct.unpack('>H', d[:2])[0]
def dd(d):
    return d[4:], struct.unpack('>I', d[:4])[0]
def color(f):
    f, c = db(f)
    return f, (c,)

# Run-length number
def number(f):
    n = 0
    while True:
        f, byte = db(f)
        if byte == 255:
            n += 255
        else:
            n += byte + 1
            break
    return f, n

# Continuously decompress the zlib stream.
z = zlib.decompressobj()

# Draw the accumulated "diffs" sent from the server.
I = Image.new('L', (1713,927))

counter = 0

# A rectangular part of the screen.
# Pixels are filled left-to-right, top-to-bottom.
class Tile(object):
    def __init__(self, I, tx, ty, tw, th):
        self.I = I
        self.tx = tx
        self.ty = ty
        self.tw = tw
        self.th = th
        self.cx = 0
        self.cy = 0
        self.done = False
    def put(self, color):
        if self.done:
            raise Exception('Tile size calc failed')

        # Convert a palette index to a random number
        # since we don't have the exact palette.
        if len(color) == 1:
            color = (color[0]*17 % 256,)

        I.putpixel((self.tx + self.cx, self.ty + self.cy), color)
        self.cx += 1
        if self.cx == self.tw:
            self.cx = 0
            self.cy += 1
            if self.cy == self.th:
                self.done = True
    def __repr__(self):
        return 'Tile(%d %d %d %d)' % (self.tx, self.ty, self.tw, self.th)

def maketiles(x, y, w, h):
    global I

    tiles = []
    tx = x
    ty = y
    while h > 0:
        th = min(64, h)
        h -= th

        while w > 0:
            tw = min(64, w)
            w -= tw
            tiles.append(Tile(I, tx, ty, tw, th))
            tx += tw
        tx = x
        ty += th
    return tiles

# Handle one frame buffer Rect.
def paste(x, y, w, h, f):
    # Skip some frames before opening the Internet.
    global counter
    counter += 1
    if counter < 110:
        pass

    tiles = maketiles(x, y, w, h)

    for tile in tiles:
        if len(f) == 0:
            break
        f, kind = db(f)
        if  kind == 128: # run-length
            while not tile.done:
                f, c = color(f)
                f, n = number(f)
                for _ in range(n):
                    tile.put(c)
        elif kind == 1: # solid fill
            f, c = color(f)
            while not tile.done:
                tile.put(c)
        elif 2 <= kind and kind <= 16: # packed palette
            palette = map(ord, f[:kind * 1])
            f = f[kind * 1:]

            if kind == 2:
                m = (tile.tw+7)/8*tile.th
                b = 1
            elif kind == 3 or kind == 4:
                m = (tile.tw+3)/4*tile.th
                b = 2
            else:
                m = (tile.tw+1)/2*tile.th
                b = 4

            bytearr = map(ord, f[:m])
            bitstrarr = map(lambda n: bin(n)[2:].rjust(8, '0'), bytearr)
            bitstr = ''.join(bitstrarr)
            f = f[m:]

            for i in range(m):
                pxstr = bitstr[:b]
                bitstr = bitstr[b:]
                c = (palette[int(pxstr, 2)],)
                tile.put(c)
        elif kind == 0:
            while not tile.done:
                f, c = color(f)
                tile.put(c)
        else:
            raise Exception("Bad TRLE kind " + str(kind))

# Handle one FrameBufferUpdate packet.
def msg(d):
    d, kind = db(d)
    if kind != 0:
        raise Exception("Not FramebufferUpdate")

    d, pad = db(d)
    d, count = dw(d)

    for _ in range(count):
        d, x = dw(d)
        d, y = dw(d)
        d, w = dw(d)
        d, h = dw(d)
        d, encoding = dd(d)
        print 'FRAME x=%d y=%d w=%d h=%d encoding=%d' % (x, y, w, h, encoding),
        if   encoding == 0xffffff11: # -239
            print 'Cursor pseudo'
            d = d[w*h:] # pixels
            d = d[(w+7)/8*h:] # mask
        elif encoding == 0xffffff18: # -232 -- not in the RFC docs.
            print 'Unknown'
        elif encoding == 0xffffff21: # -223
            print 'DesktopSize pseudo'
        elif encoding == 16:
            d, l = dd(d)
            print 'ZRLE'
            b = d[:l]
            d = d[l:]
            frame = z.decompress(b)
            paste(x, y, w, h, frame)
        else:
            raise Exception("Unknown subencoding")

    return d

# 1. Open prob.pcapng with Wireshark
# 2. Filter "tcp stream eq 0"
# 3. Select Server -> Client packets
# 4. Save the stream as binary file
with open('raw.bin') as f:
    d = f.read()

    # Skip handshake packets
    d = d[0x4b:]

trial = 0
while len(d) > 0:
    d = msg(d)
    I.save('img/%03d.png' % trial)
    trial += 1