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