#!/usr/bin/env python3
"""
conman-console.py — Drive a serial console via a conman server, read-only.

Connects to a conmand server (RFC 127-style multiplexer) over the network,
opens a named console, sends the commands from a .cmds file, captures all
output, and disconnects. Requires no expect/tcl — pure stdlib PTY.

This replaces the old sw-capture-remote.sh workflow that killed the serial
device holder (conflicting with conman/ser2net). Instead, it talks to conman
over TCP, which multiplexes safely with other sessions.

All endpoints are configurable via environment variables so this works on
any network with a conman server:

    CONMAN_SERVER   conman server host:port   (default: via CONSOLE_HOST)
    CONSOLE         console name to open       (required)
    CMDS_FILE       file of commands to send   (required)
    TIMEOUT         overall timeout in seconds (default: 45)
    CMD_DELAY       seconds between commands   (default: 3)
    WAKE_DELAY      seconds after connect      (default: 2)

Usage:
    CONMAN_SERVER=console-host:7890 \\
    python3 conman-console.py --console pfv-core-sw01 --cmds switches/pfv-core-sw01.cmds

Lines starting with '!' or '#' in the cmds file are comments (skipped).
Blank lines are skipped. The conman escape sequence (&.) is sent automatically
to disconnect. A password prompt aborts immediately (we never send creds).

Exit codes:
    0  clean run
    1  usage / setup error
    2  could not connect to conman server
    3  timeout (partial output still printed)
    4  password prompt encountered (aborted)
"""
import argparse
import os
import pty
import re
import select
import sys
import time

PWD_RE = re.compile(rb"[Pp]assword:\s*$")
MORE_RE = re.compile(rb"--\s*[Mm]ore\s*--|[Mm]ore:\s*<space>")


def main():
    ap = argparse.ArgumentParser(
        description="Drive a conman console session read-only via PTY")
    ap.add_argument("--console", required=True,
                    help="console name (e.g. pfv-core-sw01)")
    ap.add_argument("--cmds", required=True,
                    help="command file (one command per line; !/# = comment)")
    ap.add_argument("--server",
                    default=os.environ.get("CONMAN_SERVER", ""),
                    help="conman server host:port (env: CONMAN_SERVER)")
    ap.add_argument("--timeout", type=int,
                    default=int(os.environ.get("TIMEOUT", "45")),
                    help="overall timeout seconds (env: TIMEOUT)")
    ap.add_argument("--cmd-delay", type=float,
                    default=float(os.environ.get("CMD_DELAY", "3")),
                    help="seconds between commands (env: CMD_DELAY)")
    ap.add_argument("--wake-delay", type=float,
                    default=float(os.environ.get("WAKE_DELAY", "2")),
                    help="seconds after connect before first command (env: WAKE_DELAY)")
    args = ap.parse_args()

    if not args.server:
        sys.stderr.write("ERROR: --server or CONMAN_SERVER env required\n")
        return 1

    with open(args.cmds) as f:
        cmds = [l.strip() for l in f
                if l.strip() and not l.strip().startswith(("!", "#"))]

    pid, fd = pty.fork()
    if pid == 0:
        try:
            os.execvp("conman",
                      ["conman", "-d", args.server, "-f", args.console])
        except OSError as e:
            sys.stderr.write(f"ERROR: cannot exec conman: {e}\n")
            os._exit(2)
        os._exit(2)

    output = b""
    cmd_queue = list(cmds)
    sent_disconnect = False
    start = time.time()
    last_action = 0.0
    phase = "connect"

    while time.time() - start < args.timeout:
        ready, _, _ = select.select([fd], [], [], 0.5)
        if ready:
            try:
                data = os.read(fd, 8192)
            except OSError:
                break
            if not data:
                break
            output += data

            if PWD_RE.search(output.split(b"\n")[-1] if output else b""):
                sys.stderr.write("[ABORT] password prompt detected — "
                                 "never sending credentials\n")
                os.write(fd, b"&.\n")
                break

            # Handle pagination: send space to continue
            if MORE_RE.search(output[-200:] if output else b""):
                os.write(fd, b" ")
                time.sleep(0.5)

        elapsed = time.time() - start
        gap = elapsed - last_action

        if phase == "connect" and gap >= args.wake_delay:
            os.write(fd, b"\n")
            phase = "send"
            last_action = elapsed
        elif phase == "send" and gap >= args.cmd_delay:
            if cmd_queue:
                cmd = cmd_queue.pop(0)
                os.write(fd, (cmd + "\n").encode())
                last_action = elapsed
            else:
                phase = "drain"
                last_action = elapsed
        elif phase == "drain" and gap >= args.cmd_delay:
            os.write(fd, b"&.\n")
            sent_disconnect = True
            phase = "done"
            last_action = elapsed
        elif phase == "done" and gap >= 2:
            break

    if not sent_disconnect:
        try:
            os.write(fd, b"&.\n")
        except OSError:
            pass

    try:
        os.close(fd)
    except OSError:
        pass
    try:
        os.waitpid(pid, 0)
    except ChildProcessError:
        pass

    sys.stdout.buffer.write(output)
    sys.stdout.flush()
    return 0


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