From af1c7d61396ee72f9b986472c441470ce9af08a6 Mon Sep 17 00:00:00 2001 From: Cassowary Date: Tue, 26 Mar 2024 09:25:37 -0700 Subject: [PATCH] Initial checkin for fixing up multiball --- config.yml | 15 ++ multiball.cfg | 47 +++++ multiball.py | 511 ++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 573 insertions(+) create mode 100644 config.yml create mode 100644 multiball.cfg create mode 100644 multiball.py diff --git a/config.yml b/config.yml new file mode 100644 index 0000000..541cc3f --- /dev/null +++ b/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: diff --git a/multiball.cfg b/multiball.cfg new file mode 100644 index 0000000..86ab203 --- /dev/null +++ b/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 + diff --git a/multiball.py b/multiball.py new file mode 100644 index 0000000..234db48 --- /dev/null +++ b/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()