#!/usr/bin/env python

# Copyright C 2002-2008 by Isilon Systems, Inc. All rights not expressly
# granted herein are reserved.
# Created in the United States of America.
# 
# This digital media is protected by U.S. and international copyright and
# intellectual property laws. Unless otherwise specified, all information
# and screens appearing as part of this digital medium, including
# software, services, documents, text, images, icons, and logos design;
# the selection, assembly, arrangement, and design thereof; and the code
# that enables its presentation, are the sole property of Isilon Systems,
# Inc.
# 
# Except as otherwise required by applicable law, no part of any
# copyrighted material may be reproduced, modified, translated,
# distributed, published, or transmitted by any means, electronic or
# mechanical, for any purpose, without the prior written consent of Isilon
# Systems. 
#  
# Isilon and Isilon Systems are registered trademarks of Isilon Systems,
# Inc.  Legion is a trademark of Isilon Systems, Inc.  All other
# trademarks and registered trademarks are the property of their
# respective holders.

"""
 $Id: isi_rdo 425595 2014-07-28 22:51:46Z rportman $
"""

import fcntl
import getopt
import os
import posix
import pty
import re
import select
import signal
import socket
import string
import sys
import tempfile

# "isi_rdo -t ssh" should remain portable to other platforms, so
# please refrain from importing isi libraries at the top level.

PROGNAME = os.path.basename(sys.argv[0])

if sys.platform.find('linux') >= 0:
    IOCTL_TIOCSCTTY = 21518
else:
    IOCTL_TIOCSCTTY = 536900705

try:
    True
except NameError:
    True = 1
    False = 0

try:
    DIRECTORY = os.getcwd()
except:
    DIRECTORY = ''
    
HOST = None
KNOWN_HOSTS_FILE = '/dev/null'
PASSWORD = 'a'
PRE_COMMANDS = []
PREFIX = ''
USER = 'root'

bufOut = ''
bufErr = ''

ssh_only_flags = '-k, -u, -p, and -q are only valid for ssh transport'

def usage(exitCode):
    sys.stderr.write('''
Usage: %s [OPTION]... [USER@]HOST[:DIR] [COMMAND]...

Options:
  -d, --directory DIR           Execute commands from DIR on remote machine.
                                Defaults to current working directory. An empty
                                DIR will result in commands being executed in
                                the users home directory on the remote machine.
  -h, --help                    Print this message and exit.
  -k, --known-hosts-file FILE   Use FILE for ssh known hosts file instead of
                                default /dev/null.
  -l, -u, --user USER           Login as USER instead of default root.
  -o, --prefix, PREFIX          Prefix every output line with PREFIX.
  -p, --pw, --password PASSWORD Use PASSWORD instead of default 'a'.
  -P, --pre-command COMMAND     Execute COMMAND before any other commands. This
                                is useful for setting up the environment. May
                                be specified multiple times. A COMMAND of '-'
                                resets the list of commands.
  -q, --password-file FILE      Use the password read from FILE instead of
                                default 'a'.
  -t, --transport TRANSPORT     Network transport. Default is rpc.
                                Valid transports: rpc, ssh

  Note: %s

Environment Variables:

  ISI_RDO_DIRECTORY             Specifies default DIRECTORY.
  ISI_RDO_KNOWN_HOSTS_FILE      Specifies default ssh known hosts FILE.
  ISI_RDO_USER                  Specifies default USER.
  ISI_RDO_PASSWORD              Specifies default PASSWORD.
  ISI_RDO_PASSWORD_FILE         Specifies default FILE to read password from.
                                This option is overridden by ISI_RDO_PASSWORD,
                                if both are set.
  ISI_RDO_PRE_COMMANDS          Specifies COMMANDS to run before command-line
                                COMMANDS. COMMANDS are separated by newlines.

Options specified later on the command line override environment variables as
well as options specified earlier on the command line. The exception to this
rule being pre-commands.

If no COMMANDS are given on the command line, they are read from standard input.

''' % (PROGNAME, ssh_only_flags))
    sys.exit(exitCode)

def error(str):
    print >> sys.stderr, '%s: %s' % (PROGNAME, str)
    sys.exit(1)

def passwordFromFile(filename):
    try:
        if filename[:8] == '/dev/fd/':
            password = os.fdopen(int(filename[8:])).readline()
        else:
            password = file(filename).readline()
    except IOError, e:
        error(str(e))
    index = string.find(password, '\n')
    if index >= 0:
        password = password[:index]
    return password

def waitForChild(childPid):
    # wait for the chid to exit, then exit
    while True:
        childPid, status = os.waitpid(childPid, 0)
        exitWith(status)

def exitWith(status):
    if os.WIFEXITED(status):
        sys.exit(os.WEXITSTATUS(status))
    elif os.WIFSIGNALED(status):
        sys.exit(1)

def main():
    global DIRECTORY, HOST, KNOWN_HOSTS_FILE, PASSWORD, PRE_COMMANDS, PREFIX, \
        USER

    signal.signal(signal.SIGINT, signal.SIG_DFL)

    TRANSPORT_FUNCS = {
        # ------------------------------------------------------------
        # Current advertised transports
        # ------------------------------------------------------------
        'ssh'      : runSSH,
        'rpc'      : runRPC,
        
        # ------------------------------------------------------------
        # Old or experimental transports
        # ------------------------------------------------------------
        #
        # I keep these around for the sake of comparison in case
        # we ever see issues with the two main methods. They're
        # unadvertised, so users will only know about them if
        # they go poking around in this file.
        
        # New, but slower ssh method. Updating this to work more
        # like the current rpc method should speed things up, and
        # would probably fix the current ssh method's binary-mangling
        # issues.
        'sshNew'   : runSSH_new, 
        # Old, environmented, binary-mangling method. See bug 39578.
        'rpcOld2'  : runRPC_old2,
        # Oldest, environment-less rpc method
        'rpcOld1'  : runRPC_old1, 
        }

    transport = None
    for var in os.environ.keys():
        if var == 'ISI_RDO_DIRECTORY':
            DIRECTORY = os.environ[var]
        if var == 'ISI_RDO_KNOWN_HOSTS_FILE':
            KNOWN_HOSTS_FILE = os.environ[var]
        if var == 'ISI_RDO_USER':
            USER = os.environ[var]
        if var == 'ISI_RDO_PASSWORD':
            PASSWORD = os.environ[var]
        if var == 'ISI_RDO_PASSWORD_FILE' and not \
                os.environ.has_key('ISI_RDO_PASSWORD'):
            PASSWORD = passwordFromFile(os.environ[var])
        if var == 'ISI_RDO_PRE_COMMANDS':
            PRE_COMMANDS.append(os.environ[var])
        if var == 'ISI_RDO_TRANSPORT':
            transport = os.environ[var]
            if not transport in TRANSPORT_FUNCS.keys():
                print >> sys.stderr, 'ERROR: Invalid transport in ' \
                      'ISI_RDO_TRANSPORT environment var: %s' % transport
                usage(1)

    try:
        opts, args = getopt.getopt(sys.argv[1:], 'd:hk:l:o:p:P:q:u:t:',
            ['directory=', 'help', 'known-hosts-file=', 'prefix=', 'password=',
             'pw=', 'password-file=', 'pre-command=', 'user=', 'transport='])
    except getopt.GetoptError:
        usage(1)

    if transport is None:
        transport = 'rpc'

    # In compliance mode, if we are running as root (via sudo), run the remote
    # command as compadmin on the other nodes unless the -u option is set.
    import isi.misc as misc
    complianceMode = misc.sysctlbyname('isi.compliance')
    if complianceMode and not os.geteuid():
        USER = 'compadmin'

    for o, a in opts:
        if o in ['-d', '--directory']:
            DIRECTORY = a
        if o in ['-h', '--help']:
            usage(0)
        if o in ['-k', '--known-hosts-file']:
            KNOWN_HOSTS_FILE = a
        if o in ['-l', '-u', '--user']:
            USER = a
        if o in ['-p', '--pw', '--password']:
            PASSWORD = a
        if o in ['-q', '--password-file']:
            PASSWORD = passwordFromFile(a)
        if o in ['-o', '--prefix']:
            PREFIX = a
        if o in ['-P', '--pre-command']:
            if a == '-':
                PRE_COMMANDS = []
            else:
                PRE_COMMANDS.append(a)
        if o in ['-t', '--transport']:
            transport = a
            if not transport in TRANSPORT_FUNCS.keys():
                print >> sys.stderr, 'ERROR: Invalid transport: %s' % transport
                usage(1)

    if len(args) < 1:
        usage(1)

    m = re.match('^(([^@]+)@)?([^:]+)(:(.*))?$', args[0])
    if m:
        if m.group(1):
            USER = m.group(2)
        HOST = m.group(3)
        if m.group(4):
            DIRECTORY = m.group(5)
    else:
        HOST = args[0];

    try:
        tempFile = tempfile.TemporaryFile('w+b', -1, 'isi_rdo')
        if DIRECTORY != '':
            print >> tempFile, "cd '%s'" % DIRECTORY
        for command in PRE_COMMANDS:
            # Bug 99153 and bug 99154: Prepend "sudo" to commands that require
            # sudo when run as compadmin in compliance mode. sudo -l determines
            # if compadmin can sudo a command on this node, and if so prepends
            # sudo to the command.
            if complianceMode and \
                    os.WEXITSTATUS(os.system('sudo -U %s -l %s >/dev/null' % \
                    (USER, command))) == 0:
                command = "sudo " + command

            print >> tempFile, command
        if len(args) > 1:
            # Bug 99153 and bug 99154: Prepend "sudo" to commands that require sudo
            # when run as compadmin in compliance mode (explanation above).
            command = string.join(args[1:])
            if complianceMode and \
                    os.WEXITSTATUS(os.system('sudo -U %s -l %s >/dev/null' % \
                    (USER, command))) == 0:
                command = "sudo " + command
            print >> tempFile, command
        else:
            while True:
                data = os.read(0, 1024)
                if not data:
                    break
                else:
                    tempFile.write(data)
        tempFile.flush()
        tempFile.seek(0)

        TRANSPORT_FUNCS[transport](tempFile)

    except (EnvironmentError, socket.error), e:
        error(PREFIX + str(e))
        
def runSSH(tempFile):
    class Fds: pass
    fds = Fds()
    fds.ptyMaster, fds.ptySlave = pty.openpty()
    fds.pipeReadOut, fds.pipeWriteOut = os.pipe()
    fds.pipeReadErr, fds.pipeWriteErr = os.pipe()
    fds.temp = tempFile.fileno()

    childPid = os.fork()
    if childPid:
        parent(fds)
        waitForChild(childPid) # doesn't return
        assert 0

    else:
        child(fds) # doesn't return
        assert 0

def runSSH_new(tempFile):
    import isi.app.lib.procs as procs

    cmds = tempFile.read()
    cmd = procs.SshThread(username=USER, host=HOST, prefix=PREFIX,
                          knownHostsFile=KNOWN_HOSTS_FILE)

    cmd.write(cmds)
    cmd.closeStdin()

    exitWith(cmd.error())

def runRPC_old1(tempFile):
    import isi.net.rpc as rpc
    
    cmds = tempFile.read()
    cmd = rpc.RPCDaemonProxy(HOST, 'isi.app.lib.procs.CmdThread', cmds,
                             iteratorMaxLines=1000)

    for line in cmd:
        print PREFIX + line.rstrip()

    exitWith(cmd.error())

def runRPC_old2(tempFile):
    import isi.net.rpc as rpc

    cmds = tempFile.read()
    cmd = rpc.RPCDaemonProxy(HOST, 'isi.app.lib.procs.ShellThread',
                             username=USER)

    cmd.write(cmds)
    cmd.closeStdin()

    for line in cmd:
        print PREFIX + line.rstrip()

    exitWith(cmd.error())

def rpcSelectLoop(tempFile, stdin, stdout, stderr, bsize):
    global bufErr, bufOut

    for sock in [stdin, stdout, stderr]:
        sock._blocking = False
    
    class SelectOps(object):
        def __init__(self, onRead, onEmpty):
            def noop(arg=None):
                pass
            
            self.onRead = onRead or noop
            self.onEmpty = onEmpty or noop

    class Prefixer(object):
        def __init__(self, onRead):
            self.onRead = onRead

            # Always prefix the first line
            self.needPrefix = True

            # We'll want to automatically place a prefix after
            # any newline that has at least one character following it.
            self.regex = re.compile(r'\n(?=.)', re.DOTALL)

        def __call__(self, data):
            # If the last character in this chunk is a newline, we'll
            # want to add a prefix to the front of the next chunk, if
            # there is one
            willNeedPrefix = data[-1] == '\n'
 
            # Add a prefix following all non-terminating newlines
            data = self.regex.sub('\n'+PREFIX, data)

            # If this is the first line, or the final character of
            # the last data chunk was a newline, we need to add
            # a prefix to the front of this chunk.
            if self.needPrefix:
                data = PREFIX + data
            self.needPrefix = willNeedPrefix

            # Pass our prefixed data on to the stream handler
            self.onRead(data)

    def getPrefixer(onRead):
        # If there's no PREFIX to prepend to each line,
        # then don't bother using a Prefixer object, and
        # just write the data straight to the output stream.
        return PREFIX and Prefixer(onRead) or onRead

    # Make sure we always flush stdout after we write any
    # data to it. This doesn't seem necessary when running
    # isi_rdo from a shell, but when invoked via the
    # isi.app.lib.procs.{get,proc}_for_array_output()
    # functions, partial output will be buffered up and
    # not returned to the user immediately unless you
    # flush manually.
    # Also, per bug 47529, flush the output after every single \n.
    # This prevents us outputting partial line data as long as your lines
    # are less than 512B (PIPE_BUF per sys_pipe.c) in cases where the
    # recipient of stdout gets backed up (ie slow serial or stuck process).
    def outWrite(x):
        global bufOut
        bufOut += x
        while True:
            lastn = bufOut.find('\n')
            if lastn < 0:
                break
            sys.stdout.write(bufOut[0:lastn+1])
            sys.stdout.flush()
            bufOut = bufOut[lastn+1:]

    # Ditto for stderr
    def errWrite(x):
        global bufErr
        bufErr += x
        while True:
            lastn = bufErr.find('\n')
            if lastn < 0:
                break
            sys.stderr.write(bufErr[0:lastn+1])
            sys.stderr.flush()
            bufErr = bufErr[lastn+1:]

    selectees = {
        tempFile.fileno() : SelectOps(stdin.send, stdin.close),
        stdout.fileno()   : SelectOps(getPrefixer(outWrite), None),
        stderr.fileno()   : SelectOps(getPrefixer(errWrite), None),
        }

    while selectees:
        ready = select.select(selectees.keys(), [], [])[0]

        for fileno in ready:
            try:
                data = os.read(fileno, bsize)
            except EnvironmentError, e: # linux returns EIO instead of EOF
                data = ''
            if data:
                selectees[fileno].onRead(data)
            else:
                selectees[fileno].onEmpty()
                del selectees[fileno]

    if len(bufOut) > 0:
        sys.stdout.write(bufOut)
        bufOut = ''
    if len(bufErr) > 0:
        sys.stderr.write(bufErr)
        bufErr = ''

def runRPC(tempFile):
    import isi.net.rpc as rpc
    import isi.net.transport as transport

    # Create a proxy to a remote TcpShellThread object on the server.
    # We'll tell the server to pass in it's original listening IP address
    # for the TcpShellThread's listenIP argument using an
    # RPCArgumentSubstitution object. That way, if the original server
    # is listening promiscuously, so will this one.
    argSub = rpc.RPCArgumentSubstitution('original-server-address')
    cmd = rpc.RPCDaemonProxy(HOST, 'isi.app.lib.procs.TcpShellThread',
                             username=USER, listenIP=argSub)

    # Tell the thread to start accepting connections. It will return to us
    # the ip and port on which it is listening.
    (ip, port) = cmd.listen()

    # Create three new connections to this server, one each for
    # stdin, stdout, and stderr, in that order
    def makeFileStreamConnection(tip, tport):
        tp = None
        while not tp:
            try:
                tp = transport.TCPTransport(tip, tport)
                tp.connect()
            except transport.TransportPortCollisionException:
                print >> sys.stderr, 'Warning: Detected port collision on port %d, retrying...' % tport
                # The other end never got the connection due to a port
                # collision. Close it and try again.
                tp.close()
                tp = None
        return tp

    stdin = makeFileStreamConnection(ip, port)
    stdout = makeFileStreamConnection(ip, port)
    stderr = makeFileStreamConnection(ip, port)

    # Tell the command thread to fire things up
    cmd.start()

    # Marshall data between us and the command until it terminates
    # For network-bound traffic, using a block size of 8KB appears 
    # to be very marginally faster than using a block size of 1KB,
    # but appears to takes almost 1/3 as much CPU.
    rpcSelectLoop(tempFile, stdin, stdout, stderr, 8192)

    # Exit with the error code from the remote command
    exitWith(cmd.error())

def parent(fds): # parent proxies input and output
    os.close(fds.pipeWriteOut)
    os.close(fds.pipeWriteErr)
    os.close(fds.ptySlave)
    os.close(fds.temp)

    header = True
    pipeRemainderOut = ''
    pipeRemainderErr = ''
    ptyRemainder = ''
    forRead = [fds.ptyMaster, fds.pipeReadOut, fds.pipeReadErr]

    try:
        while forRead:
            ready = select.select(forRead, [], [])[0]
            if fds.ptyMaster in ready:
                try:
                    data = os.read(fds.ptyMaster, 1024)
                except EnvironmentError, e: # linux returns EIO instead of EOF
                    data = ''
                if not data:
                    forRead.remove(fds.ptyMaster)
                    continue
                data = ptyRemainder + data

                while True:
                    index = string.find(data, '\n')
                    if index < 0:
                        break
                    line = data[0:index + 1]
                    data = data[index+1:]

                if string.find(data, 'assword:') >= 0:
                    os.write(fds.ptyMaster, '%s\n' % PASSWORD)
                    data = ''

                if string.find(data, 'Enter passphrase for') == 0:
                    os.write(fds.ptyMaster, '\n')
                    data = ''

                if string.find(data, 'Response:') == 0:
                    os.write(fds.ptyMaster, '\n')
                    data = ''

                if string.find(data, 'Are you sure you want to continue') == 0:
                    os.write(fds.ptyMaster, 'yes\n')
                    data = ''

                ptyRemainder = data

            if fds.pipeReadErr in ready:
                data = os.read(fds.pipeReadErr, 1024)
                if not data:
                    forRead.remove(fds.pipeReadErr)
                    continue
                data = pipeRemainderErr + data

                while True:
                    index = string.find(data, '\n')
                    if index < 0:
                        break
                    line = data[0:index + 1]
                    data = data[index+1:]
                    if header:
                        if string.find(line, 'socket: Protocol not supported') == 0:
                            continue
                        if string.find(line, 'Pseudo-terminal will not be allocated') == 0:
                            continue
                        if string.find(line, 'The authenticity of host') == 0:
                            continue
                        if string.find(line, 'RSA key fingerprint is') == 0:
                            continue
                        if string.find(line, 'DSA key fingerprint is') == 0:
                            continue
                        if string.find(line, 'Warning: Permanently added') == 0:
                            continue
                        if string.find(line, 'Warning: Remote host denied X11 forwarding') == 0:
                            continue
                        if string.find(line, 'Permission denied') == 0:
                            print PREFIX + 'Permission denied.'
                            return
                        else:
                            header = False
                    sys.stderr.write(PREFIX + line);
                    sys.stderr.flush()

                pipeRemainderErr = data

            if fds.pipeReadOut in ready:
                data = os.read(fds.pipeReadOut, 1024)
                if not data:
                    forRead.remove(fds.pipeReadOut)
                    continue
                data = pipeRemainderOut + data

                while True:
                    index = string.find(data, '\n')
                    if index < 0:
                        break
                    line = data[0:index + 1]
                    data = data[index+1:]
                    sys.stdout.write(PREFIX + line);
                    sys.stdout.flush()

                if string.find(data, 'Are you sure you want to continue') == 0:
                    os.write(fds.ptyMaster, 'yes\n')
                    data = ''

                pipeRemainderOut = data

    finally:
        os.close(fds.pipeReadOut)
        os.close(fds.pipeReadErr)
        os.close(fds.ptyMaster)

def child(fds): # child exec's ssh
    os.close(fds.pipeReadOut)
    os.close(fds.pipeReadErr)
    os.close(fds.ptyMaster)

    os.dup2(fds.temp, 0)
    if fds.temp != 0:
        os.close(fds.temp)

    os.dup2(fds.pipeWriteOut, 1)
    os.dup2(fds.pipeWriteErr, 2)
    if fds.pipeWriteOut != 1:
        os.close(fds.pipeWriteOut)
    if fds.pipeWriteErr != 2:
        os.close(fds.pipeWriteErr)

    posix.setsid()
    fcntl.ioctl(fds.ptySlave, IOCTL_TIOCSCTTY, 0)

    try:
        os.execlp('ssh', 'ssh', '-o',
            'UserKnownHostsFile=%s' % KNOWN_HOSTS_FILE, '%s@%s' % (USER, HOST))
    except EnvironmentError, e:
        error('running ssh: %s' % str(e))

if __name__ == '__main__':
    main()
