initial - patch from pssh-1.4.3
authorthierry <thierry@41d37cc5-eb28-0410-a9bf-d37491348ade>
Fri, 17 Apr 2009 09:36:45 +0000 (09:36 +0000)
committerthierry <thierry@41d37cc5-eb28-0410-a9bf-d37491348ade>
Fri, 17 Apr 2009 09:36:45 +0000 (09:36 +0000)
scripts/pssh [new file with mode: 0755]

diff --git a/scripts/pssh b/scripts/pssh
new file mode 100755 (executable)
index 0000000..44d9dfc
--- /dev/null
@@ -0,0 +1,199 @@
+#!/usr/bin/env python
+# -*- Mode: python -*-
+#
+# Usage: pssh [OPTIONS] -h hosts.txt prog [arg0] [arg1] ..
+#
+# Parallel ssh to the set of nodes in hosts.txt. For each node, this
+# essentially does an "ssh host -l user prog [arg0] [arg1] ...". The -o
+# option can be used to store stdout from each remote node in a
+# directory.  Each output file in that directory will be named by the
+# corresponding remote node's hostname or IP address.
+#
+# Created: 16 August 2003
+#
+# $Id$
+# patched by thierry from 1.4.3 for adding -k and multiple -O
+#
+import fcntl
+import getopt
+import os
+import pwd
+import signal
+import subprocess
+import sys
+import threading
+
+basedir, bin = os.path.split(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append("%s" % basedir)
+
+from psshlib import psshutil
+from psshlib.basethread import BaseThread
+
+_DEFAULT_PARALLELISM = 32
+_DEFAULT_TIMEOUT     = 60
+
+def print_usage():
+    print "Usage: pssh [OPTIONS] -h hosts.txt prog [arg0] .."
+    print
+    print "  -h --hosts   hosts file (each line \"host[:port] [user]\")"
+    print "  -l --user    username (OPTIONAL)"
+    print "  -p --par     max number of parallel threads (OPTIONAL)"
+    print "  -o --outdir  output directory for stdout files (OPTIONAL)"
+    print "  -e --errdir  output directory for stderr files (OPTIONAL)"
+    print "  -t --timeout timeout (secs) (-1 = no timeout) per host (OPTIONAL)"
+    print "  -O --options SSH options (OPTIONAL)"
+    print "  -v --verbose turn on warning and diagnostic messages (OPTIONAL)"
+    print "  -P --print   print output as we get it (OPTIONAL)"
+    print "  -i --inline  inline aggregated output for each server (OPTIONAL)"
+    print "  -k --key     alternate key to pass to ssh -i (OPTIONAL)"
+    print
+    print "Example: pssh -h nodes.txt -l irb2 -o /tmp/foo uptime"
+    print
+
+def read_envvars(flags):
+    if os.getenv("PSSH_HOSTS"):
+        flags["hosts"] = os.getenv("PSSH_HOSTS")
+    if os.getenv("PSSH_USER"):
+        flags["user"] = os.getenv("PSSH_USER")
+    if os.getenv("PSSH_PAR"):
+        flags["par"] = int(os.getenv("PSSH_PAR"))
+    if os.getenv("PSSH_OUTDIR"):
+        flags["outdir"] = os.getenv("PSSH_OUTDIR")
+    if os.getenv("PSSH_ERRDIR"):
+        flags["errdir"] = os.getenv("PSSH_ERRDIR")
+    if os.getenv("PSSH_TIMEOUT"):
+        timeout = int(os.getenv("PSSH_TIMEOUT"))
+        if timeout != -1:
+            flags["timeout"] = timeout
+        else:
+            flags["timeout"] = None
+    if os.getenv("PSSH_OPTIONS"):
+        flags["options"] = os.getenv("PSSH_OPTIONS").split(" ")
+    if os.getenv("PSSH_VERBOSE"): # "0" or "1"
+        flags["verbose"] = int(os.getenv("PSSH_VERBOSE"))
+    if os.getenv("PSSH_PRINT"):
+       flags["print"] = int(os.getenv("PSSH_PRINT"))
+    if os.getenv("PSSH_INLINE"):
+       flags["inline"] = int(os.getenv("PSSH_INLINE"))
+    if os.getenv("PSSH_KEY"):
+       flags["inline"] = int(os.getenv("PSSH_KEY"))
+
+def parsecmdline(argv):
+    shortopts = "h:l:p:o:e:t:O:Pvik:"
+    longopts = [ "hosts=",
+                 "user=",
+                 "par=",
+                 "outdir=",
+                 "errdir=",
+                 "timeout=",
+                 "options=",
+                 "print",
+                 "verbose",
+                 "inline",
+                 "key=",
+                 ]
+    flags = { "hosts" : None,
+              "user" : None,
+              "par" : _DEFAULT_PARALLELISM,
+              "outdir" : None,
+              "errdir" : None,
+              "timeout" : _DEFAULT_TIMEOUT,
+              "options" : [],
+              "print" : None,
+              "verbose" : None,
+              "inline": None,
+              "key" : None,
+              }
+    read_envvars(flags)
+    if not flags["user"]:
+        flags["user"] = pwd.getpwuid(os.getuid())[0] # Default to current user
+    try:
+        opts, args = getopt.getopt(argv[1:], shortopts, longopts)
+    except getopt.GetoptError, e:
+        print "Error: %s\n" % str(e)
+        print_usage()
+        sys.exit(3)
+    for o, v in opts:
+        if o in ("-h", "--hosts"):
+            flags["hosts"] = v
+        elif o in ("-l", "--user"):
+            flags["user"] = v
+        elif o in ("-p", "--par"):
+            flags["par"] = int(v)
+        elif o in ("-o", "--outdir"):
+            flags["outdir"] = v
+        elif o in ("-e", "--errdir"):
+            flags["errdir"] = v
+        elif o in ("-t", "--timeout"):
+            timeout =  int(v)
+            if timeout != -1:
+               flags["timeout"] = timeout
+            else:
+               flags["timeout"] = None
+        elif o in ("-O", "--options"):
+            flags["options"].append(v)
+        elif o in ("-v", "--verbose"):
+            flags["verbose"] = 1
+        elif o in ("-P", "--print"):
+            flags["print"] = 1
+        elif o in ("-i", "--inline"):
+            flags["inline"] = 1
+        elif o in ["-k","--key"]:
+            flags ["key"]= v
+    # Required flags
+    if not flags["hosts"]:
+        print_usage()
+        sys.exit(3)
+    return args, flags
+
+def buffer_input():
+    origfl = fcntl.fcntl(sys.stdin.fileno(), fcntl.F_GETFL)
+    fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, origfl | os.O_NONBLOCK)
+    try:
+        stdin = sys.stdin.read()
+    except IOError: # Stdin contained no information
+        stdin = ""
+    fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, origfl)
+    return stdin
+
+def do_pssh(hosts, ports, users, cmdline, flags):
+    if flags["outdir"] and not os.path.exists(flags["outdir"]):
+        os.makedirs(flags["outdir"])
+    if flags["errdir"] and not os.path.exists(flags["errdir"]):
+        os.makedirs(flags["errdir"])
+    stdin = buffer_input()
+    sem = threading.Semaphore(flags["par"])
+    threads = []
+    for i in range(len(hosts)):
+        sem.acquire()
+        cmd="ssh"
+        for option in flags["options"]:
+            cmd += " -o %s"%option
+        if not flags["verbose"]:
+            cmd += " -q"
+        if flags["key"]:
+            cmd += " -i %s"%flags["key"]
+        cmd += " -p %s"  % ports[i]
+        cmd += " -l %s"  %  users[i]
+        cmd += " %s"     % hosts[i]
+        cmd += " \"%s\"" %  cmdline
+        if flags["verbose"]:
+            print 'Triggering %s'%cmd
+        t = BaseThread(hosts[i], ports[i], cmd, flags, sem, stdin)
+        t.start()
+        threads.append(t)
+    for t in threads:
+        t.join()
+   
+if __name__ == "__main__":
+    subprocess._cleanup = lambda : None
+    args, flags = parsecmdline(sys.argv)
+    if len(args) == 0:
+        print_usage()
+        sys.exit(3)
+    cmdline = " ".join(args)
+    hosts, ports, users = psshutil.read_hosts(flags["hosts"])
+    psshutil.patch_users(hosts, ports, users, flags["user"])
+    signal.signal(signal.SIGCHLD, signal.SIG_DFL)
+    os.setpgid(0, 0)
+    do_pssh(hosts, ports, users, cmdline, flags)