multiball/multiball/__main__.py

#
# 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
# - Add variables that can be set and passed to commands
# - implement various commented commands in the command list
# - implement interactive alias system


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."),
    # 'undo': (('/undo',), "Undo last host or environment command."),
    # 'alias': (('/alias', ), "Create an alias interactively.")
    ## Fixme unambiguate the above commands which list things into this command so things are more consistent
    # 'list': (('/list',), "List various things (environmet, allhosts, hosts, targets)")
    # '!': (('!', '/local'), "Run a local command and capture its 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,
        }
        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()