Fleet automation: subprocess fan-out, SSH and HTTP at scale

# Pinned clients for driving a fleet. Installed into the lab's venv.
paramiko==5.0.0
httpx==0.28.1

#!/usr/bin/env bash
# A stand-in for a remote probe. The target name picks the behaviour, so the fan-out meets fast
# hosts, a failing host, one that never stops printing, one that prints non-UTF-8 bytes, and one
# that hangs while holding a child process.
set -euo pipefail
case "$1" in
  slow-*)
    # a child that outlives a naive kill of the parent: sleep in the background, then wait on our own
    sleep 300 &
    sleep 300
    ;;
  fail-*)
    echo "unreachable"
    exit 3
    ;;
  big-*)
    yes "$1 log line"   # endless output: only a byte cap stops it before the timeout
    ;;
  bin-*)
    printf 'bad banner \377\376\n'   # not valid UTF-8
    exit 2
    ;;
  *)
    echo "ok $1"
    ;;
esac

"""Run one probe in its own process group, with a deadline and a cap on how much output we hold.
 
Whatever ends the run early (the timeout, too much output, a cancellation from the caller or a
SIGTERM), the finally block kills the probe's whole process group, so no child outlives it."""
import asyncio
import contextlib
import os
import signal
 
TIMEOUT = 2.0        # seconds one probe may take
MAX_OUTPUT = 64_000  # bytes of output we are willing to hold per probe
 
 
async def read_capped(stream: asyncio.StreamReader, limit: int) -> bytes | None:
    """Read to end of file, or return None as soon as the output passes limit bytes."""
    data = bytearray()
    while chunk := await stream.read(65536):
        data += chunk
        if len(data) > limit:
            return None
    return bytes(data)
 
 
async def run_probe(target: str) -> tuple[str, str]:
    proc = await asyncio.create_subprocess_exec(
        "./probe.sh", target,
        stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT,
        start_new_session=True,  # own process group, so killpg reaches every child it starts
    )
    finished = False
    try:
        async with asyncio.timeout(TIMEOUT):
            out = await read_capped(proc.stdout, MAX_OUTPUT)
            if out is None:
                return "too much output", f"over {MAX_OUTPUT} bytes, probe killed"
            rc = await proc.wait()
        finished = True
        text = out.decode(errors="replace").strip()  # a target's bytes may not be UTF-8
        return ("ok" if rc == 0 else f"exit {rc}"), text
    except TimeoutError:
        return "timed out", f"after {TIMEOUT:.0f}s, group killed"
    finally:
        if not finished:  # timeout, output cap, cancellation or SIGTERM: stop the whole group
            with contextlib.suppress(ProcessLookupError):  # the group may be gone already
                os.killpg(proc.pid, signal.SIGKILL)  # signal the whole process group
            # read what is left in the pipe (its writers are dead, so it ends) and reap the probe;
            # proc.wait() alone can hang here, because it also waits for the pipe to close
            await proc.communicate()

"""Fan a probe out across many targets: at most CONCURRENCY at once, a result per target, a non-zero
exit on partial failure, and a clean stop on SIGTERM that leaves no probe running."""
import asyncio
import signal
import sys
 
from runprobe import run_probe
 
CONCURRENCY = 5  # probes in flight at once, however many targets there are
TARGETS = ([f"web{i:02d}" for i in range(12)]
           + ["fail-db1", "fail-db2", "big-log1", "bin-fw1", "slow-gw1", "slow-gw2"])
 
 
async def bounded(target: str, sem: asyncio.Semaphore) -> tuple[str, str, str]:
    async with sem:
        return (target, *await run_probe(target))
 
 
async def main() -> int:
    # SIGTERM (systemctl stop, timeout, kill) or Ctrl-C cancels this task. The TaskGroup then cancels
    # every probe task, and each one's finally kills its process group.
    main_task, received = asyncio.current_task(), []
 
    def stop(sig: signal.Signals) -> None:
        received.append(sig)
        main_task.cancel()
 
    for sig in (signal.SIGTERM, signal.SIGINT):
        asyncio.get_running_loop().add_signal_handler(sig, stop, sig)
    sem = asyncio.Semaphore(CONCURRENCY)
    tasks = []
    try:
        async with asyncio.TaskGroup() as tg:  # an unexpected error in one task cancels the rest
            tasks = [tg.create_task(bounded(t, sem)) for t in TARGETS]
    except asyncio.CancelledError:
        if not received:
            raise
        stopped = sum(t.cancelled() for t in tasks)
        print(f"stopped by {received[0].name}: {stopped} unfinished probe(s) cancelled", file=sys.stderr)
        return 128 + received[0]  # 143 for SIGTERM, 130 for SIGINT, as a shell reports them
    results = [t.result() for t in tasks]
    bad = [r for r in results if r[1] != "ok"]
    print(f"{len(results) - len(bad)} ok, {len(bad)} not ok")
    for target, status, detail in bad:
        print(f"  {target:9} {status}: {detail}")
    return 1 if bad else 0  # partial failure is a non-zero exit
 
 
if __name__ == "__main__":
    sys.exit(asyncio.run(main()))

#!/usr/bin/env bash
# Build two throwaway SSH targets in the current directory. Run it as yourself, not with sudo.
# Host A will listen on 127.0.0.1:18823 and host B on 127.0.0.1:18824, each with its own host key.
# known_hosts pins host A's key for BOTH ports on purpose: host B then presents a key that does
# not match its pin, which is the failure ssh_fleet.py must refuse.
set -euo pipefail
dir=$PWD
user=$(id -un)
if [[ -e hostkey_a ]]; then
  echo "targets already set up in $dir (remove hostkey_* and id_lab* to start again)" >&2
  exit 1
fi
 
ssh-keygen -q -t ed25519 -f hostkey_a -N '' -C lab-host-a
ssh-keygen -q -t ed25519 -f hostkey_b -N '' -C lab-host-b
ssh-keygen -q -t ed25519 -f id_lab -N '' -C lab-client  # lab only: no passphrase
# restrict: no pty, no forwarding; from=: the key is refused from anywhere but this machine
printf 'restrict,from="127.0.0.1" %s\n' "$(cat id_lab.pub)" > authorized_keys
chmod 0600 hostkey_a hostkey_b id_lab authorized_keys
 
pin_a=$(cut -d' ' -f1,2 hostkey_a.pub)
printf '[127.0.0.1]:18823 %s\n[127.0.0.1]:18824 %s\n' "$pin_a" "$pin_a" > known_hosts
printf '127.0.0.1 18823\n127.0.0.1 18824\n' > targets.txt
 
for target in a:18823 b:18824; do
  id=${target%%:*} port=${target##*:}
  cat > "sshd_$id.conf" <<CONF
Port $port
ListenAddress 127.0.0.1
HostKey $dir/hostkey_$id
PidFile $dir/sshd_$id.pid
AuthorizedKeysFile $dir/authorized_keys
AllowUsers $user
PasswordAuthentication no
KbdInteractiveAuthentication no
UsePAM yes
StrictModes no
PrintMotd no
CONF
done
echo "keys, known_hosts, targets.txt, sshd_a.conf and sshd_b.conf are in $dir"

"""How every script in this lesson connects: pinned host keys, the lab key, bounded connect."""
import getpass
 
import paramiko
 
 
def connect(host: str, port: int) -> paramiko.SSHClient:
    client = paramiko.SSHClient()
    client.load_host_keys("known_hosts")                        # only these keys are trusted
    client.set_missing_host_key_policy(paramiko.RejectPolicy())  # an unknown host key fails closed
    try:
        client.connect(host, port=port, username=getpass.getuser(),
                       key_filename="id_lab",  # lab only: an unencrypted key file on disk
                       allow_agent=False, look_for_keys=False,
                       timeout=5, banner_timeout=5, auth_timeout=5)
    except BaseException:
        client.close()
        raise
    return client

"""Run one command across an SSH fleet: a bounded thread pool, a deadline on every host, capped
output, a result per host and a non-zero exit on partial failure."""
import shlex
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
 
import paramiko
 
from labssh import connect
 
CONCURRENCY = 10     # hosts in flight at once; paramiko blocks, so each one gets a thread
DEADLINE = 10        # seconds the command may run, enforced on the host by timeout(1)
MAX_OUTPUT = 64_000  # bytes of output kept per host
 
 
def check_host(host: str, port: int, command: str) -> tuple[bool, str]:
    name = f"{host}:{port}"
    try:
        with connect(host, port) as client:  # closes the connection on every path
            remote = f"timeout {DEADLINE} sh -c {shlex.quote(command)}"
            _, stdout, _ = client.exec_command(remote, timeout=DEADLINE + 5)
            stdout.channel.set_combine_stderr(True)  # one stream, so an unread stderr cannot stall us
            out = stdout.read(MAX_OUTPUT + 1)         # drain the output before the exit status
            if len(out) > MAX_OUTPUT:
                return False, f"{name}: ERROR output over {MAX_OUTPUT} bytes"
            rc = stdout.channel.recv_exit_status()
    except paramiko.BadHostKeyException as err:
        return False, f"{name}: ERROR host key mismatch ({err.expected_key.get_name()})"
    except Exception as err:  # refused, timed out, auth failed: a result for this host, not a crash
        return False, f"{name}: ERROR {type(err).__name__}: {err}"
    text = out.decode(errors="replace").strip()  # a host's bytes may not be UTF-8
    if rc != 0:
        return False, f"{name}: exit {rc} {text}".rstrip()
    return True, f"{name}: {text}"
 
 
def main() -> int:
    command = sys.argv[1]
    targets = [line.split() for line in Path("targets.txt").read_text().splitlines() if line.strip()]
    with ThreadPoolExecutor(max_workers=CONCURRENCY) as pool:
        results = list(pool.map(lambda t: check_host(t[0], int(t[1]), command), targets))
    for _, line in results:
        print(line)
    ok = sum(good for good, _ in results)
    print(f"{ok}/{len(results)} hosts ok")
    return 0 if ok == len(results) else 1  # any unreachable, mismatched or failing host is a non-zero exit
 
 
if __name__ == "__main__":
    sys.exit(main())

"""exec_command runs through the remote login shell, so the list-form protection does not apply.
 
Any untrusted value put into a remote command must be quoted with shlex.quote, or a shell
metacharacter in it starts a second command on the far side. Usage: remote_quote.py VALUE unsafe|safe"""
import shlex
import sys
 
from labssh import connect
 
value, mode = sys.argv[1], sys.argv[2]
if mode == "unsafe":
    command = f"wc -l {value}"               # value pasted straight into the remote command line
else:
    command = f"wc -l {shlex.quote(value)}"  # value folded into a single, inert argument
 
with connect("127.0.0.1", 18823) as client:
    _, stdout, stderr = client.exec_command(command, timeout=5)
    out = stdout.read().decode(errors="replace")  # a few bytes each: reading one stream, then the
    err = stderr.read().decode(errors="replace")  # other, is safe only for small output (ssh_drain.py)
    rc = stdout.channel.recv_exit_status()
 
print(f"sent   : {command}")
print(f"result : rc={rc} out={out.strip()!r} err={err.strip()!r}")

"""Read a remote command's output before its exit status, without stalling on stderr.
 
stdout and stderr share one SSH flow-control window, and paramiko reopens it only as you read.
The command writes about 6.9 MB to stderr and one line to stdout. Usage:
  ssh_drain.py status-first | stdout-first | combined"""
import sys
 
from labssh import connect
 
MODE = sys.argv[1]
COMMAND = "seq 1 1000000 >&2; echo done"
 
with connect("127.0.0.1", 18823) as client:
    _, stdout, stderr = client.exec_command(COMMAND, timeout=5)  # 5 s without data raises
    chan = stdout.channel
    try:
        if MODE == "status-first":
            rc = chan.recv_exit_status()    # waits for the command, which waits for a reader
            out = stdout.read()
        elif MODE == "stdout-first":
            out = stdout.read() + stderr.read()  # stdout waits for EOF while stderr fills the window
            rc = chan.recv_exit_status()
        else:
            chan.set_combine_stderr(True)   # stderr joins stdout: one stream to drain
            out = stdout.read()
            rc = chan.recv_exit_status()    # only after the output has been read
    except TimeoutError:
        held = len(chan.recv_stderr(8 * 2**20))  # what arrived on stderr and was never read
        print(f"{MODE}: stalled, no data for 5 s; {held} bytes of stderr were waiting unread")
        sys.exit(1)
 
print(f"{MODE}: read {len(out)} bytes, exit status {rc}")

"""A pooled HTTP client with connection limits and a timeout, reusing sockets across many requests.
 
httpx over requests here for two reasons: it exposes the pool size directly as httpx.Limits, and it
has one API for the sync client and the async client, so the same code shape fits the asyncio fan-out
above. requests.Session also pools, but the pool size is set indirectly through a mounted adapter."""
import httpx
 
BASE = "http://127.0.0.1:18825"
 
 
def main() -> None:
    limits = httpx.Limits(max_connections=5, max_keepalive_connections=5)
    timeout = httpx.Timeout(5.0, connect=2.0)  # a timeout on every request; never wait forever
    with httpx.Client(base_url=BASE, limits=limits, timeout=timeout) as client:
        ok = sum(client.get("/ok").status_code == 200 for _ in range(50))
        fail = client.get("/fail")  # a 5xx is surfaced here; retrying it belongs to "Failure design"
        stats = client.get("/stats").json()
    print(f"{ok}/50 ok, /fail -> {fail.status_code}")
    print(f"server saw {stats['connections']} connection(s) for {stats['requests']} requests")
 
 
if __name__ == "__main__":
    main()