From 614ea8b964707996f27ebe3862a3bd6b6b05f439 Mon Sep 17 00:00:00 2001 From: Cassowary Date: Sat, 4 Oct 2025 17:41:01 -0700 Subject: [PATCH] Colorizing Start! - Add styles and colors to most output. - Update TODO list - Add ability to interrupt "waiting" condition with Control-C --- TODO.txt | 4 +++ multiball/__init__.py | 2 +- multiball/__main__.py | 66 +++++++++++++++++++++++++------------------ multiball/fabtools.py | 60 +++++++++++++++++++++++++-------------- multiball/style.py | 26 +++++++++++++++++ 5 files changed, 108 insertions(+), 50 deletions(-) create mode 100644 multiball/style.py diff --git a/TODO.txt b/TODO.txt index 8051fed..66d7aac 100644 --- a/TODO.txt +++ b/TODO.txt @@ -1,4 +1,8 @@ TODO +- fix handling of control-c +- redo the command loop to make some fuckin sense why do we use eceptions to do modal stuff +- make styles optional +- expose styles in config file or personal config - add local variable support that can be interpolated into commands - keep track of previous command outputs for last command, and any previous command with target list and command line - Allow any given command to be tagged with a label, and the outputs would be stored in the output dictionary under that name (for scripting); diff --git a/multiball/__init__.py b/multiball/__init__.py index 3d18726..906d362 100644 --- a/multiball/__init__.py +++ b/multiball/__init__.py @@ -1 +1 @@ -__version__ = "0.5.0" +__version__ = "0.6.0" diff --git a/multiball/__main__.py b/multiball/__main__.py index 2198318..7cd6279 100644 --- a/multiball/__main__.py +++ b/multiball/__main__.py @@ -11,24 +11,27 @@ import logging import os import pprint import re -import shutil import shlex +import shutil 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 +import yaml from prompt_toolkit.completion import WordCompleter -from prompt_toolkit.history import InMemoryHistory, FileHistory +from prompt_toolkit.history import FileHistory, InMemoryHistory from prompt_toolkit.shortcuts import PromptSession +from prompt_toolkit.styles import Style from . import __version__ from .fabtools import HostSet -from .tools import load_sshconfig_server_list, load_bare_server_list, load_config +from .tools import (load_bare_server_list, load_config, + load_sshconfig_server_list) +from .style import print, STYLES class MultiballCommandException(BaseException): pass @@ -113,19 +116,19 @@ class Multiball: self.ssh_config = Path(sc).expanduser() newhosts = load_sshconfig_server_list( Path(sc).expanduser()) - print(f"Loaded {len(newhosts)} from {sc}") + 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}") + 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.") + f"Warning: no [config] section in {configFile}. Using fallbacks.") self.allhosts = load_sshconfig_server_list(self.ssh_config) # setup state and environment @@ -161,17 +164,17 @@ class Multiball: """ if args.lower().startswith("off"): self.safety = False - print("** Safety OFF. Now unsafe.") + print("** Safety OFF. Now unsafe.") else: self.safety = True - print("** Safety ON.") + 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) + print(f"{args}") def command_environment(self, command, args): """ @@ -200,12 +203,14 @@ class Multiball: """ Prompt user with a yes/no question. Default responses result in a CommandAbort exception. """ - result = prompt_toolkit.prompt( - "<" + args + " [Yys/Nn] > ") + result = prompt_toolkit.prompt([("class:promptyn", + "<" + args + " [Yys/Nn]>"), + ("class:message", + " ")], style=STYLES) if result and (result[0] in 'Yys'): return True else: - print("(no)") + print("(no)") raise MBCommandAbort def command_arguments(self, command, args): @@ -217,9 +222,9 @@ class Multiball: """ if not self._is_scripting or not self._script_frame_arguments: if not self._is_scripting: - print("does nothing outside of script") + print("does nothing outside of script") else: - print("arguments required") + print("arguments required") raise MBCommandAbort return @@ -230,9 +235,9 @@ class Multiball: """ if not self._is_scripting or self.safety: if not self._is_scripting: - print("does nothing outside of script") + print("does nothing outside of script") else: - print("cannot run in safety mode") + print("cannot run in safety mode") raise MBCommandAbort return @@ -245,7 +250,7 @@ class Multiball: # FIXME: We need to make aliases store multiple entries, and save should take the whole alias not just # the last command. if (not self.last_run): - print("Nothing to log.") + print("Nothing to log.") return tstamp = datetime.datetime.now().strftime("%Y-%m-%d-%H%M%Z") outfilename = f"{tstamp}.multiball.log" @@ -254,7 +259,7 @@ class Multiball: if (not outfilename.endswith(".log")): outfilename = outfilename + ".log" - print(f"Writing {outfilename} for command `{self.last_run['cmd']}`.") + print(f"Writing {outfilename} for command `{self.last_run['cmd']}`.") with (open(outfilename, 'w')) as out: out.writelines([f"command: {self.last_run['cmd']}\n", "-------------------------------------------\n" @@ -265,12 +270,13 @@ class Multiball: def _print_targetlist(self, tlist=None, keyword='Targeting'): if not tlist: tlist = self.targethosts - print(f"{keyword} {len(self.targethosts)} hosts:") + 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]) + m = ', '.join(item[0]) + ' - ' + item[1] + print(f"{m}") print("Other entries run command on targets.") def _run_command(self, command: str): @@ -382,10 +388,10 @@ class Multiball: except MBCommandRun: # remote command! if self.safety: - print(f"would run `{command}` on targethosts") - print("Safety is ON. `/safety off` to turn off.") + print(f"would run `{command}` on targethosts") + print("Safety is ON. `/safety off` to turn off.") else: - print(f"running `{command}` on targethosts") + print(f"running `{command}` on targethosts") self._run_remote_command(command) def cmd_prompt(self): @@ -397,10 +403,14 @@ class Multiball: completer = WordCompleter(list(self.allhosts) + allcommands) if self.safety: prompt = "!" + pclass = "class:safe" else: prompt = '$' - command = self.session.prompt( - f'{len(self.targethosts)} -> {prompt} ', completer=completer) + pclass = "class:unsafe" + command = self.session.prompt([("class:count", f'{len(self.targethosts)}'), + ("class:prompt", '->'), + (pclass, prompt), + ("class:default", " ")], completer=completer, style=STYLES) try: self.do_cmd(command) except MBCommandAbort: @@ -409,7 +419,7 @@ class Multiball: def multiball(*args, **kwargs): - print(f"Welcome to Multiball {__version__}. type /help for help.") + print(f"Welcome to Multiball {__version__}. type /help for help.") mb = Multiball() # loop on prompt while (True): @@ -426,7 +436,7 @@ def main(): sys.exit(0) except Exception as inst: # fixme log exception - print(f"Exception! {inst}") + print(f"Exception! {inst} 10.0): - tupdate = time.time() - print("Still waiting on ", [thread.host for thread in threads]) - # FIXME have the thread killed if it's still waiting after $TIMEOUT + if ((time.time() - tupdate) > 10.0): + tupdate = time.time() + m = "Still waiting on " + str([thread.host for thread in threads]) + print_style(m, style="alert") + # FIXME have the thread killed if it's still waiting after $TIMEOUT + except KeyboardInterrupt as inst: + # cancel remaining servers + print_style("(break)", style="error") + for thread in threads: + thread.connection.close() + prog.close() # Gather up results by output @@ -109,10 +127,10 @@ class HostSet: self._lprint('-----> [{}]'.format(' '.join(connection.original_host for connection in connections))) self._lprint(result) - self._lprint('#####> RUN SUMMARY <#####') + self._lprint('#####> RUN SUMMARY <#####', style="alert") if (success): - self._lprint(' Succeeded: ', ', '.join(connection.original_host for connection in success)) + self._lprint(' Succeeded: ', ', '.join(connection.original_host for connection in success), style="good") if (commandfail): - self._lprint(' Command fail: ', ', '.join(connection.original_host for connection in commandfail)) + self._lprint(' Command fail: ', ', '.join(connection.original_host for connection in commandfail), style="warning") if (error): - self._lprint('Failed with Error: ', ','.join(connection.original_host for connection in error)) + self._lprint('Failed with Error: ', ','.join(connection.original_host for connection in error), style="error") diff --git a/multiball/style.py b/multiball/style.py new file mode 100644 index 0000000..927951c --- /dev/null +++ b/multiball/style.py @@ -0,0 +1,26 @@ +from prompt_toolkit import print_formatted_text, HTML +from prompt_toolkit.styles import Style +from prompt_toolkit.formatted_text import FormattedText + +STYLES = Style.from_dict({ + "safe": "#012 bg:#00ff88", + "unsafe": "#012 bg:#ff2200", + "message": "#666699 bg:#000", + "count": "#6666ff bg:#000", + "prompt": "#777700 bg:#000", + "alert": "bold #f0f bg:#000", + "error": "bold #f00 bg:#000", + "warning": "bold #f50 bg:#000", + "good": "bold #090 bg:#000", + "promptyn": "bold #000 bg:#ff0", + "default": "#880 bg:#000", + "bright": "#ff0 bg:#000", + }) + + +def print(message): + print_formatted_text(HTML(f"{message}"), style=STYLES) + + +def print_style(message, style): + print_formatted_text(FormattedText([(f"class:{style}", message)]), style=STYLES)