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