#!/usr/bin/env python3
"""Run ssh under a real pty of a known size and capture what it draws.

  ssh_pty.py <keyfile> <user> [seconds] [keys-to-send]

Keys are sent after 4s, then output is captured until the timeout.
"""
import os, pty, sys, time, select, struct, fcntl, termios, signal

key, user = sys.argv[1], sys.argv[2]
seconds = float(sys.argv[3]) if len(sys.argv) > 3 else 12.0
to_send = sys.argv[4] if len(sys.argv) > 4 else ""

pid, fd = pty.fork()
if pid == 0:
    os.execvp("ssh", [
        "ssh", "-tt",
        "-o", "StrictHostKeyChecking=no",
        "-o", "UserKnownHostsFile=/dev/null",
        "-o", "LogLevel=ERROR",
        "-p", "2222", "-i", key, f"{user}@127.0.0.1",
    ])

fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", 40, 120, 0, 0))

out = b""
deadline = time.time() + seconds
sent = False
while time.time() < deadline:
    if not sent and time.time() > deadline - seconds + 5:
        if to_send:
            os.write(fd, to_send.encode())
        sent = True
    r, _, _ = select.select([fd], [], [], 0.5)
    if r:
        try:
            chunk = os.read(fd, 65536)
        except OSError:
            break
        if not chunk:
            break
        out += chunk

try:
    os.kill(pid, signal.SIGKILL)
except ProcessLookupError:
    pass
os.waitpid(pid, 0)
sys.stdout.buffer.write(out)