#!/usr/bin/python
#
# supermover -- wrapper around a number of data movement tools
# Copyright 2009, 2011, 2013 University of Tennessee
# Copyright 2015 Ohio Supercomputer Center
#
# License:  GNU GPL v2; see ../COPYING for details.
# Revision info:
# $HeadURL$
# $Revision$
# $Date$
#
# Usage: supermover [args] src_url dest_url
#        ***or***
#        supermover [args] -f file
#
# Arguments:
#        -c cfg -- use alternate config file; default is $SUPERMOVER_CFG,
#         or failing that, /usr/local/etc/supermover.cfg
#        -f urllist -- use file urllist as the list of src/dest urls
#
# Supported URL protocols:
#         file (default)
#         gsiftp
#         gsiscp
#         hpss
#         scp
#         rsync
#         rsync+ssh
#         rsync+gsissh
#         tar
#         tar+gzip
#         tar+bzip2
#
# Examples:
#         supermover /tmp/foo hpss:///home/user/foo
#         supermover hpss:///home/troy/bar gsiftp://gridftp.random.org/tmp/bar
#         supermover scp://login.random.org/tmp/baz file:///tmp/baz
#         supermover -f urllist
#
# URL list format:
# src_url1 dest_url1
# [src_url2 dest_url2]
# ...
#
# Config file format:
# src_proto:dest_proto = cmd [args]
#
# Macros for config file args:
# %u - src URL
# %U - dest URL
# %h - src host
# %H - dest host
# %p - src port
# %P - dest port
# %f - src path
# %F - dest path

from urlparse import urlparse
import getopt
import datetime
import os.path
import sys

# Globals
debug = False
successes = 0
failures = 0
proto_list = ["file"]


def get_methods(file):
    global debug
    global proto_list
    methods = {}
    if ( not os.path.exists(file) ):
        sys.stderr.write("Config file "+file+" does not exist.\n")
    elif ( not os.access(file,os.R_OK) ):
        sys.stderr.write("Config file "+file+" is not readable.\n")
    else:
        fd = open(file,"r")
        lines = fd.readlines()
        for line in lines:
            token = line.split("=",1)
            if ( token!=[] and len(token)==2 and not token[0].startswith("#") ):
                protopair = token[0].strip()
                method = token[1].strip()
                methods[protopair] = method
                protos = protopair.split(":")
                for proto in protos:
                    if ( proto not in proto_list ):
                        proto_list.append(proto)
        fd.close()
    return methods


def do_xfer(src,dest,methods,log):
    global debug
    global successes
    global failures
    global proto_list
    src_url = urlparse(src,"file")
    src_proto = src_url[0]
    if ( src_proto=="" ):
        src_proto = "file"
    src_hostport = src_url[1].split(":",2)
    src_host = src_hostport[0]
    if ( len(src_hostport)>1 ):
        src_port = src_hostport[1]
    else:
        src_port = None
    src_path = src_url[2]
    if ( (src_proto in proto_list) and
         src_path.startswith("//") and src_host=='' ):
        # if we get here, urlparse has screwed the pooch on
        # parsing the URL, so try some heuristics to recover
        new_hostpath = (src_path[2:]).split("/",2)
        src_hostport = new_hostpath[0].split(":",2)
        src_host = src_hostport[0]
        if ( len(src_hostport)>1 ):
            src_port = src_hostport[1]
        else:
            src_port = None
        src_path = src_url[2].split(src_host)[1]
    dest_url = urlparse(dest)
    dest_proto = dest_url[0]
    if ( dest_proto=="" ):
        dest_proto = "file"
    dest_hostport = dest_url[1].split(":",2)
    dest_host = dest_hostport[0]
    if ( len(dest_hostport)>1 ):
        dest_port = dest_hostport[1]
    else:
        dest_port = None
    dest_path = dest_url[2]
    if ( (dest_proto in proto_list) and
         dest_path.startswith("//") and dest_host=='' ):
        # if we get here, urlparse has screwed the pooch on
        # parsing the URL, so try some heuristics to recover
        new_hostpath = (dest_path[2:]).split("/",2)
        dest_hostport = new_hostpath[0].split(":",2)
        dest_host = dest_hostport[0]
        if ( len(dest_hostport)>1 ):
            dest_port = dest_hostport[1]
        else:
            dest_port = None
        dest_path = dest_url[2].split(dest_host)[1]
    if ( src_proto.startswith("gsi") or dest_proto.startswith("gsi") ):
        uid = os.getuid()
        if (
             # KRB5
             ( not os.path.exists("/tmp/krb5cc_"+str(os.getuid())) ) and
             ( "KRB5CCNAME" in os.environ and not os.path.exists(os.environ["KRB5CCNAME"]) ) and
             # KRB4
             ( not os.path.exists("/tmp/tkt"+str(os.getuid())) ) and
             # Globus/x.509
             ( not os.path.exists("/tmp/x509up_u"+str(os.getuid())) ) and
             ( not os.path.exists(os.path.expanduser("~")+"/.globus/x509up_u"+str(os.getuid())) ) and
             ( "X509_USER_PROXY" in os.environ and not os.path.exists(os.environ["X509_USER_PROXY"]) )
           ):
            failures = failures+1
            raise RuntimeException("GSI protocol requested but no GSI proxy credential found; please run myproxy-logon, grid-proxy-init, or kinit")
    method = methods[src_proto+":"+dest_proto]
    if ( method is not None ):
        cmd = method
        cmd = cmd.replace("%u",src)
        cmd = cmd.replace("%h",src_host)
        if ( src_port is not None ):
            cmd = cmd.replace("%p",src_port)
        cmd = cmd.replace("%f",src_path)
        cmd = cmd.replace("%U",dest)
        cmd = cmd.replace("%H",dest_host)
        if ( dest_port is not None ):
            cmd = cmd.replace("%P",dest_port)
        cmd = cmd.replace("%F",dest_path)
        log.write(str(datetime.datetime.now())+":  initiating transfer of "+src+" to "+dest+"\n")
        log.write(str(datetime.datetime.now())+":  "+cmd+"\n")
        log.flush()
        if ( debug is False ):
            pipe = os.popen(cmd+" 2>1")
            while True:
                line = pipe.readline()
                if ( not line ):
                    break
                else:
                    log.write(str(datetime.datetime.now())+":  "+line)
            returncode = pipe.close()
            if ( returncode is None or returncode==0 ):
                log.write(str(datetime.datetime.now())+":  completed transfer of "+src+" to "+dest+"\n")
                log.flush()
                successes += 1
            else:
                log.write(str(datetime.datetime.now())+":  failed transfer of "+src+" to "+dest+", return code "+str(returncode)+"\n")
                log.flush()
                failures += 1
    else:
        raise RuntimeException("No method defined for protocol pair "+src_proto+":"+dest_proto)


def usage(methods):
    sys.stderr.write("Usage:  supermover [args] srcURL destURL\n")
    sys.stderr.write("   or:  supermover [args] -f urlPairList\n\n")
    sys.stderr.write("Arguments:\n")
    sys.stderr.write(" -c cfgfile, --config=cfgfile\n\tRead config file \"cfgfile\" (default is "+cfgfile+")\n")
    sys.stderr.write(" -d, --debug\n\tEnable extra logging for debugging\n")
    sys.stderr.write(" -f urlPairList, --file=urlPairList\n\tRead src/edest URL pairs from file urlPairList rather than\n\tcommand line arguments\n")
    sys.stderr.write(" -h, --help\n\tPrint this help message\n")
    sys.stderr.write(" -l logfile, --log=logfile\n\tWrite log messages to file \"logfile\" (default is stdout)\n")
    sys.stderr.write("\nSupported protocols:\n\t")
    sys.stderr.write("\n\t".join(proto_list))
    sys.stderr.write("\n")
    sys.stderr.write("\nSupported transfers (src_proto:dest_prot):\n\t")
    sys.stderr.write("\n\t".join(methods.keys()))
    sys.stderr.write("\n\n")
    sys.exit(0)


cfgfile = "/usr/local/etc/supermover.cfg"
if ( os.path.exists("/etc/supermover.cfg") ):
    cfgfile = "/etc/supermover.cfg"
if ( "SUPERMOVER_CFG" in os.environ ):
    cfgfile = os.environ["SUPERMOVER_CFG"]
log = sys.stdout
urllist = None
show_usage = False

try:
    opts, args = getopt.getopt(sys.argv[1:],
                               "c:df:hl:",
                               ["config=s","debug","file=s","help","log=s"])
except getopt.GetoptError, err:
    sys.stderr.write(str(err)+"\n\n")
    methods = get_methods(cfgfile)
    usage(methods)
    sys.exit(-1)
    
if opts == [] and args == []:
    show_usage = True
    
if ( not show_usage ):
    for opt in opts:
        if ( opt[0]=="-h" or opt[0]=="--help" ):
            show_usage = True
        if ( opt[0]=="-d" or opt[0]=="--debug" ):
            debug = True
        if ( opt[0]=="-c" or opt[0]=="--config" ):
            cfgfile = opt[1]
        if ( opt[0]=="-f" or opt[0]=="--file" ):
            urllist = opt[1]
        if ( opt[0]=="-l" or opt[0]=="--log" ):
            logfile = opt[1]
            try:
                log = open(logfile,'w')
            except IOError, (errno, strerror):
                sys.stderr.write("Can't open "+logfile+" for writing:  "+strerror+" (errno="+str(errno)+")\n")
                pass          

methods = get_methods(cfgfile)
if ( debug is True ):
    for method in methods.keys():
        log.write("method["+method+"] = "+methods[method]+"\n")
if ( show_usage ):
    usage(methods)
    sys.exit(0)
log.flush()

if ( urllist is not None ):
    if ( not os.path.exists(urllist) ):
        raise IOError("URL list file "+urllist+" does not exist")
    elif ( not os.path.isfile(urllist) ):
        raise IOError("URL list \"file\" "+urllist+" is not actually a file")
    else:
        fd = open(urllist,"r")
        lines = fd.readlines()
        for line in lines:
            token = line.split()
            if ( token!=[] and len(token)==2 and not token[0].startswith("#") ):
                src = token[0]
                dest = token[1]
                do_xfer(src,dest,methods,log)
        fd.close()
else:
    if ( len(args)!=2 ):
        usage()
    else:
        src = args[0]
        dest = args[1]
        do_xfer(src,dest,methods,log)

log.write(str(datetime.datetime.now())+":  "+str(successes)+" successes, "+str(failures)+" failures\n")
if ( failures>0 ):
    sys.exit(-failures)
else:
    sys.exit(0)