broadcast ssh - my entry into the world of parallel or cluster ssh tools

I threatened before when writing up my review of the parallel or cluster ssh tools out there that I would write my own. And after a quick review of paramiko, especially the demo scripts included, this turned out to be a pretty quick hack.

The basics of my ssh client are this:

• prompt for username and password - ssh keys not required (although I would like to add support for using them if available in the future)
• interactive use - I often want to look at things, which leads me to want to look at other things; in other words, I don't want to be constrained by having a list of commands to execute up front
• parallel - I have a bunch of commands and a bunch of machines, a simple loop executing each command on each machine isn't going to cut it

Those last two together make things a little tricky, the interactive bit means you have two choices.  You can be line oriented, prompt for a command and send it to each host or you can be completely interactive and send each character.

I started with the former, but soon found that if you only send commands, you have practically no environment setup (that is all done in the shell, remember...).  So, despite the difficulties in dealing with the terminal, that is what this implementation does.

A terminal based solution has its own issues, the line buffering of output from each host has to be dealt with to avoid interleaved garbage as the results of each command.  But in the end that wasn't too hard.

In the end, usage is simple, cut-n-paste the code, save it as bssh.py and run something like:

./bssh.py host1 host2 host3 host4

Enter a username and password at the prompts, and away you go, just enter commands as you would with any other shell. You will of course need paramiko and its dependency pycrypto installed and available.

#!/usr/bin/env python

import sys, getpass, paramiko, socket, tty, termios, traceback
from select import select

user=raw_input('Username: ') or getpass.getuser()
password=getpass.getpass()
port=22

def main(hosts):
  clients = {}
  try:
    for h in hosts:
      client = paramiko.SSHClient()
      client.load_system_host_keys()
      client.set_missing_host_key_policy(paramiko.WarningPolicy())
      print 'connecting to', h, 'as', user
      client.connect(h, port, user, password)
      print 'invoking shell on', h
      clients[h] = {
        'client': client,
        'channel': client.invoke_shell(),
        'output': ''
      }

    oldtty = termios.tcgetattr(sys.stdin)
    try:
      tty.setraw(sys.stdin.fileno())
      tty.setcbreak(sys.stdin.fileno())
      for h in clients:
        clients[h]['channel'].settimeout(0.0)

      while clients:
        channels = [clients[h]['channel'] for h in clients]
        rd, wr, err = select([sys.stdin] + channels, [], [])
        closed = []
        for h in clients:
          ch = clients[h]['channel']
          if ch in rd:
            try:
              x = ch.recv(1024)
              if len(x) == 0:
                ch.close()
                clients[h]['client'].close()
                closed.append(h)
              else:
                x = clients[h]['output'] + x
                lines = x.splitlines()
                if x.endswith('\n'):
                  clients[h]['output'] = ''
                else:
                  clients[h]['output'] = lines[-1]
                  del lines[-1]

                for line in lines:
                  sys.stdout.write('%s> %s\n\r' % (h, line))
                  sys.stdout.flush()
            except socket.timeout:
              pass

        for h in closed:
          del clients[h]

        if sys.stdin in rd:
          x = sys.stdin.read(1)
          for h in clients:
            clients[h]['channel'].send(x)
          sys.stdout.write(x)
          sys.stdout.flush()
          
    finally:
      termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
  except:
    traceback.print_exc()
    for h in clients:
      try:
        clients[h]['channel'].close()
      except:
        pass
      try:
        clients[h]['client'].close()
      except:
        pass

if __name__ == '__main__':
  main(sys.argv[1:])

Of course, the hundred (or so) lines above could use some work, things like using different usernames and passwords for various hosts, ssh key (or other) authentication, better (or maybe just shorter) labels to prefix each output line, and I'm sure there are more. But for now, it's a workable solution that has payed dividends already in my productivity.

The bottom line here is: paramiko rocks. From the demos to the shear flexibility of being able to program ssh sessions at the protocol level, I'm sure this is going to be an oft-used addition to my toolbox.