scripts/pssh

   1 #!/usr/bin/env python
   2 # -*- Mode: python -*-
   3 #
   4 # Usage: pssh [OPTIONS] -h hosts.txt prog [arg0] [arg1] ..
   5 #
   6 # Parallel ssh to the set of nodes in hosts.txt. For each node, this
   7 # essentially does an "ssh host -l user prog [arg0] [arg1] ...". The -o
   8 # option can be used to store stdout from each remote node in a
   9 # directory.  Each output file in that directory will be named by the
  10 # corresponding remote node's hostname or IP address.
  11 #
  12 # Created: 16 August 2003
  13 #
  14 # $Id$
  15 # patched by thierry from 1.4.3 for adding -k and multiple -O
  16 #
  17 import fcntl
  18 import getopt
  19 import os
  20 import pwd
  21 import signal
  22 import subprocess
  23 import sys
  24 import threading
  25
  26 basedir, bin = os.path.split(os.path.dirname(os.path.abspath(sys.argv[0])))
  27 sys.path.append("%s" % basedir)
  28
  29 from psshlib import psshutil
  30 from psshlib.basethread import BaseThread
  31
  32 _DEFAULT_PARALLELISM = 32
  33 _DEFAULT_TIMEOUT     = 60
  34
  35 def print_usage():
  36     print "Usage: pssh [OPTIONS] -h hosts.txt prog [arg0] .."
  37     print
  38     print "  -h --hosts   hosts file (each line \"host[:port] [user]\")"
  39     print "  -l --user    username (OPTIONAL)"
  40     print "  -p --par     max number of parallel threads (OPTIONAL)"
  41     print "  -o --outdir  output directory for stdout files (OPTIONAL)"
  42     print "  -e --errdir  output directory for stderr files (OPTIONAL)"
  43     print "  -t --timeout timeout (secs) (-1 = no timeout) per host (OPTIONAL)"
  44     print "  -O --options SSH options (OPTIONAL)"
  45     print "  -v --verbose turn on warning and diagnostic messages (OPTIONAL)"
  46     print "  -P --print   print output as we get it (OPTIONAL)"
  47     print "  -i --inline  inline aggregated output for each server (OPTIONAL)"
  48     print "  -k --key     alternate key to pass to ssh -i (OPTIONAL)"
  49     print
  50     print "Example: pssh -h nodes.txt -l irb2 -o /tmp/foo uptime"
  51     print
  52
  53 def read_envvars(flags):
  54     if os.getenv("PSSH_HOSTS"):
  55         flags["hosts"] = os.getenv("PSSH_HOSTS")
  56     if os.getenv("PSSH_USER"):
  57         flags["user"] = os.getenv("PSSH_USER")
  58     if os.getenv("PSSH_PAR"):
  59         flags["par"] = int(os.getenv("PSSH_PAR"))
  60     if os.getenv("PSSH_OUTDIR"):
  61         flags["outdir"] = os.getenv("PSSH_OUTDIR")
  62     if os.getenv("PSSH_ERRDIR"):
  63         flags["errdir"] = os.getenv("PSSH_ERRDIR")
  64     if os.getenv("PSSH_TIMEOUT"):
  65         timeout = int(os.getenv("PSSH_TIMEOUT"))
  66         if timeout != -1:
  67             flags["timeout"] = timeout
  68         else:
  69             flags["timeout"] = None
  70     if os.getenv("PSSH_OPTIONS"):
  71         flags["options"] = os.getenv("PSSH_OPTIONS").split(" ")
  72     if os.getenv("PSSH_VERBOSE"): # "0" or "1"
  73         flags["verbose"] = int(os.getenv("PSSH_VERBOSE"))
  74     if os.getenv("PSSH_PRINT"):
  75        flags["print"] = int(os.getenv("PSSH_PRINT"))
  76     if os.getenv("PSSH_INLINE"):
  77        flags["inline"] = int(os.getenv("PSSH_INLINE"))
  78     if os.getenv("PSSH_KEY"):
  79        flags["inline"] = int(os.getenv("PSSH_KEY"))
  80
  81 def parsecmdline(argv):
  82     shortopts = "h:l:p:o:e:t:O:Pvik:"
  83     longopts = [ "hosts=",
  84                  "user=",
  85                  "par=",
  86                  "outdir=",
  87                  "errdir=",
  88                  "timeout=",
  89                  "options=",
  90                  "print",
  91                  "verbose",
  92                  "inline",
  93                  "key=",
  94                  ]
  95     flags = { "hosts" : None,
  96               "user" : None,
  97               "par" : _DEFAULT_PARALLELISM,
  98               "outdir" : None,
  99               "errdir" : None,
 100               "timeout" : _DEFAULT_TIMEOUT,
 101               "options" : [],
 102               "print" : None,
 103               "verbose" : None,
 104               "inline": None,
 105               "key" : None,
 106               }
 107     read_envvars(flags)
 108     if not flags["user"]:
 109         flags["user"] = pwd.getpwuid(os.getuid())[0] # Default to current user
 110     try:
 111         opts, args = getopt.getopt(argv[1:], shortopts, longopts)
 112     except getopt.GetoptError, e:
 113         print "Error: %s\n" % str(e)
 114         print_usage()
 115         sys.exit(3)
 116     for o, v in opts:
 117         if o in ("-h", "--hosts"):
 118             flags["hosts"] = v
 119         elif o in ("-l", "--user"):
 120             flags["user"] = v
 121         elif o in ("-p", "--par"):
 122             flags["par"] = int(v)
 123         elif o in ("-o", "--outdir"):
 124             flags["outdir"] = v
 125         elif o in ("-e", "--errdir"):
 126             flags["errdir"] = v
 127         elif o in ("-t", "--timeout"):
 128             timeout =  int(v)
 129             if timeout != -1:
 130                flags["timeout"] = timeout
 131             else:
 132                flags["timeout"] = None
 133         elif o in ("-O", "--options"):
 134             flags["options"].append(v)
 135         elif o in ("-v", "--verbose"):
 136             flags["verbose"] = 1
 137         elif o in ("-P", "--print"):
 138             flags["print"] = 1
 139         elif o in ("-i", "--inline"):
 140             flags["inline"] = 1
 141         elif o in ["-k","--key"]:
 142             flags ["key"]= v
 143     # Required flags
 144     if not flags["hosts"]:
 145         print_usage()
 146         sys.exit(3)
 147     return args, flags
 148
 149 def buffer_input():
 150     origfl = fcntl.fcntl(sys.stdin.fileno(), fcntl.F_GETFL)
 151     fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, origfl | os.O_NONBLOCK)
 152     try:
 153         stdin = sys.stdin.read()
 154     except IOError: # Stdin contained no information
 155         stdin = ""
 156     fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, origfl)
 157     return stdin
 158
 159 def do_pssh(hosts, ports, users, cmdline, flags):
 160     if flags["outdir"] and not os.path.exists(flags["outdir"]):
 161         os.makedirs(flags["outdir"])
 162     if flags["errdir"] and not os.path.exists(flags["errdir"]):
 163         os.makedirs(flags["errdir"])
 164     stdin = buffer_input()
 165     sem = threading.Semaphore(flags["par"])
 166     threads = []
 167     for i in range(len(hosts)):
 168         sem.acquire()
 169         cmd="ssh"
 170         for option in flags["options"]:
 171             cmd += " -o %s"%option
 172         if not flags["verbose"]:
 173             cmd += " -q"
 174         if flags["key"]:
 175             cmd += " -i %s"%flags["key"]
 176         cmd += " -p %s"  % ports[i]
 177         cmd += " -l %s"  %  users[i]
 178         cmd += " %s"     % hosts[i]
 179         cmd += " \"%s\"" %  cmdline
 180         if flags["verbose"]:
 181             print 'Triggering %s'%cmd
 182         t = BaseThread(hosts[i], ports[i], cmd, flags, sem, stdin)
 183         t.start()
 184         threads.append(t)
 185     for t in threads:
 186         t.join()
 187
 188 if __name__ == "__main__":
 189     subprocess._cleanup = lambda : None
 190     args, flags = parsecmdline(sys.argv)
 191     if len(args) == 0:
 192         print_usage()
 193         sys.exit(3)
 194     cmdline = " ".join(args)
 195     hosts, ports, users = psshutil.read_hosts(flags["hosts"])
 196     psshutil.patch_users(hosts, ports, users, flags["user"])
 197     signal.signal(signal.SIGCHLD, signal.SIG_DFL)
 198     os.setpgid(0, 0)
 199     do_pssh(hosts, ports, users, cmdline, flags)