gpMgmt/bin/gpssh

#!/usr/bin/env python3

'''
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 
          -s        : source gpdb environment while login
          cmd       : the command to execute. If not present, 
                      go into interactive mode
'''
import os
import sys

# this inserts the vendored pexpect from gpMgmt/bin/lib/pexpect/  into the path.
# Note that IDEs may show the local macos version of pexpect, which is NOT THE SAME
sys.path.insert(1, sys.path[0] + '/lib')

import getopt
import atexit
import signal
import pexpect
import time
import configparser
from gppylib.util import ssh_utils
from gppylib.gpparseopts import OptParser
from gppylib.commands.gp import get_coordinatordatadir


#
# 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
    opt['-s'] = False
    argcmd = None
    session = None
    DELAY_BEFORE_SEND = None
    PROMPT_VALIDATION_TIMEOUT = None
    SYNC_RETRIES = None

GV = __globals__()


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


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


def parseCommandLine():
    try:
        (options, args) = getopt.getopt(sys.argv[1:], '?evsh:f:D:u:d:t:', ['version'])
    except Exception as e:
        usage('Error: ' + str(e))
    for (switch, val) in options:
        if (switch == '-?'):
            usage(0)
        elif (switch[1] in 'evDs'):
            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
        elif (switch == '-d'):
            GV.DELAY_BEFORE_SEND = float(val)
        elif (switch == '-t'):
            GV.PROMPT_VALIDATION_TIMEOUT = float(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 parseConfigFile():
    if GV.DELAY_BEFORE_SEND is not None and GV.PROMPT_VALIDATION_TIMEOUT is not None:
        GV.SYNC_RETRIES = 3
        if GV.opt['-v']:
            print('Skip parsing gpssh.conf as both -d and -t are being used')
        return

    try:
        config = configparser.RawConfigParser()
        config.read(get_coordinatordatadir() + '/gpssh.conf')
        if GV.DELAY_BEFORE_SEND is None:
            GV.DELAY_BEFORE_SEND = config.getfloat('gpssh', 'delaybeforesend')
        if GV.PROMPT_VALIDATION_TIMEOUT is None:
            GV.PROMPT_VALIDATION_TIMEOUT = config.getfloat('gpssh', 'prompt_validation_timeout')
        if GV.SYNC_RETRIES is None:
            GV.SYNC_RETRIES = config.getint('gpssh', 'sync_retries')
    except Exception:
        if GV.opt['-v']:
            print('[WARN] Reference default values as $COORDINATOR_DATA_DIRECTORY/gpssh.conf could not be found')

        if GV.DELAY_BEFORE_SEND is None:
            GV.DELAY_BEFORE_SEND = 0.05
        if GV.PROMPT_VALIDATION_TIMEOUT is None:
            GV.PROMPT_VALIDATION_TIMEOUT = 1.0
        if GV.SYNC_RETRIES is None:
            GV.SYNC_RETRIES = 3

    if GV.DELAY_BEFORE_SEND < 0.0:
        print('[ERROR] delaybeforesend cannot be negative')
        sys.exit(1)

    if GV.PROMPT_VALIDATION_TIMEOUT <= 0.0:
        print('[ERROR] prompt_validation_timeout cannot be negative or zero')
        sys.exit(1)

    if GV.SYNC_RETRIES < 0:
        print('[ERROR] sync_retries cannot be negative')
        sys.exit(1)

    if GV.opt['-v']:
        print('Using delaybeforesend %s and prompt_validation_timeout %s' % (GV.DELAY_BEFORE_SEND,
                                                                             GV.PROMPT_VALIDATION_TIMEOUT))


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 permissions 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 as 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.DELAY_BEFORE_SEND, GV.PROMPT_VALIDATION_TIMEOUT, GV.SYNC_RETRIES)
                GV.session.echoCommand = GV.opt['-e']
            if GV.opt['-s']:
                GV.session.executeCommand("source {0}/cloudberry-env.sh".format(os.environ["GPHOME"]))
            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 as e:
            print('Error: %s' % e)
            pass
        except KeyboardInterrupt:
            print('\n[Interrupt...]')
            GV.session.reset()
            pass


def main():
    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.''' % os.path.split(sys.argv[0])[-1])

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

        # Parse gpssh.conf file
        parseConfigFile()

        # 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.DELAY_BEFORE_SEND, GV.PROMPT_VALIDATION_TIMEOUT, GV.SYNC_RETRIES)
                GV.session.echoCommand = GV.opt['-e']
                if GV.opt['-s']:
                    GV.session.executeCommand("source {0}/cloudberry-env.sh".format(os.environ["GPHOME"]))
                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 as 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...')


#############
if __name__ == '__main__':
    main()