Colorizing Start!

- Add styles and colors to most output.
- Update TODO list
- Add ability to interrupt "waiting" condition with Control-C

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

multiball/__init__.py

@@ -1 +1 @@
-__version__ = "0.5.0"
+__version__ = "0.6.0"

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"<message>Loaded {len(newhosts)} from {sc}</message>")
                     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"<message>Loaded {len(newhosts)} from {sc}</message>")
                     self.allhosts.update(newhosts)
 
         else:
             self.ssh_config = Path("~/.ssh/ssh_config").expanduser()
             print(
-                f"Warning: no [config] section in {configFile}. Using fallbacks.")
+                f"<error>Warning: no [config] section in {configFile}. Using fallbacks.</error>")
             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("<unsafe>** Safety OFF. Now unsafe.</unsafe>")
         else:
             self.safety = True
-            print("** Safety ON.")
+            print("<safe>** Safety ON.</safe>")
 
     def command_echo(self, command, args):
         """
         Echo string
         fixme put string interpolation here so that state variables can be printed
         """
-        print(args)
+        print(f"<message>{args}</message>")
 
     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(
+        result = prompt_toolkit.prompt([("class:promptyn",
-            "<" + args + " [Yys/Nn] > ")
+                                         "<" + args + " [Yys/Nn]>"),
+                                        ("class:message",
+                                         " ")], style=STYLES)
         if result and (result[0] in 'Yys'):
             return True
         else:
-            print("(no)")
+            print("<message>(no)</message>")
             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("<alert>does nothing outside of script</alert>")
             else:
-                print("arguments required")
+                print("<error>arguments required</error>")
             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("<alert>does nothing outside of script</alert>")
             else:
-                print("cannot run in safety mode")
+                print("<error>cannot run in safety mode</error>")
             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("<alert>Nothing to log.</alert>")
             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"<alert>Writing {outfilename} for command `{self.last_run['cmd']}`.</alert>")
         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"<message>{keyword} {len(self.targethosts)} hosts:</message>")
         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"<bright>{m}</bright>")
         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(f"<bright>would run `{command}` on targethosts</bright>")
-                print("Safety is ON. `/safety off` to turn off.")
+                print("<alert>Safety is ON. `/safety off` to turn off.</alert>")
             else:
-                print(f"running `{command}` on targethosts")
+                print(f"<bright>running `{command}` on targethosts</bright>")
                 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(
+            pclass = "class:unsafe"
-            f'{len(self.targethosts)} -> {prompt} ', completer=completer)
+        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"<alert>Welcome to Multiball {__version__}. type /help for help.</alert>")
     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"<error>Exception! {inst}</error")
         print(traceback.format_exc())
         sys.exit(-1)
 

multiball/fabtools.py

@@ -9,6 +9,7 @@ import paramiko
 from fabric2 import Config, Connection, SerialGroup, ThreadingGroup
 from tqdm import tqdm
 
+from .style import print_style
 
 class RunResult(Enum):
     Success = 0
@@ -20,6 +21,7 @@ class RunResult(Enum):
 def thread_run(connection, command, result_lock, result_queue):
     runresult = RunResult.Success
     try:
+        res = connection.open()
         res = connection.run(command, warn=True, hide=True)
         if (res.exited != 0):
             runresult = RunResult.CommandFailure
@@ -32,15 +34,22 @@ def thread_run(connection, command, result_lock, result_queue):
     except Exception as inst:
         res = f"Unknown error for host: {inst}"
         runresult = RunResult.Error
+    finally:
+        connection.close()
+
     with result_lock:
         result_queue.append((connection, res, runresult))
 
 
+
 # A set of hosts we can target with a series of commands and also can collect output of each command
 class HostSet:
-    def _lprint(self, *args):
+    def _lprint(self, *args, **kwargs):
         s = " ".join([str(a) for a in args]) + '\n'
         self.results.write(s)
+        if ("style" in kwargs):
+           print_style(s.rstrip(), style=kwargs["style"])
+        else:
             print(*args)
 
     def __init__(self, hostlist: list, ssh_config_path: Path=None):
@@ -62,11 +71,13 @@ class HostSet:
         for connection in self.connections:
             t = Thread(target=thread_run, args=[connection, command, reslock, resq])
             t.host = connection.original_host
+            t.connection = connection
             t.start()
             threads.append(t)
 
         # display status about threads, and join them when they finish
         while (True):
+            try:
                 nt = []
                 for i in threads:
                     if not i.is_alive():
@@ -81,8 +92,15 @@ class HostSet:
                 else:
                     if ((time.time() - tupdate) > 10.0):
                         tupdate = time.time()
-                    print("Still waiting on ", [thread.host for thread in threads])
+                        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")

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"<default>{message}</default>"), style=STYLES)
+
+
+def print_style(message, style):
+    print_formatted_text(FormattedText([(f"class:{style}", message)]), style=STYLES)