rshall

#!/usr/bin/env python3
# rshall (part of ossobv/vcutil) // wdoekes/2024-2026 // Public Domain
#
# Run script on many hosts at once and collect output.
#
# This can be used for one-off scripts that you don't want to build any
# ansible/chef/puppet scripts for.
#
# Invocation:
#
#   rshall [-c] [-f] [-t alt_targets_name] script_name [alt_project_name]
#
# It will read "<script_name>.targets" for target hosts anda "<script_name>"
# for the script to execute. The author suggests using ".rsh" as the
# script_name suffix (for "remote-shell-script").
#
# Example: you want to confirm that the DNS is set the same everywhere.
# rshall can collect the info using multiple ssh connections.
#
# - Create 'dns-config.rsh':
#
#   #!/bin/sh
#   sed -e '/^nameserver/!d;s/[[:blank:]]/;/' /etc/resolv.conf
#   resolvectl status 2>/dev/null |
#       awk '/DNS Servers:/{gsub(/[^:]*: */, "");print "resolvectl;" $0}'
#
# - Test your script:
#
#   $ sudo sh -c ./dns-config.rsh
#   nameserver;127.0.0.53
#   resolvectl;8.8.8.8 1.1.1.1
#
# - Create 'dns-config.rsh.targets':
#
#   #hostname;#ip4
#   foo.com;1.2.3.4
#   myuser@bar.com;55.55.55.55
#
# - Run the script on all target hosts, using rshall:
#
#   $ USERNAME=root rshall dns-config.rsh
#   ...
#   2 success, 0 fails, 0 earlier success, 0 earlier fails
#   See dns-config.rsh.out for output
#
# - Check the output:
#
#   $ cat dns-config.rsh.out
#   #hostname;#ip4;#result
#   foo.com;1.2.3.4;nameserver;127.0.0.53
#   foo.com;1.2.3.4;resolvectl;8.8.8.8 1.1.1.1
#   myuser@bar.com;55.55.55.55;nameserver;127.0.0.53
#   myuser@bar.com;55.55.55.55;resolvectl;8.8.8.8 1.1.1.1
#
# Caveats:
#
# - Right now, you're stuck with semi-colon delimited CSV for targets (input)
#   and output. This is a fair compromise between readability and automation.
# - Configuring the USERNAME is cumbersome if you're not using the same
#   username everywhere. You can set the user as "myuser@hostname" instead.
# - Any output on stderr is now considered a failure. This abides by the
#   principle of least surprise. If you want to ignore stderr, redirect
#   stderr to /dev/null where appropriate.
# - The script must start with #!/bin/sh for now. Except if you're using
#   "!#/usr/bin/rshall", in which case #!/bin/sh is silently assumed.
#
# Environment variables:
#
# - USERNAME sets the default ssh username. The host-entry can override this
#   by using username@host.
# - SUDO_COMMAND defaults to "sudo". You can set it to "pkexec" or
#   "sudo -u someoneelse" to change the default sudo call.
# - PRE_SUDO_COMMAND is a free-form shell command to run before becoming root
#   using sudo. (Technically, it replaces SUDO_COMMAND with a complicated
#   shell spaghetti. They cannot be combined at the moment. If you also want
#   a different SUDO_COMMAND you'll have to create the SUDO_COMMAND spaghetti
#   yourself.)
# - For now, we enforce that PRE_SUDO_COMMAND is set. If you don't want to use
#   it, set it to the empty string. This makes it harder to forget to set it.
#
import sys

from argparse import ArgumentParser
from collections import namedtuple
from ipaddress import AddressValueError, IPv4Address
from multiprocessing import Process, Pipe
from os import environ, path, rename, unlink
from shlex import quote as shell_escape, split as shell_split
from subprocess import CalledProcessError, PIPE, Popen, TimeoutExpired
from time import sleep
from warnings import warn


Target = namedtuple('Target', 'hostname ip')
ProcInfo = namedtuple('ProcInfo', 'target process pipe')
ProcResult = namedtuple('ProcResult', 'target is_success value')

VERBOSE = False
CONCURRENT_JOBS = 40
CONNECT_TIMEOUT = 10
SCRIPT_TIMEOUT = 60


class Runner:
    STATUS_TICK = 0
    STATUS_NEW = 1
    STATUS_SUCCESS = 2
    STATUS_FAILURE = 3

    def __init__(self):
        self._script = '''\
#!/bin/sh
set -eu
echo "I am a sample script on $(hostname)"
'''
        self._targets = []
        self.successes = {}
        self.failures = {}
        self.concurrency = CONCURRENT_JOBS

        self._procinfos = []

    def set_script(self, script):
        self._script = script

    def add_targets(self, targets):
        self._targets.extend(targets)

    def _cycle(self, state):
        if VERBOSE:
            return

        if state == self.STATUS_NEW:
            print('+', end='', file=sys.stderr)
        elif state == self.STATUS_SUCCESS:
            print('-', end='', file=sys.stderr)
        elif state == self.STATUS_FAILURE:
            print('x', end='', file=sys.stderr)
        elif state == self.STATUS_TICK:
            print('.', end='', file=sys.stderr)
        else:
            raise NotImplementedError

        sys.stderr.flush()

    def _reap(self):
        for idx, procinfo in reversed(list(enumerate(self._procinfos))):
            if not procinfo.process.is_alive():
                result = procinfo.pipe.recv()
                assert result.target == procinfo.target, (procinfo, result)
                if result.is_success:
                    self.successes[result.target] = result.value
                else:
                    self.failures[result.target] = result.value
                status = procinfo.process.join()
                procinfo.pipe.close()
                assert status is None, (status, procinfo)
                self._procinfos.pop(idx)

                if VERBOSE:
                    maybe_ok = ('OK' if result.is_success else 'FAIL')
                    print(result.target.hostname, maybe_ok, file=sys.stderr)

                if result.is_success:
                    self._cycle(self.STATUS_SUCCESS)
                else:
                    self._cycle(self.STATUS_FAILURE)

    def run(self):
        for target in self._targets:
            # BUG: Some processes might have data in pipe, but we don't fetch
            # it yet. That causes an infinite stall. The current workaround is
            # to cap the output buffer to below 64KiB. We don't want such big
            # output anyway.
            self._reap()  # try reaping reapable
            while len(self._procinfos) >= self.concurrency:
                self._reap()
                sleep(0.1)

            parent_conn, child_conn = Pipe()
            proc = Process(
                target=run_remote, args=(
                    parent_conn, child_conn, target, self._script))
            self._procinfos.append(
                ProcInfo(target=target, process=proc, pipe=parent_conn))
            proc.start()
            child_conn.close()
            self._cycle(self.STATUS_NEW)
            sleep(0.5)  # some sleep before spawning too much at once

        while self._procinfos:
            self._reap()
            sleep(0.5)
            self._cycle(self.STATUS_TICK)

        if not VERBOSE:
            print('', file=sys.stderr)  # last _cycle


def get_sudo_command(env):
    if env.get('SUDO_COMMAND'):
        # Prefer not to do this. This command could be ugly and error-prone.
        # Unless you really just want to use a different sudo.
        # SUDO_COMMAND=pkexec
        if env.get('PRE_SUDO_COMMAND'):
            warn('Both SUDO_COMMAND and PRE_SUDO_COMMAND set, using former')
        sudo_cmd = tuple(shell_split(env['SUDO_COMMAND']))

    elif env.get('PRE_SUDO_COMMAND'):
        # This command runs _before_ the script. The only use for it right
        # now is to call a here-I-am script.
        # PRE_SUDO_COMMAND='osso-change write-log \#1234 june-test'
        pre_sudo_str = env['PRE_SUDO_COMMAND']
        sudo_cmd = (
            'sh',
            '-c',
            # NOTE: Complicated shell-escape. Be careful what you put here
            # and make it silent.
            # We expect all targets to have the PRE_SUDO_COMMAND working.
            shell_escape(f'{pre_sudo_str} && exec sudo "$@"'),
            '--',
        )

    elif 'PRE_SUDO_COMMAND' not in env:
        # Set this thing to empty to continue. For some systems we might need
        # PRE_SUDO_COMMAND for bookkeeping. Force it to be set.
        print('error: Did you forget PRE_SUDO_COMMAND?', file=sys.stderr)
        exit(1)

    else:
        sudo_cmd = ('sudo',)

    return sudo_cmd


def run_remote(parent_conn, child_conn, target, script):
    quoted_script = shell_escape(script)

    parent = child_conn
    parent_conn.close()

    out, err = '', ''

    env = environ.copy()  # keep SSH_AGENT stuff

    # Remove language stuff that causes different output.
    for key in list(env.keys()):
        if key.startswith('LC_') or key == 'LANG':
            del env[key]
    env['LC_ALL'] = 'C'

    # The script doesn't have a controlling terminal. OpenSSH might
    # fall back to ssh-askpass (which may not be installed) to prompt
    # for passwords or Yubikey touch confirmation.
    # Without, we still require touch, but we won't log missing ssh-askpass.
    # #env.pop('DISPLAY', None)
    # #env.pop('WAYLAND_DISPLAY', None)
    env['SSH_ASKPASS'] = '/bin/true'

    # The hostname may be "hostname" or "user@hostname".
    if '@' in target.hostname:
        username, ip = target.hostname.split('@')
    else:
        username, ip = env['USERNAME'], target.hostname
    if target.ip:
        ip = target.ip  # always prefer IP if available

    ssh_command = (
        'ssh', '-l', username,
        '-oConnectTimeout={}'.format(CONNECT_TIMEOUT),
        '-oLogLevel=error',  # suppress banners on stderr
        '-oPasswordAuthentication=no',
        '-oUpdateHostKeys=yes',
        '-oStrictHostKeyChecking=accept-new',
        ip, '--',
    )
    sudo_command = get_sudo_command(env)
    shell_command = ('sh', '-c', quoted_script)
    command = ssh_command + sudo_command + shell_command

    try:
        proc = Popen(
            command, stdin=None, stdout=PIPE, stderr=PIPE, text=True, env=env)
        try:
            out, err = proc.communicate(timeout=SCRIPT_TIMEOUT)
        except TimeoutExpired:
            proc.kill()
            out, err = proc.communicate(timeout=3)  # should be instant?
            err += '\n\nTIMEOUT'
        if proc.wait() != 0 or err:
            if proc.returncode == 0 and all(
                    # If you have multiple yubikeys, you might get:
                    #   sign_and_send_pubkey: signing failed for ED25519-SK
                    #     "/home/jdoe/.ssh/id_yubi1234t": device not found
                    line.startswith('sign_and_send_pubkey: signing failed for')
                    for line in err.rstrip().split('\n')):
                # Don't fail on just the "sign_and_send_pubkey" stderr.
                pass
            else:
                # Fail on ANYTHING seen on stderr.
                raise CalledProcessError(
                    proc.returncode, ['ssh', target.ip, '<script>'], out, err)
    except CalledProcessError as e:
        err = f'{e}: err={e.stderr!r}; out={e.output!r}'
        result = ProcResult(target=target, is_success=False, value=err)
    except Exception as e:
        err = f'{target}: {e}'
        result = ProcResult(target=target, is_success=False, value=err)
    else:
        result = ProcResult(target=target, is_success=True, value=out)

    # BUG: If we send too much data here, the pipe stalls. Default cap is about
    # 64KiB. Half of it is already way too much though.
    if len(result.value) > 32768:
        result = ProcResult(
            target=target, is_success=False, value=(
                f'output exceeded (len={len(result.value)}, '
                f'{result.is_success=}, '
                f'value={repr(result.value[0:500]+"...")})'))

    parent.send(result)
    parent.close()


class Project:
    def __init__(
            self, script_name, alt_project_name=None,
            alt_targets_name=None, args=None):
        # A bit convoluted how we set the parameters.
        self.script_name = script_name      # foobar.rsh
        self.project_name = script_name     # foobar.rsh

        # If alt_project_name is set, use that. If it included the ".targets"
        # suffix, drop that.
        if alt_project_name:
            if (alt_project_name.endswith('.targets')
                    and len(alt_project_name) > 8
                    and path.exists(alt_project_name)):
                alt_project_name = alt_project_name[0:-8]
            self.project_name = alt_project_name
            # Check that user is not trying to run a different script after
            # all.
            if self.script_name != self.project_name and path.exists(
                    self.script_name):
                raise ValueError(
                    'seeing two scripts; unsure which one you want to run')

        # Target file is script.rsh.targets, or project.targets, or custom.
        self.targets_name = (
            alt_targets_name or f'{self.project_name}.targets')

        self.script = self._load_script(self.script_name)
        self.targets = self._load_targets(self.targets_name)
        self.skipped = 0

        # Process other args.
        self.to_stdout = (args and args.stdout)
        if args and args.force:
            if not self.to_stdout:
                self._unlink_or_ignore(self.get_successes_name())
            self._unlink_or_ignore(self.get_failures_name())

        # Load earlier successes and drop from targets
        if self.to_stdout:
            self.successes_skip = 0
        else:
            successes = self._load_lines(self.get_successes_name())
            self.successes_skip = self.prune_targets(
                self._lines_to_targets(successes))
        # Load earlier failures and drop from targets
        failures = self._load_lines(self.get_failures_name())
        self.failures_skip = self.prune_targets(
            self._lines_to_targets(failures))

    def get_successes_name(self):
        if self.to_stdout:
            return object()  # something that fails to open
        return f'{self.project_name}.out'

    def get_failures_name(self):
        return f'{self.project_name}.err'

    def prune_targets(self, remove_from_targets):
        new_targets = []
        remove_from_targets = set(remove_from_targets)
        skipped = 0
        for target in self.targets:
            if target in remove_from_targets:
                skipped += 1
            else:
                new_targets.append(target)
        self.targets = new_targets
        return skipped

    def update_successes(self, new_successes):
        if not self.to_stdout:
            successes = self._load_lines(self.get_successes_name())
        else:
            successes = []

        if not successes:
            successes.append('#hostname;#ip4;#result')

        for target, result in sorted(new_successes.items()):
            for line in result.rstrip('\n').split('\n'):
                successes.append('{};{}'.format(';'.join(target), line))

        if self.to_stdout:
            print('\n'.join(successes))
        else:
            self._save_lines(self.get_successes_name(), successes)

    def update_failures(self, new_failures):
        if new_failures:
            failures = self._load_lines(self.get_failures_name())
            if not failures:
                failures.append('#hostname;#ip4;#reason')

            for target, result in sorted(new_failures.items()):
                for line in result.rstrip('\n').split('\n'):
                    failures.append('{};{}'.format(';'.join(target), line))

            self._save_lines(self.get_failures_name(), failures)

    @classmethod
    def _load_targets(cls, filename):
        """
        Load foobar.rsh.targets file, containing CSV "#hostname;#ip4"
        """
        targets = cls._lines_to_targets(cls._load_lines(filename))
        if not targets:
            raise ValueError(f'expected at least one target in {filename}')
        return targets

    @staticmethod
    def _load_script(filename):
        """
        Load foobar.rsh file, containing "#!/bin/sh .. echo foo; echo bar"

        As an exception, instead of "#!/bin/sh" the shebang may also point to
        "rshall" so the script can be run directly.
        """
        try:
            with open(filename) as fp:
                script = fp.read()
                first_line = script.split('\n', 1)[0].strip()
                if '\n' in script and first_line in (
                        '#!/usr/bin/env rshall', '#!/usr/bin/rshall'):
                    # Replace their "rshall" shebang with regular /bin/sh.
                    script = '#!/bin/sh' + script[script.index('\n'):]
                elif first_line != '#!/bin/sh':
                    # This is the expected syntax.
                    raise NotImplementedError('we only support sh-syntax now')
        except FileNotFoundError:
            print(
                f'(no {filename} script found, using example/default script)',
                file=sys.stderr)
            return None
        else:
            return script

    @staticmethod
    def _load_lines(filename):
        """
        Load foobar.rsh.out OR foobar.rsh.err file, with output lines
        """
        lines = []
        try:
            with open(filename) as fp:
                data = fp.read().rstrip()
                if data:
                    lines = data.split('\n')
        except FileNotFoundError:
            pass
        return lines

    @staticmethod
    def _save_lines(filename, lines):
        """
        Write foobar.rsh.out OR foobar.rsh.err file, with output lines
        """
        with open(f'{filename}.new', 'w') as fp:
            fp.write('\n'.join(lines) + '\n')
        rename(f'{filename}.new', filename)

    @staticmethod
    def _lines_to_targets(lines):
        """
        Take the first two ';'-delimited arguments and create Targets

        The file might look like:

            the.awesome.host;1.2.3.4;whatever

        This will return [Target(hostname='the.awesome.host', ip='1.2.3.4')]
        """
        targets = []

        for line in lines:
            line = line.strip()
            if line.startswith('#'):
                pass
            elif line and ';' in line:
                hostname, ip = line.split(';', 2)[0:2]
                try:
                    ip = str(IPv4Address(ip))
                except AddressValueError:
                    raise ValueError(f'unparsable IP in line: {line}')
                target = Target(hostname=hostname, ip=ip)
                targets.append(target)
            else:
                assert not line, (
                    f'expected "<hostname>;<ip>[;..]", got {line!r}')

        return targets

    @staticmethod
    def _unlink_or_ignore(filename):
        try:
            unlink(filename)
        except OSError as e:
            if e.args[0] != 2:  # ENOFILE (FileNotFoundError)
                raise


def main():
    parser = ArgumentParser(
        prog='rshall',
        description='''
            Run script on many hosts at once and collect output.
            Used for one-off scripts that you don't want to build any
            ansible/chef/puppet scripts for.
        ''')
    parser.add_argument(
        '-c', '--stdout', action='store_true',
        help='Use stdout instead of the .out file.')
    parser.add_argument(
        '-f', '--force', action='store_true',
        help='Remove existing out/err files and then run.')
    parser.add_argument(
        '-t', '--targets', metavar='targets_name',
        help='Use specified targets file.')
    parser.add_argument(
        'script_name', help='The rsh script to process.')
    parser.add_argument(
        'project_name', nargs='?', help=(
            'Alternative file to operate on: project_name.{targets,out,err}.'))
    args = parser.parse_args()

    project = Project(
        args.script_name, args.project_name, args.targets, args=args)
    runner = Runner()

    runner.add_targets(project.targets)
    if project.script is not None:
        runner.set_script(project.script)

    try:
        runner.run()
    finally:
        project.update_successes(runner.successes)
        project.update_failures(runner.failures)

    total_successes = (len(runner.successes) + project.successes_skip)
    total_failures = (len(runner.failures) + project.failures_skip)

    if not args.stdout:
        print(
            f'{len(runner.successes)} success, {len(runner.failures)} fails, '
            f'{project.successes_skip} earlier success, '
            f'{project.failures_skip} earlier fails',
            file=sys.stderr)
        if total_successes:
            print(
                (f'See {project.get_successes_name()} for '
                 f'{total_successes} outputs'), file=sys.stderr)

    if total_failures:
        print(
            (f'See {project.get_failures_name()} for '
             f'{total_failures} errors'), file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()