import secrets
import socket
import struct
import sys


def read_line(conn):
    data = bytearray()
    while not data.endswith(b"\n"):
        part = conn.recv(1)
        if not part:
            raise ConnectionError("client disconnected")
        data.extend(part)
    return bytes(data[:-1])


def handle(conn):
    base = 0x555500000000 + (secrets.randbelow(0x10000) << 12)
    canary = secrets.randbits(56) << 8
    conn.sendall(f"entry={base + 0x1470:#x}\n".encode())
    conn.sendall(b"> ")
    if read_line(conn) != b"peek 64":
        conn.sendall(b"bad peek\n")
        return
    conn.sendall(b"word=" + struct.pack("<Q", canary).hex().encode() + b"\n> ")
    line = read_line(conn)
    if not line.startswith(b"submit "):
        conn.sendall(b"bad command\n")
        return
    try:
        payload = bytes.fromhex(line[7:].decode())
    except ValueError:
        conn.sendall(b"bad hex\n")
        return
    valid = len(payload) == 88
    valid &= payload[64:72] == struct.pack("<Q", canary)
    valid &= payload[80:88] == struct.pack("<Q", base + 0x11f0)
    conn.sendall(b"flag{glasshouse_training}\n" if valid else b"stack check failed\n")


def serve(host, port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
        listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        listener.bind((host, port))
        listener.listen()
        print("listening on", listener.getsockname(), flush=True)
        while True:
            conn, _ = listener.accept()
            with conn:
                try:
                    handle(conn)
                except ConnectionError:
                    pass


if __name__ == "__main__":
    host = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1"
    port = int(sys.argv[2]) if len(sys.argv) > 2 else 31341
    serve(host, port)