#!/usr/bin/env python

# based on isi_rdo

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

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

KNOWN_HOSTS_FILE = '/dev/null'
PASSWORD = 'a'
OPTIONS = []
PREFIX = ''
USER = 'root'

def usage(exitCode):
    sys.stderr.write('''
Usage: %s [OPTION]... [[USER@]HOST1:]FILE1 [[USER@]HOST2:]FILE2

Options:
  -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.
  -p, --pw, --password PASSWORD Use PASSWORD instead of default 'a'.
  -q, --password-file FILE      Use the password read from FILE instead of
                                default 'a'.
  -o, --option OPTIONS          Pass OPTIONS to scp (quotes optional).
                                Multiple instances allowed.
                                E.g., -o pr,v ==> "scp -pr -v"

Environment Variables:

  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.

Options specified later on the command line override environment variables as
well as options specified earlier on the command line.

''' % PROGNAME)
    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 getHostname():
    try:
        f = os.popen("/bin/hostname", "r")
        host_name = f.readline()
    except OSError, (errno, strerror):
        error("error running hostname: %s" % os.strerror(errno))
    f.close()
    host_name = host_name.strip()
    return host_name

def waitForChild(childPid):
    # wait for the chid to exit, then exit
    while True:
        childPid, status = os.waitpid(childPid, 0)
        if os.WIFEXITED(status):
            sys.exit(os.WEXITSTATUS(status))
        elif os.WIFSIGNALED(status):
            sys.exit(1)

def main():
    global KNOWN_HOSTS_FILE, PASSWORD, USER, OPTIONS

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

    for var in os.environ.keys():
	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])

    try:
        opts, args = getopt.getopt(sys.argv[1:], 'hk:l:o:p:q:u:',
            ['help', 'known-hosts-file=', 'option=', 'password=',
             'pw=', 'password-file=', 'user='])
    except getopt.GetoptError:
        usage(1)

    for o, a in opts:
        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', '--option']:
            OPTIONS.extend(a.split(','))

    if len(args) != 2:
        usage(1)

    HOST = getHostname()
    class Scp: pass
    scp = Scp()
    scp.host_self = HOST
    scp.options = OPTIONS
    scp.user = [ USER, USER ]
    scp.host = [ HOST, HOST ]
    scp.file = args
    for i in range(2):
        m = re.match('^(([^@]+)@)?(([^:/]+):)?(.*)$', scp.file[i])
        if m:
            if m.group(1):
                scp.user[i] = m.group(2)
            if m.group(3):
                scp.host[i] = m.group(4)
                scp.file[i] = m.group(5)
        else:
            usage(1)

    # Fix options: if any need leading '-', add it
    scp.options = map(lambda x:x.strip()[0]=='-' and x or '-%s' % x, scp.options)

    try:

        class Fds: pass
        fds = Fds()
        fds.ptyMaster, fds.ptySlave = pty.openpty()
        fds.pipeRead, fds.pipeWrite = os.pipe()

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

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

    except EnvironmentError, e:
        error(str(e))

def parent(fds, scp): # parent proxies input and output
    os.close(fds.pipeWrite)
    os.close(fds.ptySlave)

    header = True
    pipeRemainder = ''
    ptyRemainder = ''
    forRead = [fds.ptyMaster, fds.pipeRead]

    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.pipeRead in ready:
                data = os.read(fds.pipeRead, 1024)
                if not data:
                    forRead.remove(fds.pipeRead)
                    continue
                data = pipeRemainder + 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
                    print PREFIX + line,
                sys.stdout.flush()

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

                pipeRemainder = data

    finally:
        os.close(fds.pipeRead)
        os.close(fds.ptyMaster)

def child(fds, scp): # child exec's ssh
    os.close(fds.pipeRead)
    os.close(fds.ptyMaster)

    os.dup2(fds.pipeWrite, 1)
    os.dup2(fds.pipeWrite, 2)
    if fds.pipeWrite not in [1, 2]:
        os.close(fds.pipeWrite)

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

    try:
        files = {}
        for f in xrange(2):
            files[f] = '%s@%s:%s' % (scp.user[f], scp.host[f], scp.file[f])
            if scp.host[f] == scp.host_self:
                files[f] = '%s' % (scp.file[f])
            elif scp.user[f] == None:
                files[f] = '%s:%s' % (scp.host[f], scp.file[f])

        # check if we're trying to cp to ourselves (same file), avoid error:
        # cp: FILE1 and FILE2 are identical (not copied).
        if files[0].strip() == files[1].strip():
            sys.exit(0)

        args = ['scp']
        if scp.options:
            args.extend(scp.options)
        args.extend(['-o', 'UserKnownHostsFile=%s' % KNOWN_HOSTS_FILE,
                     files[0], files[1]])
        os.execvp('scp', args)
    except EnvironmentError, e:
        error('running scp: %s' % str(e))

if __name__ == '__main__':
    main()
