#!/usr/bin/env python3
"""Drive a command on a real pty from a small expect/send script.

The private tier refuses to run without a terminal -- `[ -t 0 ]` is one of the
guards under test -- so every assertion about the password prompt, the retry
loop, the token file and the clone needs a pty to exist at all. `expect` is not
installed everywhere; python3 is, and the whole of what is needed here is
"read until this pattern, then write that line".

  ptydrive.py [--timeout SEC] --script FILE -- cmd [args...]

Script lines, blank lines and #-comments ignored:

  expect <python regex>    wait for it in everything the child has written
  send   <text>            write it, plus a newline
  sendraw <text>           write it with no newline
  close                    close the child's input (EOF)

Everything the child writes is copied to stdout, so the caller greps the
transcript exactly as it would grep any other command's output. Exit status is
the child's, or 3 for "a pattern never arrived" -- which names the pattern on
stderr rather than letting the run die of a later, unrelated timeout.
"""
import argparse
import os
import pty
import re
import select
import sys
import time

ap = argparse.ArgumentParser()
ap.add_argument("--timeout", type=float, default=20.0)
ap.add_argument("--script", required=True)
ap.add_argument("cmd", nargs=argparse.REMAINDER)
a = ap.parse_args()
cmd = a.cmd[1:] if a.cmd and a.cmd[0] == "--" else a.cmd
if not cmd:
    sys.exit("ptydrive.py: no command")

steps = []
for raw in open(a.script):
    line = raw.rstrip("\n")
    if not line.strip() or line.lstrip().startswith("#"):
        continue
    verb, _, rest = line.partition(" ")
    steps.append((verb, rest))

pid, fd = pty.fork()
if pid == 0:
    os.execvp(cmd[0], cmd)
    os._exit(127)

buf = ""
out = []
deadline = time.time() + a.timeout
eof = False


def pump(block=0.2):
    """Read whatever is there. Returns False once the child's side is gone."""
    global buf, eof
    r, _, _ = select.select([fd], [], [], block)
    if not r:
        return True
    try:
        chunk = os.read(fd, 65536)
    except OSError:
        chunk = b""
    if not chunk:
        eof = True
        return False
    text = chunk.decode("utf-8", "replace")
    buf += text
    out.append(text)
    return True


status = 0
try:
    for verb, arg in steps:
        if verb == "expect":
            pat = re.compile(arg)
            end = time.time() + a.timeout
            while not pat.search(buf):
                if time.time() > end:
                    sys.stderr.write("ptydrive.py TIMEOUT: never saw /%s/\n" % arg)
                    status = 3
                    raise SystemExit
                if not pump():
                    if pat.search(buf):
                        break
                    sys.stderr.write("ptydrive.py EOF before /%s/\n" % arg)
                    status = 4
                    raise SystemExit
            # Consume up to and including the match so a later `expect` for the
            # same text waits for a NEW one -- the retry loop asks for the same
            # three prompts twice, and matching the first round twice would
            # prove nothing.
            buf = buf[pat.search(buf).end():]
        elif verb in ("send", "sendraw"):
            data = arg + ("\n" if verb == "send" else "")
            os.write(fd, data.encode())
        elif verb == "close":
            os.close(fd)
            fd = -1
        else:
            sys.exit("ptydrive.py: unknown verb %r" % verb)
except SystemExit:
    pass

# Drain whatever is left, then reap.
while not eof and time.time() < deadline:
    if not pump(0.1):
        break
try:
    _, wstatus = os.waitpid(pid, 0)
    child = os.waitstatus_to_exitcode(wstatus)
except ChildProcessError:
    child = 0
sys.stdout.write("".join(out))
sys.stdout.flush()
sys.exit(status if status else (child if child >= 0 else 128 - child))