Initial checkin for fixing up multiball

config.yml

@@ -0,0 +1,15 @@
+# GSCK cumin.yml
+# Basically the default config to make cumin work.
+
+# Cumin main configuration
+# #
+# # By default Cumin load the configuration from /etc/cumin/config.yaml, but it can be overriden by command line argument
+# #
+transport: clustershell  # Default transport to use, can be overriden by command line argument
+log_file: ~/.cumin/cumin.log  # Absolute or relative path for the log file, expands ~ into the user's home directory
+# # If set, use this backend to parse the query first and only if it fails, fallback to parse it with the general
+# # multi-query grammar [optional]
+default_backend: direct
+
+# # Environment variables that will be defined [optional]
+# environment:

multiball.cfg

@@ -0,0 +1,47 @@
+[config]
+history=multiball.history
+sshconfigs=~/.ssh/autonomic_config
+hostlists=
+
+[startup]
+# Fixme we gotta figure these out, but maybe we do it with grouping instead.
+#-filteralias iww *iww*
+#-filteralias autonomic *autonomic*
+
+[cmd.upgrade]
+-unsafe
+-targets
+-confirm Targeting these, ok?
+-env DEBIAN_FRONTEND=noninteractive
+sudo -E apt -y upgrade
+-confirm Do they need to be rebooted?
+sudo reboot
+
+
+[cmd.install-pkg]
+-arguments
+-unsafe
+-targets
+-confirm Installing `{@}` on these, ok?
+-env DEBIAN_FRONTEND=noninteractive
+sudo -E apt -y install {@}
+-confirm Do they need to be rebooted?
+sudo reboot
+
+[cmd.upgrade-pkg]
+-arguments
+-unsafe
+-targets
+-confirm Installing `{@}` on these, ok?
+-env DEBIAN_FRONTEND=noninteractive
+sudo -E apt -y install --only-upgrade {@}
+
+[cmd.last-updated]
+-safe sudo zgrep -B1 'apt-get upgrade' /var/log/apt/* | grep 'Start' | cut -d'z' -f2 | sort | tail -n1 | cut -d' ' -f2 || true
+
+[cmd.upgradable]
+-safe apt list --upgradable 2>&1 | grep upgradable | grep -v WARNING || true
+
+[cmd.ping]
+-safe echo ping
+

multiball.py

@@ -0,0 +1,511 @@
+#
+# Multiball - Run things on a buncha servers interactively.
+#
+
+#
+# TODO
+# - switch to Fabric so we can get away from cumin's bullshit
+# - keep track of previous command outputs for last command, and any previous command with target list and command line
+# - calling commands and alieses from command line
+# - tagging and filtering commands (so #noupgrade gets a tag, and we can filter them)
+# - saving the last host set in a .file
+# - allow server groups
+# - allow grepping of output from a command, and then populating the target list from matchintg on it
+# - add 'watch' to run a command repeatedly until it succeeds (interruptable)
+# - add the number of servers for each group of services too
+# - notice when you don't sudo in front of command and it errors and ask you if you meant to sudo
+# - autocomplete for environment variable names (requires heirarchical completer)
+# - implement target aliases (-filteralias) which gives a label to an argument to -hosts
+# - ad-hoc host groupings with assigned names, and a host grouping stack
+# - Assign each output group to a number that can be easily selected with a -target etc. But also allow matching
+#   against those groups. Output groupings survive until the next command.
+# - allow scripts that use -safe to prompt for safety / restore safety after running
+
+
+
+
+import argparse
+import datetime
+import fnmatch
+import logging
+import os
+import pprint
+import re
+import shutil
+import shlex
+import sys
+import traceback
+from itertools import islice, zip_longest
+from pathlib import Path
+from typing import Iterable, Optional, Sequence, Union
+
+import cumin
+import yaml
+from cumin import query, transport, transports
+import prompt_toolkit
+from prompt_toolkit.completion import WordCompleter
+from prompt_toolkit.history import InMemoryHistory, FileHistory
+from prompt_toolkit.shortcuts import PromptSession
+
+
+class MultiballCommandException(BaseException):
+    pass
+
+
+class MBCommandAbort(MultiballCommandException):
+    pass
+
+
+class MBCommandExit(MultiballCommandException):
+    pass
+
+
+class MBCommandRun(MultiballCommandException):
+    pass
+
+
+class MBExit(MultiballCommandException):
+    pass
+
+
+# Dictionary of top-level commands and help strings (and eventually we'll note what kind of argumest so we can construct
+# heirarchical completion from them)
+COMMANDS = {
+    'help': (("-help", "help", "?"), "This help"),
+    'targethosts': (('-targethosts', '-targethosts', '-targets', '-hosts'), "Show the current list of target hosts."),
+    'clear': (('-clear',), "Clear target host list."),
+    'host': (('-host', '-add', '-target'), "Set the target host list. -host clears the list first. Supports wildcards."),
+    'remove': (('-remove', '-removehost', '-removetarget'), "Remove hosts from target host list. Supports wildcards."),
+    'all': (('-all',), "Reset the target host list to all known hosts."),
+    'allhosts': (('-allhosts', '-alltargets'), "Show list of all known hosts."),
+    'exit': (('-exit', 'exit'), "Exit."),
+    'environment': (('-environment', '-env', '-set'), "Set (or print) environment to send to remote host before running commands."),
+    'clearenv': (('-clearenv',), "Clear environment (entire or single variable)"),
+    'safety': (('-safety', 'safety'), "Turn on safety or turn off with `off`."),
+    'safe': (('-safe'), "Run this command as if it were safe even if safety is on."),
+    'confirm': (('-confirm',), "[scripting] Prompt for confirmation"),
+    'echo': (('-echo', ), "[scripting] print string"),
+    'arguments': (('-arguments',), "[scripting] Abort if no arguments are specified."),
+    # fixme add interpolation
+    'unsafe': (('-unsafe'), "[scriptiong] Abort command if safety is on."),
+    # 'history': (('-history', '-h'), "Show the last 10 commands in the history."),
+    # 'filteralias': (('-filteralias',), "Create a filter alias which allows quick filtering.")
+}
+
+
+# load utility functions
+def load_sshconfig_server_list(sshconfig: Path,
+                               hostfilter: Optional[re.Pattern] = None,
+                               lineSkip: Optional[re.Pattern] = None) -> set:
+    """
+    Load a list of servers from an ssh configuration file.
+    """
+
+    output = set()
+    with sshconfig.open() as inconfig:
+        for line in [l.strip() for l in inconfig]:
+            if line.startswith("Host "):
+                if lineSkip and lineSkip.match(line):
+                    continue
+                hostinfo = line.split('#', 2)[0]
+                hostname = hostinfo.split()[-1]
+                if hostfilter and not hostfilter.match(hostname):
+                    continue
+                output.add(hostname)
+    return output
+
+
+def load_bare_server_list(serverlist: Path,
+                          hostfilter: Optional[re.Pattern] = None,
+                          lineSkip: Optional[re.Pattern] = None) -> set:
+    """
+    Load a list of servers from a simple text file.
+    """
+    output = set()
+    with sshconfig.open() as inconfig:
+        for line in [l.strip() for l in inconfig]:
+            if not line:
+                continue
+            if lineSkip and lineSkip.match(line):
+                continue
+
+            hostinfo = line.split('#', 2)[0]
+            hostname = hostinfo.split()[-1]
+            if hostfilter and not hostfilter.match(hostname):
+                continue
+            output.add(hostname)
+    return output
+
+
+def load_config(config_path: Path):
+    output = dict()
+    with config_path.open() as inconfig:
+        section = 'global'
+        cmd = False
+        lineno = 0
+        for rline in inconfig:
+            lineno += 1
+
+            line = rline.strip()
+
+            if not line:
+                continue
+            if line.startswith('#'):
+                continue
+
+            if line.startswith('['):
+                section = line[1:-1]
+                if section.startswith("cmd") or section == 'startup':
+                    cmd = True
+                    output.setdefault(section, list())
+                else:
+                    cmd = False
+                    output.setdefault(section, dict())
+                continue
+
+            if cmd:
+                output[section].append(line)
+            else:
+                if not '=' in line:
+                    print(f"config parse error, line {lineno}: {line}")
+                    continue
+                kv = line.split('=', 2)
+                if len(kv) == 1:
+                    value = ''
+                else:
+                    value = kv[1]
+                output[section][kv[0]] = value
+        return output
+
+
+class Multiball:
+    def __init__(self, configFile: Optional[Path] = None):
+        # load configuration
+        self.handlers = {
+            'environment': self.command_environment,
+            'confirm': self.command_confirm,
+            'arguments': self.command_arguments,
+            'unsafe': self.command_unsafe,
+            'echo': self.command_echo,
+            'safety': self.command_safety,
+            'safe': self.command_safe,
+            'clearenv': self.command_clearenv,
+        }
+        self.cuminconfig = cumin.Config("config.yml")
+        configdict = load_config(Path("multiball.cfg"))
+        self.allhosts: set = set()
+        self.history: Union[InMemoryHistory, FileHistory] = InMemoryHistory()
+        self.environment = {}
+        self.safety = True
+
+        if config := configdict.get('config', {}):
+            if 'history' in config:
+                self.history = FileHistory(config['history'])
+
+            if sshconfigs := config.get('sshconfigs', ''):
+                for sc in sshconfigs.split(','):
+                    newhosts = load_sshconfig_server_list(
+                        Path(sc).expanduser())
+                    print(f"Loaded {len(newhosts)} from {sc}")
+                    self.allhosts.update(newhosts)
+
+            if bareconfigs := config.get('hostlists'):
+                for sc in bareconfigs.split(','):
+                    newhosts = load_bare_server_list(Path(sc).expanduser())
+                    print(f"Loaded {len(newhosts)} from {sc}")
+                    self.allhosts.update(newhosts)
+
+        else:
+            print(
+                f"Warning: no [config] section in {configFile}. Using fallbacks.")
+            self.allhosts = load_sshconfig_server_list(
+                Path("~/.ssh/ssh_config").expanduser())
+
+        # setup state and environment
+        self.targethosts = set(self.allhosts)
+        self._is_scripting = False
+        self._script_frame_arguments: Sequence[str] = []
+        self.session: PromptSession = PromptSession(history=self.history)
+        self.aliases = {}
+        # make aliases
+        for k in configdict:
+            if k.startswith('cmd.'):
+                # print(f"command alias: {k}")
+                alias = k.split('.', 1)[1]
+                self.aliases[alias] = list(configdict[k])
+
+        # call startup commands
+        if startup := configdict.get('startup', []):
+            # fixme execute startup and handle exceptions correctly
+            for cmd in startup:
+                self._run_command(cmd)
+
+
+    def command_safe(self, command, args):
+        try:
+            self.safety = False
+            self.do_cmd(args)
+        finally:
+            self.safety = True
+
+    def command_safety(self, command, args):
+        """
+        Toggle the safety.
+        """
+        if args.lower().startswith("off"):
+            self.safety = False
+            print("** Safety OFF. Now unsafe.")
+        else:
+            self.safety = True
+            print("** Safety ON.")
+
+    def command_echo(self, command, args):
+        """
+        Echo string
+        fixme put string interpolation here so that state variables can be printed
+        """
+        print(args)
+
+    def command_environment(self, command, args):
+        """
+        Set an environment variable in the environment object, or print out the environment vithout any arguments.
+        """
+        if (args):
+            if '=' in args:
+                cmd = [x.strip() for x in args.split('=', 1)]
+                key = cmd[0]
+                if (len(cmd) > 1) and (cmd[1]):
+                    self.environment[key] = cmd[1]
+                else:
+                    if cmd[0] in self.environment:
+                        del self.environment[key]
+            else:
+                for a in args.split():
+                    if a in self.environment:
+                        print(f"{a}={self.environment[a]}")
+                    else:
+                        print(f"{a} not in environment.")
+        else:
+            for k, v in self.environment.items():
+                print(f"{k}={v}")
+
+    def command_confirm(self, command, args):
+        """
+        Prompt user with a yes/no question. Default responses result in a CommandAbort exception.
+        """
+        result = prompt_toolkit.prompt(
+            "<" + args + " [Yys/Nn] > ")
+        if result and (result[0] in 'Yys'):
+            return True
+        else:
+            print("(no)")
+            raise MBCommandAbort
+
+    def command_arguments(self, command, args):
+        """
+        Rasiess CommandAbort if the scriptenv doesn't contain arguments.
+
+        FIXME make a command parser in here that constrains the scriptenv arguments to a specific set
+        or format, and populates the environment with the variables so they can be placed in other commands.
+        """
+        if not self._is_scripting or not self._script_frame_arguments:
+            if not self._is_scripting:
+                print("does nothing outside of script")
+            else:
+                print("arguments required")
+            raise MBCommandAbort
+
+        return
+
+    def command_unsafe(self, command, args):
+        """
+        Raises CommandAbort if safty is on.
+        """
+        if not self._is_scripting or self.safety:
+            if not self._is_scripting:
+                print("does nothing outside of script")
+            else:
+                print("cannot run in safety mode")
+            raise MBCommandAbort
+
+        return
+
+    def command_clearenv(self, command, args):
+        print("cleared environment")
+        self.environment = dict()
+
+    def _print_targetlist(self, tlist=None, keyword='Targeting'):
+        if not tlist:
+            tlist = self.targethosts
+        print(f"{keyword} {len(self.targethosts)} hosts:")
+        print(' '.join(tlist))
+
+    def _print_help(self):
+        for item in COMMANDS.values():
+            print(', '.join(item[0]) + ' - ' + item[1])
+        print("Other entries run command on targets.")
+
+    def _run_command(self, command: str):
+        if not command or command.split()[0] in COMMANDS['targethosts'][0]:
+            self._print_targetlist()
+            return
+
+        # fixme we might be able to use pattern matching to look up the command abstract and make this whole section
+        # more readable
+        commands = command.split(maxsplit=1)
+        abst = 'REMOTECOMMAND'  # default is run a remote command
+        for a, c in COMMANDS.items():
+            if commands[0] in c[0]:
+                abst = a
+        command = commands[0]
+        if (len(commands) > 1):
+            cargs = ' '.join(commands[1:])
+        else:
+            cargs = ''
+
+        if abst == 'exit':
+            raise MBExit
+        if abst == 'help':
+            self._print_help()
+        elif abst in self.handlers:
+            self.handlers[abst](command[0], cargs)
+        elif abst == 'clear':
+            print("cleared targets")
+            self.targethosts = set()
+        elif abst == 'host':
+            if len(commands) > 1:
+                targethostlist = set()
+                # iterate potential targethosts and match with globber
+                for phost in cargs.split():
+                    if ('?' in phost) or ('*' in phost) or ('[' in phost):
+                        # glob
+                        for h in self.allhosts:
+                            if fnmatch.fnmatch(h, phost):
+                                targethostlist.add(h)
+                    else:
+                        targethostlist.add(phost)
+                if command in ('-host', '-target'):
+                    print("note: clearing previous selection")
+                    self.targethosts = targethostlist
+                else:
+                    self.targethosts.update(targethostlist)
+                self._print_targetlist()
+        elif abst == 'remove':
+            if len(commands) > 1:
+                targethostlist = set()
+                # iterate potential targethosts and match with globber
+                for phost in cargs.split():
+                    if ('?' in phost) or ('*' in phost) or ('[' in phost):
+                        # glob
+                        for h in self.allhosts:
+                            if fnmatch.fnmatch(h, phost):
+                                targethostlist.add(h)
+                    else:
+                        targethostlist.add(phost)
+                self.targethosts = self.targethosts - targethostlist
+                self._print_targetlist()
+        elif abst == 'all':
+            self.targethosts = self.allhosts
+            self._print_targetlist()
+        elif abst == 'allhosts':
+            self._print_targetlist(self.allhosts, "all known")
+        elif command in self.aliases:
+            print(f"running alias `{command}`")
+            self._run_alias(command, cargs)
+        else:
+            raise MBCommandRun
+
+    def _run_alias(self, command, cargs):
+        self._is_scripting = True
+        self._script_frame_arguments = (cargs, shlex.split(cargs))
+        try:
+            for cline in self.aliases[command]:
+                self._is_scripting = True
+                # fixme more extensive variable filling
+                if '{@}' in cline:
+                    cline = cline.replace('{@}', cargs.strip())
+                try:
+                    self.do_cmd(cline)
+                except MBCommandAbort:
+                    # end the alias on abort
+                    return
+        finally:
+            self._is_scripting = False
+            self._script_frame_arguments = None
+
+    def _run_remote_command(self, command):
+        target = transports.Target(list(self.targethosts))
+        worker = transport.Transport.new(self.cuminconfig, target)
+        worker.reporter = cumin.transports.clustershell.NullReporter
+        remotecmd = []
+        for k, v in self.environment.items():
+            remotecmd.append(f"export {k}='{v}'")
+        remotecmd.append(command)
+        worker.commands = [';'.join(remotecmd)]
+        worker.handler = 'async'
+        exit_code = worker.execute()  # Execute the command on all hosts in parallel
+        for nodes, output in worker.get_results():  # Cycle over the results
+            nstr = ' '.join(nodes)
+            print("Results from:")
+            print(nstr)
+            print('-----')
+            print(output.message().decode())
+            print('-----')
+
+    def do_cmd(self, command):
+        try:
+            self._run_command(command)
+        except MBCommandRun:
+            # remote command!
+            if self.safety:
+                print(f"would run `{command}` on targethosts")
+                print("Safety is ON. `-safety off` to turn off.")
+            else:
+                print(f"running `{command}` on targethosts")
+                self._run_remote_command(command)
+
+    def cmd_prompt(self):
+        self._is_scripting = False
+        allcommands = []
+        for item in COMMANDS.values():
+            allcommands.extend(item[0])
+        allcommands.extend(list(self.aliases.keys()))
+        completer = WordCompleter(list(self.allhosts) + allcommands)
+        if self.safety:
+            prompt = "!"
+        else:
+            prompt = '$'
+        command = self.session.prompt(
+            f'{len(self.targethosts)} -> {prompt} ', completer=completer)
+        try:
+            self.do_cmd(command)
+        except MBCommandAbort:
+            # Nothing of importance happens here.
+            ...
+
+
+def multiball(*args, **kwargs):
+    print("Welcome to Multiball. type -help for help.")
+    mb = Multiball()
+    # loop on prompt
+    while (True):
+        try:
+            mb.cmd_prompt()
+        except MBExit:
+            return
+
+
+def main():
+    try:
+        sys.exit(multiball(sys.argv))
+    except EOFError:
+        sys.exit(0)
+    except Exception as inst:
+        # fixme log exception
+        print(f"Exception! {inst}")
+        print(traceback.format_exc())
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    main()