Port to fabric2; add example config

LICENCE

@@ -0,0 +1 @@
+No Nazis and GPLv2.

README.txt

@@ -0,0 +1,7 @@
+ __  __       _ _   _ _           _ _ _ _ 
+|  \/  |_   _| | |_(_) |__   __ _| | | | |
+| |\/| | | | | | __| | '_ \ / _` | | | | |
+| |  | | |_| | | |_| | |_) | (_| | | |_|_|
+|_|  |_|\__,_|_|\__|_|_.__/ \__,_|_|_(_|_)
+
+Run commands across a set of hosts interactively!

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
+

multiball/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.0"

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

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)

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

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"
+