src/link/pexpect/examples/hive.py

   1 #!/usr/bin/env python
   2
   3 """hive -- Hive Shell
   4
   5 This lets you ssh to a group of servers and control them as if they were one.
   6 Each command you enter is sent to each host in parallel. The response of each
   7 host is collected and printed. In normal synchronous mode Hive will wait for
   8 each host to return the shell command line prompt. The shell prompt is used to
   9 sync output.
  10
  11 Example:
  12
  13     $ hive.py --sameuser --samepass host1.example.com host2.example.net
  14     username: myusername
  15     password:
  16     connecting to host1.example.com - OK
  17     connecting to host2.example.net - OK
  18     targetting hosts: 192.168.1.104 192.168.1.107
  19     CMD (? for help) > uptime
  20     =======================================================================
  21     host1.example.com
  22     -----------------------------------------------------------------------
  23     uptime
  24     23:49:55 up 74 days,  5:14,  2 users,  load average: 0.15, 0.05, 0.01
  25     =======================================================================
  26     host2.example.net
  27     -----------------------------------------------------------------------
  28     uptime
  29     23:53:02 up 1 day, 13:36,  2 users,  load average: 0.50, 0.40, 0.46
  30     =======================================================================
  31
  32 Other Usage Examples:
  33
  34 1. You will be asked for your username and password for each host.
  35
  36     hive.py host1 host2 host3 ... hostN
  37
  38 2. You will be asked once for your username and password.
  39    This will be used for each host.
  40
  41     hive.py --sameuser --samepass host1 host2 host3 ... hostN
  42
  43 3. Give a username and password on the command-line:
  44
  45     hive.py user1:pass2@host1 user2:pass2@host2 ... userN:passN@hostN
  46
  47 You can use an extended host notation to specify username, password, and host
  48 instead of entering auth information interactively. Where you would enter a
  49 host name use this format:
  50
  51     username:password@host
  52
  53 This assumes that ':' is not part of the password. If your password contains a
  54 ':' then you can use '\\:' to indicate a ':' and '\\\\' to indicate a single
  55 '\\'. Remember that this information will appear in the process listing. Anyone
  56 on your machine can see this auth information. This is not secure.
  57
  58 This is a crude script that begs to be multithreaded. But it serves its
  59 purpose.
  60
  61 Noah Spurrier
  62
  63 $Id: hive.py 509 2008-01-05 21:27:47Z noah $
  64 """
  65
  66 # TODO add feature to support username:password@host combination
  67 # TODO add feature to log each host output in separate file
  68
  69 import sys, os, re, optparse, traceback, types, time, getpass
  70 import pexpect, pxssh
  71 import readline, atexit
  72
  73 #histfile = os.path.join(os.environ["HOME"], ".hive_history")
  74 #try:
  75 #    readline.read_history_file(histfile)
  76 #except IOError:
  77 #    pass
  78 #atexit.register(readline.write_history_file, histfile)
  79
  80 CMD_HELP="""Hive commands are preceded by a colon : (just think of vi).
  81
  82 :target name1 name2 name3 ...
  83
  84     set list of hosts to target commands
  85
  86 :target all
  87
  88     reset list of hosts to target all hosts in the hive.
  89
  90 :to name command
  91
  92     send a command line to the named host. This is similar to :target, but
  93     sends only one command and does not change the list of targets for future
  94     commands.
  95
  96 :sync
  97
  98     set mode to wait for shell prompts after commands are run. This is the
  99     default. When Hive first logs into a host it sets a special shell prompt
 100     pattern that it can later look for to synchronize output of the hosts. If
 101     you 'su' to another user then it can upset the synchronization. If you need
 102     to run something like 'su' then use the following pattern:
 103
 104     CMD (? for help) > :async
 105     CMD (? for help) > sudo su - root
 106     CMD (? for help) > :prompt
 107     CMD (? for help) > :sync
 108
 109 :async
 110
 111     set mode to not expect command line prompts (see :sync). Afterwards
 112     commands are send to target hosts, but their responses are not read back
 113     until :sync is run. This is useful to run before commands that will not
 114     return with the special shell prompt pattern that Hive uses to synchronize.
 115
 116 :refresh
 117
 118     refresh the display. This shows the last few lines of output from all hosts.
 119     This is similar to resync, but does not expect the promt. This is useful
 120     for seeing what hosts are doing during long running commands.
 121
 122 :resync
 123
 124     This is similar to :sync, but it does not change the mode. It looks for the
 125     prompt and thus consumes all input from all targetted hosts.
 126
 127 :prompt
 128
 129     force each host to reset command line prompt to the special pattern used to
 130     synchronize all the hosts. This is useful if you 'su' to a different user
 131     where Hive would not know the prompt to match.
 132
 133 :send my text
 134
 135     This will send the 'my text' wihtout a line feed to the targetted hosts.
 136     This output of the hosts is not automatically synchronized.
 137
 138 :control X
 139
 140     This will send the given control character to the targetted hosts.
 141     For example, ":control c" will send ASCII 3.
 142
 143 :exit
 144
 145     This will exit the hive shell.
 146
 147 """
 148
 149 def login (args, cli_username=None, cli_password=None):
 150
 151     # I have to keep a separate list of host names because Python dicts are not ordered.
 152     # I want to keep the same order as in the args list.
 153     host_names = []
 154     hive_connect_info = {}
 155     hive = {}
 156     # build up the list of connection information (hostname, username, password, port)
 157     for host_connect_string in args:
 158         hcd = parse_host_connect_string (host_connect_string)
 159         hostname = hcd['hostname']
 160         port     = hcd['port']
 161         if port == '':
 162             port = None
 163         if len(hcd['username']) > 0:
 164             username = hcd['username']
 165         elif cli_username is not None:
 166             username = cli_username
 167         else:
 168             username = raw_input('%s username: ' % hostname)
 169         if len(hcd['password']) > 0:
 170             password = hcd['password']
 171         elif cli_password is not None:
 172             password = cli_password
 173         else:
 174             password = getpass.getpass('%s password: ' % hostname)
 175         host_names.append(hostname)
 176         hive_connect_info[hostname] = (hostname, username, password, port)
 177     # build up the list of hive connections using the connection information.
 178     for hostname in host_names:
 179         print 'connecting to', hostname
 180         try:
 181             fout = file("log_"+hostname, "w")
 182             hive[hostname] = pxssh.pxssh()
 183             hive[hostname].login(*hive_connect_info[hostname])
 184             print hive[hostname].before
 185             hive[hostname].logfile = fout
 186             print '- OK'
 187         except Exception, e:
 188             print '- ERROR',
 189             print str(e)
 190             print 'Skipping', hostname
 191             hive[hostname] = None
 192     return host_names, hive
 193
 194 def main ():
 195
 196     global options, args, CMD_HELP
 197
 198     if options.sameuser:
 199         cli_username = raw_input('username: ')
 200     else:
 201         cli_username = None
 202
 203     if options.samepass:
 204         cli_password = getpass.getpass('password: ')
 205     else:
 206         cli_password = None
 207
 208     host_names, hive = login(args, cli_username, cli_password)
 209
 210     synchronous_mode = True
 211     target_hostnames = host_names[:]
 212     print 'targetting hosts:', ' '.join(target_hostnames)
 213     while True:
 214         cmd = raw_input('CMD (? for help) > ')
 215         cmd = cmd.strip()
 216         if cmd=='?' or cmd==':help' or cmd==':h':
 217             print CMD_HELP
 218             continue
 219         elif cmd==':refresh':
 220             refresh (hive, target_hostnames, timeout=0.5)
 221             for hostname in target_hostnames:
 222                 if hive[hostname] is None:
 223                     print '/============================================================================='
 224                     print '| ' + hostname + ' is DEAD'
 225                     print '\\-----------------------------------------------------------------------------'
 226                 else:
 227                     print '/============================================================================='
 228                     print '| ' + hostname
 229                     print '\\-----------------------------------------------------------------------------'
 230                     print hive[hostname].before
 231             print '=============================================================================='
 232             continue
 233         elif cmd==':resync':
 234             resync (hive, target_hostnames, timeout=0.5)
 235             for hostname in target_hostnames:
 236                 if hive[hostname] is None:
 237                     print '/============================================================================='
 238                     print '| ' + hostname + ' is DEAD'
 239                     print '\\-----------------------------------------------------------------------------'
 240                 else:
 241                     print '/============================================================================='
 242                     print '| ' + hostname
 243                     print '\\-----------------------------------------------------------------------------'
 244                     print hive[hostname].before
 245             print '=============================================================================='
 246             continue
 247         elif cmd==':sync':
 248             synchronous_mode = True
 249             resync (hive, target_hostnames, timeout=0.5)
 250             continue
 251         elif cmd==':async':
 252             synchronous_mode = False
 253             continue
 254         elif cmd==':prompt':
 255             for hostname in target_hostnames:
 256                 try:
 257                     if hive[hostname] is not None:
 258                         hive[hostname].set_unique_prompt()
 259                 except Exception, e:
 260                     print "Had trouble communicating with %s, so removing it from the target list." % hostname
 261                     print str(e)
 262                     hive[hostname] = None
 263             continue
 264         elif cmd[:5] == ':send':
 265             cmd, txt = cmd.split(None,1)
 266             for hostname in target_hostnames:
 267                 try:
 268                     if hive[hostname] is not None:
 269                         hive[hostname].send(txt)
 270                 except Exception, e:
 271                     print "Had trouble communicating with %s, so removing it from the target list." % hostname
 272                     print str(e)
 273                     hive[hostname] = None
 274             continue
 275         elif cmd[:3] == ':to':
 276             cmd, hostname, txt = cmd.split(None,2)
 277             if hive[hostname] is None:
 278                 print '/============================================================================='
 279                 print '| ' + hostname + ' is DEAD'
 280                 print '\\-----------------------------------------------------------------------------'
 281                 continue
 282             try:
 283                 hive[hostname].sendline (txt)
 284                 hive[hostname].prompt(timeout=2)
 285                 print '/============================================================================='
 286                 print '| ' + hostname
 287                 print '\\-----------------------------------------------------------------------------'
 288                 print hive[hostname].before
 289             except Exception, e:
 290                 print "Had trouble communicating with %s, so removing it from the target list." % hostname
 291                 print str(e)
 292                 hive[hostname] = None
 293             continue
 294         elif cmd[:7] == ':expect':
 295             cmd, pattern = cmd.split(None,1)
 296             print 'looking for', pattern
 297             try:
 298                 for hostname in target_hostnames:
 299                     if hive[hostname] is not None:
 300                         hive[hostname].expect(pattern)
 301                         print hive[hostname].before
 302             except Exception, e:
 303                 print "Had trouble communicating with %s, so removing it from the target list." % hostname
 304                 print str(e)
 305                 hive[hostname] = None
 306             continue
 307         elif cmd[:7] == ':target':
 308             target_hostnames = cmd.split()[1:]
 309             if len(target_hostnames) == 0 or target_hostnames[0] == all:
 310                 target_hostnames = host_names[:]
 311             print 'targetting hosts:', ' '.join(target_hostnames)
 312             continue
 313         elif cmd == ':exit' or cmd == ':q' or cmd == ':quit':
 314             break
 315         elif cmd[:8] == ':control' or cmd[:5] == ':ctrl' :
 316             cmd, c = cmd.split(None,1)
 317             if ord(c)-96 < 0 or ord(c)-96 > 255:
 318                 print '/============================================================================='
 319                 print '| Invalid character. Must be [a-zA-Z], @, [, ], \\, ^, _, or ?'
 320                 print '\\-----------------------------------------------------------------------------'
 321                 continue
 322             for hostname in target_hostnames:
 323                 try:
 324                     if hive[hostname] is not None:
 325                         hive[hostname].sendcontrol(c)
 326                 except Exception, e:
 327                     print "Had trouble communicating with %s, so removing it from the target list." % hostname
 328                     print str(e)
 329                     hive[hostname] = None
 330             continue
 331         elif cmd == ':esc':
 332             for hostname in target_hostnames:
 333                 if hive[hostname] is not None:
 334                     hive[hostname].send(chr(27))
 335             continue
 336         #
 337         # Run the command on all targets in parallel
 338         #
 339         for hostname in target_hostnames:
 340             try:
 341                 if hive[hostname] is not None:
 342                     hive[hostname].sendline (cmd)
 343             except Exception, e:
 344                 print "Had trouble communicating with %s, so removing it from the target list." % hostname
 345                 print str(e)
 346                 hive[hostname] = None
 347
 348         #
 349         # print the response for each targeted host.
 350         #
 351         if synchronous_mode:
 352             for hostname in target_hostnames:
 353                 try:
 354                     if hive[hostname] is None:
 355                         print '/============================================================================='
 356                         print '| ' + hostname + ' is DEAD'
 357                         print '\\-----------------------------------------------------------------------------'
 358                     else:
 359                         hive[hostname].prompt(timeout=2)
 360                         print '/============================================================================='
 361                         print '| ' + hostname
 362                         print '\\-----------------------------------------------------------------------------'
 363                         print hive[hostname].before
 364                 except Exception, e:
 365                     print "Had trouble communicating with %s, so removing it from the target list." % hostname
 366                     print str(e)
 367                     hive[hostname] = None
 368             print '=============================================================================='
 369
 370 def refresh (hive, hive_names, timeout=0.5):
 371
 372     """This waits for the TIMEOUT on each host.
 373     """
 374
 375     # TODO This is ideal for threading.
 376     for hostname in hive_names:
 377         hive[hostname].expect([pexpect.TIMEOUT,pexpect.EOF],timeout=timeout)
 378
 379 def resync (hive, hive_names, timeout=2, max_attempts=5):
 380
 381     """This waits for the shell prompt for each host in an effort to try to get
 382     them all to the same state. The timeout is set low so that hosts that are
 383     already at the prompt will not slow things down too much. If a prompt match
 384     is made for a hosts then keep asking until it stops matching. This is a
 385     best effort to consume all input if it printed more than one prompt. It's
 386     kind of kludgy. Note that this will always introduce a delay equal to the
 387     timeout for each machine. So for 10 machines with a 2 second delay you will
 388     get AT LEAST a 20 second delay if not more. """
 389
 390     # TODO This is ideal for threading.
 391     for hostname in hive_names:
 392         for attempts in xrange(0, max_attempts):
 393             if not hive[hostname].prompt(timeout=timeout):
 394                 break
 395
 396 def parse_host_connect_string (hcs):
 397
 398     """This parses a host connection string in the form
 399     username:password@hostname:port. All fields are options expcet hostname. A
 400     dictionary is returned with all four keys. Keys that were not included are
 401     set to empty strings ''. Note that if your password has the '@' character
 402     then you must backslash escape it. """
 403
 404     if '@' in hcs:
 405         p = re.compile (r'(?P<username>[^@:]*)(:?)(?P<password>.*)(?!\\)@(?P<hostname>[^:]*):?(?P<port>[0-9]*)')
 406     else:
 407         p = re.compile (r'(?P<username>)(?P<password>)(?P<hostname>[^:]*):?(?P<port>[0-9]*)')
 408     m = p.search (hcs)
 409     d = m.groupdict()
 410     d['password'] = d['password'].replace('\\@','@')
 411     return d
 412
 413 if __name__ == '__main__':
 414     try:
 415         start_time = time.time()
 416         parser = optparse.OptionParser(formatter=optparse.TitledHelpFormatter(), usage=globals()['__doc__'], version='$Id: hive.py 509 2008-01-05 21:27:47Z noah $',conflict_handler="resolve")
 417         parser.add_option ('-v', '--verbose', action='store_true', default=False, help='verbose output')
 418         parser.add_option ('--samepass', action='store_true', default=False, help='Use same password for each login.')
 419         parser.add_option ('--sameuser', action='store_true', default=False, help='Use same username for each login.')
 420         (options, args) = parser.parse_args()
 421         if len(args) < 1:
 422             parser.error ('missing argument')
 423         if options.verbose: print time.asctime()
 424         main()
 425         if options.verbose: print time.asctime()
 426         if options.verbose: print 'TOTAL TIME IN MINUTES:',
 427         if options.verbose: print (time.time() - start_time) / 60.0
 428         sys.exit(0)
 429     except KeyboardInterrupt, e: # Ctrl-C
 430         raise e
 431     except SystemExit, e: # sys.exit()
 432         raise e
 433     except Exception, e:
 434         print 'ERROR, UNEXPECTED EXCEPTION'
 435         print str(e)
 436         traceback.print_exc()
 437         os._exit(1)