tools/bin/gpssh

#!/usr/bin/env python
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
# 
#   http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

'''
gpssh -- ssh access to multiple hosts at once

Usage: gpssh [--version] [-?v] [-h host] [-f hostfile] [cmd]

          --version : print version information
          -?        : print this help screen
          -v        : verbose mode
          -e        : echo commands as they are executed
          -h host   : the host to connect to (multiple -h is okay)
          -f file   : a file listing all hosts to connect to
          -D        : do not filter multi-homed hosts 
          cmd       : the command to execute. If not present, 
                      go into interactive mode
'''
import os
import sys

progname = os.path.split(sys.argv[0])[-1]
sys.path.append(sys.path[0] + '/lib')

if sys.version_info < (2, 5, 0):
    sys.exit(
'''Error: %s is supported on Python versions 2.5 or greater
Please upgrade python installed on this machine.''' % progname)


import getopt
import atexit
import signal
import pexpect
import time
from gppylib.util import ssh_utils
from gppylib.gpparseopts import OptParser
from gppylib.gpcoverage import GpCoverage


#
# all the command line options
#
class __globals__:
    script_name = os.path.split(__file__)[-1]
    USER = os.environ.get('LOGNAME') or os.environ.get('USER')
    opt = {}
    opt['-v'] = False
    opt['-e'] = False
    opt['-h'] = []
    opt['-f'] = False
    opt['-D'] = False
    argcmd = None
    session = None


GV = __globals__()

################
def usage(exitarg):
    parser = OptParser()
    try:
        parser.print_help()
    except:
        print __doc__
    sys.exit(0)


#############
def print_version():
    print '%s version $Revision$' % GV.script_name
    sys.exit(0)

#############
def parseCommandLine():
    try:
        (options, args) = getopt.getopt(sys.argv[1:], '?evh:f:D:u:', ['version'])
    except Exception, e:
        usage('[ERROR] ' + str(e))

    for (switch, val) in options:
        if (switch == '-?'):              usage(0)
        elif (switch[1] in 'evD'):    GV.opt[switch] = True
        elif (switch[1] in 'f'):      GV.opt[switch] = val
        elif (switch == '-h'):        GV.opt[switch].append(val)
        elif (switch == '--version'): print_version()
        elif (switch[1] in 'u'):      GV.USER = val

    hf = (len(GV.opt['-h']) and 1 or 0) + (GV.opt['-f'] and 1 or 0)
    if hf != 1:
        usage('Error: please specify at least one of -h or -f args, but not both')

    if (len(args) >= 1):
        GV.argcmd = " ".join(args)
           
def sessionCleanup():
    while True:
        try:
            return_code = 0
            if GV.session: 
                if GV.session.verbose: print '\n[Cleanup...]'; 
                return_code = GV.session.close()
            GV.session = None
            return return_code
        except KeyboardInterrupt:
            pass

sigint_time = 0
def sigint_handle(signum, frame):
    global sigint_time
    now = time.time()
    if now - sigint_time >= 3: 
        sigint_time = now
        raise KeyboardInterrupt
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    print '\n[Exiting...]'
    sys.exit(1)

def sighup_handle(signum, frame):
    sys.exit(1)

def interactive():
    try:
        import readline
        #Read in the saved command history, if any
        histfile = os.path.join(os.environ["HOME"], ".gshist")
        # Set the maximum number of commands to 100
        readline.set_history_length(500)
        try:
            readline.read_history_file(histfile)            
        except IOError:
            pass
        
        #MPP-4054 - let's check the permissons before we register
        try:
            f=open(histfile,'a')
            atexit.register(readline.write_history_file, histfile)
            f.close()                        
        except IOError:
            print "\n[WARN] Unable to write to gpssh history file: '%s'. Please check permissions." % histfile
        
            
    except Exception, e:
        print "Note: command history unsupported on this machine ..."

    atexit.register(sessionCleanup)
    signal.signal(signal.SIGINT, sigint_handle)
    signal.signal(signal.SIGHUP, sighup_handle)
    while True:
        try:
            if not GV.session:
                GV.session = ssh_utils.Session()
                GV.session.verbose=GV.opt['-v']
                GV.session.login(GV.opt['-h'], GV.USER)
                GV.session.echoCommand=GV.opt['-e']
            GV.session.cmdloop()
        except pexpect.EOF:
            print '\n[Unexpected EOF from some hosts...]'
            pass
        except ssh_utils.Session.SessionCmdExit:
                print ''
                break
        except ssh_utils.Session.SessionError, e:
                print 'Error: %s' % e
                pass
        except KeyboardInterrupt:
            print '\n[Interrupt...]'
            GV.session.reset()
            pass

#############
coverage = GpCoverage()
coverage.start()

try:
    parseCommandLine()              #Read options from the command line

    #Acquire the list of hosts from command line arguments
    hostlist = ssh_utils.HostList()
    for h in GV.opt['-h']:
        hostlist.add(h)
    if GV.opt['-f']:
        hostlist.parseFile(GV.opt['-f'])

    #Filter out non-unique hostnames unless the -D option is provided
    if not GV.opt['-D']:
        GV.opt['-h'] = hostlist.filterMultiHomedHosts()
    else:
        GV.opt['-h'] = hostlist.get()

    if len(GV.opt['-h']) == 0:
        usage('Error: missing hosts in -h and/or -f arguments')

    #If a single command was passed to us, implement a single command session
    if GV.argcmd:
            try:
                GV.session = ssh_utils.Session()
                GV.session.verbose=GV.opt['-v']
                GV.session.login(GV.opt['-h'], GV.USER)
                GV.session.echoCommand=GV.opt['-e']
                output=GV.session.executeCommand(GV.argcmd)
                GV.session.writeCommandOutput(output)
                if GV.session.verbose: print '[INFO] completed successfully'
                sys.stdout.flush()
            except ssh_utils.Session.SessionError, e:
                print 'Error: %s' % e
                pass

    else: #Otherwise, implement an interactive session
        interactive()
    
    return_code = sessionCleanup()
    sys.exit(return_code)

except KeyboardInterrupt:
    sessionCleanup()
    sys.exit('\nInterrupted...')
finally:
    coverage.stop()
    coverage.generate_report()