diff --git a/LICENCE b/LICENCE new file mode 100644 index 0000000..4f037d9 --- /dev/null +++ b/LICENCE @@ -0,0 +1 @@ +No Nazis and GPLv2. diff --git a/README.txt b/README.txt new file mode 100644 index 0000000..0e5abea --- /dev/null +++ b/README.txt @@ -0,0 +1,7 @@ + __ __ _ _ _ _ _ _ _ _ +| \/ |_ _| | |_(_) |__ __ _| | | | | +| |\/| | | | | | __| | '_ \ / _` | | | | | +| | | | |_| | | |_| | |_) | (_| | | |_|_| +|_| |_|\__,_|_|\__|_|_.__/ \__,_|_|_(_|_) + +Run commands across a set of hosts interactively! diff --git a/multiball.cfg.example b/multiball.cfg.example new file mode 100644 index 0000000..66c5797 --- /dev/null +++ b/multiball.cfg.example @@ -0,0 +1,44 @@ +[config] +history=multiball.history +sshconfigs=~/.ssh/config +hostlists= + +[startup] + +[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/__init__.py b/multiball/__init__.py new file mode 100644 index 0000000..6c8e6b9 --- /dev/null +++ b/multiball/__init__.py @@ -0,0 +1 @@ +__version__ = "0.0.0" diff --git a/multiball/__main__.py b/multiball/__main__.py new file mode 100644 index 0000000..382d31e --- /dev/null +++ b/multiball/__main__.py @@ -0,0 +1,416 @@ +# +# Multiball - Run things on a buncha servers interactively. +# + +# +# TODO +# - 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 yaml +import prompt_toolkit +from prompt_toolkit.completion import WordCompleter +from prompt_toolkit.history import InMemoryHistory, FileHistory +from prompt_toolkit.shortcuts import PromptSession + +from .fabtools import HostSet +from .tools import load_sshconfig_server_list, load_bare_server_list, load_config + + +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.") +} + +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, + } + configdict = load_config(Path("multiball.cfg")) + self.allhosts: set = set() + self.history: Union[InMemoryHistory, FileHistory] = InMemoryHistory() + self.environment = {} + self.safety = True + self.ssh_config = None + + if config := configdict.get('config', {}): + if 'history' in config: + self.history = FileHistory(config['history']) + + if sshconfigs := config.get('sshconfigs', ''): + for sc in sshconfigs.split(','): + if self.ssh_config is None: + self.ssh_config = Path(sc).expanduser() + 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: + self.ssh_config = Path("~/.ssh/ssh_config").expanduser() + print( + f"Warning: no [config] section in {configFile}. Using fallbacks.") + self.allhosts = load_sshconfig_server_list(self.ssh_config) + + # 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): + remotecmd = [] + for k, v in self.environment.items(): + remotecmd.append(f"export {k}='{v}'") + remotecmd.append(command) + h = HostSet(self.targethosts, self.ssh_config) + h.run(';'.join(remotecmd)) + + 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() diff --git a/multiball/fabtools.py b/multiball/fabtools.py new file mode 100644 index 0000000..307f2ce --- /dev/null +++ b/multiball/fabtools.py @@ -0,0 +1,67 @@ +from pathlib import Path +from threading import Thread, Lock + +from tqdm import tqdm + +from fabric2 import ThreadingGroup, SerialGroup, Config, Connection + +def thread_run(connection, command, result_lock, result_queue): + res = connection.run(command, warn=True, hide=True) + with result_lock: + result_queue.append((connection, res)) + +# A set of hosts we can target with a series of commands and also can collect output of each command +class HostSet: + def __init__(self, hostlist: list, ssh_config_path: Path=None): + if ssh_config_path is None: + ssh_config_path = Path("~/.ssh/config") + ssh_config_path = ssh_config_path.expanduser() + config = Config({"ssh_config_path": str(ssh_config_path)}) + self.connections = [] + for host in hostlist: + self.connections.append(Connection(host, config=config)) + + def run(self, command: str): + resq = [] + reslock = Lock() + threads = [] + prog = tqdm(total=len(self.connections), unit="hosts") + for connection in self.connections: + t = Thread(target=thread_run, args=[connection, command, reslock, resq]) + t.start() + threads.append(t) + + # display status about threads, and join them when they finish + while (True): + nt = [] + for i in threads: + if not i.is_alive(): + i.join() + prog.update() + #print('.', end='', flush=True) + else: + nt.append(i) + threads = nt + if len(threads) == 0: + break + prog.close() + + # Gather up results by output + gathered = {} + for connection, result in resq: + rstr = str(result) + if not rstr in gathered: + gathered[rstr] = [] + gathered[rstr].append(connection) + + # display results + for result, connections in gathered.items(): + print('-----> [{}]'.format(' '.join(connection.original_host for connection in connections))) + print(result) + # ## gather_output + # for connection in self.connections: + # # import pdb + # # pdb.set_trace() + # res = connection.run(command) + # print('---- {} ----'.format(connection.original_host)) + # print(res) diff --git a/multiball/tools.py b/multiball/tools.py new file mode 100644 index 0000000..fb16bce --- /dev/null +++ b/multiball/tools.py @@ -0,0 +1,91 @@ +import re + +import os.path + +from pathlib import Path +from typing import Optional + + +# 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 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5431e3c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,22 @@ +[build-system] +requires = ["pdm-backend"] +build-backend = "pdm.backend" + + +[project] +name = "multiball" +dynamic = ["version"] +description = "Run commands accross a set of hosts interactively." +authors = [{name = "Cassowary", email="cassowary@riseup.net"}] +dependencies = ["tqdm", "yaml-1.3", "prompt_toolkit", "fabric2"] +requires-python = ">=3.8" +readme = {file = "README.txt", content-type = "text/plain"} +license = {text = "LICENSE"} + +[tool.pdm.version] +source = "file" +path = "multiball/__init__.py" + +[project.scripts] +multiball = "multiball.__main__:main" + \ No newline at end of file