# # 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()