PFVCluster/scripts/sw-capture.py

#!/usr/bin/env python3
"""
sw-capture.py - drive a Dell switch over a serial console and log all output.

Read-only. Sends ONLY the commands listed in the supplied .cmds file
(comment lines starting with '!' and blank lines are skipped). Handles
`--More--` pagination by sending a space. Aborts cleanly on any password
prompt (we never supply credentials).

Pure stdlib (termios + select). No pyserial/expect required.

Exit codes:
   0  clean run, every command saw a prompt again
   2  could not synchronize with a prompt during wake
   3  one or more commands timed out (log still written)
   4  password prompt encountered (aborted)

Usage:
    sw-capture.py --device /dev/ttyUSB2 --baud 9600 \\
                  --cmds pfv-core-sw01.cmds --log /root/pfv-core-sw01.log
"""
import argparse
import os
import re
import select
import sys
import termios
import time

PROMPT_RE = re.compile(rb'[>#]\s*$')              # ends in # or > + spaces
MORE_RE   = re.compile(rb'--\s*More\s*--')        # pagination prompt
PWD_RE    = re.compile(rb'[Pp]assword:\s*$')      # enable / login password

BAUDS = {
    '9600':    termios.B9600,
    '19200':   termios.B19200,
    '38400':   termios.B38400,
    '57600':   termios.B57600,
    '115200':  termios.B115200,
}


def log(msg, level='INFO'):
    sys.stderr.write(f'[{level}] {msg}\n')
    sys.stderr.flush()


def open_port(device, baud):
    """Open the serial device raw at the requested baud, 8N1, no flow ctrl."""
    fd = os.open(device, os.O_RDWR | os.O_NOCTTY)
    try:
        attrs = termios.tcgetattr(fd)
    except termios.error:
        log(f'{device} is not a termios-capable device', 'WARN')
        return fd

    # raw input
    attrs[0] &= ~(termios.IGNBRK | termios.BRKINT | termios.PARMRK |
                  termios.ISTRIP | termios.INLCR | termios.IGNCR |
                  termios.ICRNL | termios.IXON)
    # raw output
    attrs[1] &= ~termios.OPOST
    # 8N1, enable receiver, ignore modem control lines
    attrs[2] &= ~(termios.CSIZE | termios.PARENB | termios.CSTOPB)
    attrs[2] |= termios.CS8 | termios.CREAD | termios.CLOCAL
    # raw local
    attrs[3] &= ~(termios.ECHO | termios.ECHONL | termios.ICANON |
                  termios.ISIG | termios.IEXTEN)
    # non-blocking-ish reads (select is the primary gate)
    attrs[6][termios.VMIN] = 0
    attrs[6][termios.VTIME] = 1

    b = BAUDS.get(str(baud))
    if b is None:
        raise SystemExit(f'unsupported baud: {baud}')
    # Set ispeed/ospeed directly on the attribute list. (Equivalent to
    # termios.cfsetispeed/cfsetospeed, which are missing on some Python
    # builds — e.g. the one on pfv-tsys4.)
    attrs[4] = b   # ispeed
    attrs[5] = b   # ospeed
    termios.tcsetattr(fd, termios.TCSANOW, attrs)
    return fd


def read_chunk(fd, timeout):
    """Read whatever arrives within `timeout`. Extends briefly on activity."""
    buf = b''
    deadline = time.time() + timeout
    while True:
        remaining = deadline - time.time()
        if remaining <= 0:
            return buf
        r, _, _ = select.select([fd], [], [], min(0.5, remaining))
        if not r:
            if buf:
                return buf
            continue
        try:
            chunk = os.read(fd, 4096)
        except OSError:
            return buf
        if not chunk:
            return buf
        buf += chunk
        # keep collecting as long as bytes are flowing
        deadline = time.time() + 0.3


def drain(fd, timeout=1.0):
    total = 0
    while True:
        b = read_chunk(fd, timeout=timeout)
        if not b:
            return total
        total += len(b)


def send(fd, s):
    if isinstance(s, str):
        s = s.encode()
    os.write(fd, s)


def wait_for(fd, regex, timeout, on_more=None, on_pwd=None):
    """Read until `regex` matches the tail of the buffer, or timeout."""
    buf = b''
    deadline = time.time() + timeout
    while time.time() < deadline:
        remaining = deadline - time.time()
        chunk = read_chunk(fd, timeout=min(1.0, remaining))
        if chunk:
            buf += chunk
            tail64 = buf[-64:]
            tail32 = buf[-32:]
            tail128 = buf[-128:]
            if on_more and MORE_RE.search(tail64):
                on_more(fd)
                continue
            if on_pwd and PWD_RE.search(tail32):
                on_pwd(buf)
                return buf, 'pwd'
            if regex.search(tail128):
                return buf, 'ok'
    return buf, 'timeout'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--device', required=True)
    ap.add_argument('--baud', type=int, default=9600)
    ap.add_argument('--cmds', required=True)
    ap.add_argument('--log',  required=True)
    ap.add_argument('--per-cmd-timeout', type=float, default=45.0)
    ap.add_argument('--wake-timeout',    type=float, default=15.0)
    ap.add_argument('--session-max',     type=float, default=600.0)
    args = ap.parse_args()

    cmds = []
    with open(args.cmds) as f:
        for raw in f:
            s = raw.rstrip('\n').strip()
            if not s or s.startswith('!'):
                continue
            cmds.append(s)
    log(f'parsed {len(cmds)} commands from {args.cmds}')

    logf = open(args.log, 'wb', buffering=0)

    def w(b):
        if isinstance(b, str):
            b = b.encode()
        logf.write(b)

    w(f'==== sw-capture {time.strftime("%Y-%m-%d %H:%M:%S")} ====\n')
    w(f'device={args.device} baud={args.baud} cmds={args.cmds} '
      f'n={len(cmds)} per_cmd_timeout={args.per_cmd_timeout}\n\n')

    fd = open_port(args.device, args.baud)
    log(f'opened {args.device} @ {args.baud} 8N1 raw')

    session_start = time.time()
    abort = False

    def on_more(fd_):
        log('--More-- -> space')
        send(fd_, b' ')

    def on_pwd(buf):
        nonlocal abort
        abort = True
        log('password prompt detected (enable or login) - aborting; '
            'no credentials supplied', 'ERROR')
        w(buf)
        w(b'\n[PASSWORD PROMPT - ABORTED]\n')

    # WAKE: nudge with Ctrl-C + Enter, look for any prompt
    drain(fd, 0.5)
    synced = False
    wake_deadline = time.time() + args.wake_timeout
    attempt = 0
    while time.time() < wake_deadline:
        attempt += 1
        send(fd, b'\x03')
        time.sleep(0.2)
        send(fd, b'\r')
        buf, status = wait_for(fd, PROMPT_RE, timeout=3.0,
                               on_more=on_more, on_pwd=on_pwd)
        w(buf)
        if status == 'pwd':
            logf.close(); os.close(fd); sys.exit(4)
        if status == 'ok':
            synced = True
            log(f'prompt synced after {attempt} attempt(s)')
            break
    if not synced:
        w(b'\n[NO PROMPT - ABORT]\n')
        log('no prompt detected during wake window', 'ERROR')
        logf.close(); os.close(fd); sys.exit(2)

    # RUN commands verbatim from the .cmds list
    failures = 0
    for idx, cmd in enumerate(cmds, 1):
        if time.time() - session_start > args.session_max:
            log('session_max exceeded - stopping early', 'ERROR')
            w(b'\n[SESSION_MAX - STOP]\n')
            break
        if abort:
            break
        log(f'[{idx}/{len(cmds)}] {cmd}')
        send(fd, cmd + '\r')
        buf, status = wait_for(fd, PROMPT_RE,
                                timeout=args.per_cmd_timeout,
                                on_more=on_more, on_pwd=on_pwd)
        w(buf)
        if status == 'pwd':
            failures += 1
            break
        if status == 'timeout':
            log(f'timeout after: {cmd}', 'WARN')
            failures += 1
            # try to resync: Ctrl-C + drain
            send(fd, b'\x03')
            time.sleep(0.3)
            drain(fd, 0.5)

    w(f'\n==== end {time.strftime("%Y-%m-%d %H:%M:%S")} '
      f'failures={failures} ====\n')
    logf.close()
    os.close(fd)
    log(f'done -> {args.log}  failures={failures}')
    sys.exit(0 if failures == 0 else 3)


if __name__ == '__main__':
    main()