#!/usr/bin/env python

"""
Pegasus utility for transfer of files during workflow enactment

Usage: pegasus-transfer [options]
"""

##
#  Copyright 2007-2013 University Of Southern California
#
#  Licensed 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.
##

import errno
import logging
import math
import optparse
import os
import Queue
import random
import re
import signal
import stat
import string
import subprocess
import sys
import tempfile
import threading
import time
import cmd
import ConfigParser

try:
    # Python 3.0 and later
    import urllib.parse as urllib
except ImportError:
    # Fall back to Python 2's urllib
    import urllib as urllib


__author__ = "Mats Rynge <rynge@isi.edu>"


# --- regular expressions -----------------------------------------------------

re_parse_comment = re.compile(r'^# +[\w]+ +[\w]+ +([\w\-_]+)')
re_parse_url = re.compile(r'([\w]+)://([\w\.\-:@]*)(/[\S ]*)?')


# --- classes -----------------------------------------------------------------

class Transfer:
    """
    Represents a single transfer request.
    """

    pair_id        = 0          # the id of the pair in the input (nth pair)
    src_sitename   = "UNKNOWN"  # name from the Pegasus site catalog
    src_proto      = ""         # 
    src_host       = ""         # 
    src_path       = ""         #
    dst_sitename   = "UNKNOWN"  # name from the Pegasus site catalog
    dst_proto      = ""         #
    dst_host       = ""         #
    dst_path       = ""         #
    allow_grouping = True       # can this transfer be grouped with others?

    def __init__(self, pair_id):
        """
        Initializes the transfer class
        """
        self.pair_id = pair_id
        
    def set_src_sitename(self, sitename):
        # the site name is used to match against shell variables, so we have
        # have to replace dashes with underscores (as we do in the planner)
        self.src_sitename = string.replace(sitename, "-", "_")

    def set_dst_sitename(self, sitename):
        # the site name is used to match against shell variables, so we have
        # have to replace dashes with underscores (as we do in the planner)
        self.dst_sitename = string.replace(sitename, "-", "_")

    def set_src(self, url):
        self.src_proto, self.src_host, self.src_path = self.parse_url(url)
    
    def set_dst(self, url):
        self.dst_proto, self.dst_host, self.dst_path = self.parse_url(url)

    def parse_url(self, url):
        proto = ""
        host = ""
        path = ""
            
        # the url should be URL decoded to work with our shell callouts
        url = urllib.unquote(url)

        # default protocol is file://
        if string.find(url, ":") == -1:
            logger.debug("URL without protocol (" + url + 
                        ") - assuming file://")
            url = "file://" + url

        # file url is a special cases as it can contain relative paths and
        # env vars
        if string.find(url, "file:") == 0:
            proto = "file"
            # file urls can either start with file://[\w]*/ or file: (no //)
            path = re.sub("^file:(//)?", "", url)
            path = expand_env_vars(path)
            return proto, host, path
        
        # symlink url is a special cases as it can contain relative paths and
        # env vars
        if string.find(url, "symlink:") == 0:
            proto = "symlink"
            # symlink urls can either start with symlink://[\w]*/ or
            # symlink: (no //)
            path = re.sub("^symlink:(//)?", "", url)
            path = expand_env_vars(path)
            return proto, host, path

        # other than file/symlink urls
        r = re_parse_url.search(url)
        if not r:
            raise RuntimeError("Unable to parse URL: %s" % (url))
        
        # Parse successful
        proto = r.group(1)
        host = r.group(2)
        path = r.group(3)
        
        if path is None:
            path = ""
        
        # no double slashes in urls
        path = re.sub('//+', '/', path)
        
        return proto, host, path

    def src_url(self):
        return "%s://%s%s" % (self.src_proto, self.src_host, self.src_path)
    
    def src_url_encoded(self):
        return "%s://%s%s" % (self.src_proto, self.src_host, urllib.quote(self.src_path))

    def src_url_srm(self):
        """
        srm-copy is using broken urls - wants an extra / 
        """
        if self.src_proto != "srm":
            return "%s://%s/%s" % (self.src_proto, self.src_host, self.src_path)
        return self.src_url()

    def dst_url(self):
        return "%s://%s%s" % (self.dst_proto, self.dst_host, self.dst_path)
    
    def dst_url_encoded(self):
        return "%s://%s%s" % (self.dst_proto, self.dst_host, urllib.quote(self.dst_path))

    def dst_url_srm(self):
        """
        srm-copy is using broken urls - wants an extra / 
        """
        if self.dst_proto != "srm":
            return "%s://%s/%s" % (self.dst_proto, self.dst_host, self.dst_path)
        return self.dst_url()
    
    def dst_url_dirname(self):
        dn = os.path.dirname(self.dst_path)
        return "%s://%s%s" % (self.dst_proto, self.dst_host, dn)

    def groupable(self):
        """
        currently only gridftp allows for grouping
        """
        return self.allow_grouping and \
               (self.src_proto == "gsiftp" or self.dst_proto == "gsiftp")

    def __cmp__(self, other):
        """
        compares first on protos, then on hosts, then on paths - useful
        for grouping similar types of transfers
        """
        if cmp(self.src_proto, other.src_proto) != 0:
            return cmp(self.src_proto, other.src_proto)
        if cmp(self.dst_proto, other.dst_proto) != 0:
            return cmp(self.dst_proto, other.dst_proto)
        if cmp(self.src_host, other.src_host) != 0:
            return cmp(self.src_host, other.src_host)
        if cmp(self.dst_host, other.dst_host) != 0:
            return cmp(self.dst_host, other.dst_host)
        if cmp(self.src_path, other.src_path) != 0:
            return cmp(self.src_path, other.src_path)
        if cmp(self.dst_path, other.dst_path) != 0:
            return cmp(self.dst_path, other.dst_path)
        return 0


class Singleton(type):
    """Implementation of the singleton pattern"""
    _instances = {}
    def __call__(cls, *args, **kwargs):
        if cls not in cls._instances:
            cls._instances[cls] = \
                super(Singleton, cls).__call__(*args, **kwargs)
            cls.lock = threading.Lock()
        return cls._instances[cls]


class Tools(object):
    """Singleton for detecting and maintaining tools we depend on
    """
    
    __metaclass__ = Singleton
    
    _info = {}

    def find(self, executable, version_arg=None, version_regex=None):

        self.lock.acquire()
        try:
            if executable in self._info:
                if self._info[executable] is None:
                    return None
                return self._info[executable]
            
            logger.debug("Trying to detect availability/location of tool: %s"
                         %(executable))
            
            # initialize the global tool info for this executable
            self._info[executable] = {}
            self._info[executable]['full_path'] = None
            self._info[executable]['version'] = None
            self._info[executable]['version_major'] = None
            self._info[executable]['version_minor'] = None
            self._info[executable]['version_patch'] = None
        
            # figure out the full path to the executable
            full_path = backticks("which " + executable + " 2>/dev/null") 
            full_path = full_path.rstrip('\n')
            if full_path == "":
                logger.info("Command '%s' not found in the current environment"
                            %(executable))
                self._info[executable] = None
                return self._info[executable]
            self._info[executable]['full_path'] = full_path
        
            # version
            if version_regex is None:
                version = "N/A"
            else:
                version = backticks(executable + " " + version_arg + " 2>&1")
                version = version.replace('\n', "")
                re_version = re.compile(version_regex)
                result = re_version.search(version)
                if result:
                    version = result.group(1)
                self._info[executable]['version'] = version
        
            # if possible, break up version into major, minor, patch
            re_version = re.compile("([0-9]+)\.([0-9]+)(\.([0-9]+)){0,1}")
            result = re_version.search(version)
            if result:
                self._info[executable]['version_major'] = int(result.group(1))
                self._info[executable]['version_minor'] = int(result.group(2))
                self._info[executable]['version_patch'] = result.group(4)
            if self._info[executable]['version_patch'] is None or \
               self._info[executable]['version_patch'] == "":
                self._info[executable]['version_patch'] = None
            else:
                self._info[executable]['version_patch'] = \
                    int(self._info[executable]['version_patch'])
        
            logger.info("Tool found: %s   Version: %s   Path: %s" 
                        % (executable, version, full_path))
            return self._info[executable]['full_path']
        finally:
            self.lock.release()


    def full_path(self, executable):
        """ Returns the full path to a given executable """
        self.lock.acquire()
        try:
            if executable in self._info and self._info[executable] is not None:
                return self._info[executable]['full_path']
            return None
        finally:
            self.lock.release()


    def major_version(self, executable):
        """ Returns the detected major version given executable """
        self.lock.acquire()
        try:
            if executable in self._info and self._info[executable] is not None:
                return self._info[executable]['version_major']
            return None
        finally:
            self.lock.release()
                

class TimedCommand(object):
    """ Provides a shell callout with a timeout """
        
    def __init__(self, cmd, timeout_secs = 5*60*60, should_log = False):
        self.cmd = cmd
        self.timeout_secs = timeout_secs
        self.should_log = should_log
        self.process = None

    def run(self):
        def target():
            self.process = subprocess.Popen(self.cmd, shell=True, preexec_fn=os.setpgrp)
            self.process.communicate()

        if self.should_log or logger.isEnabledFor(logging.DEBUG):
            logger.info(self.cmd)
            
        sys.stdout.flush()
        
        thread = threading.Thread(target=target)
        thread.start()

        thread.join(self.timeout_secs)
        if thread.isAlive():
            # do our best to kill the whole process group
            try:
                # os.killpg did not work in all environments
                #os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
                kill_cmd = "kill -TERM -%d" %(os.getpgid(self.process.pid))
                kp = subprocess.Popen(kill_cmd, shell=True)
                kp.communicate()
                self.process.terminate()
            except:
                pass
            thread.join()
            raise RuntimeError("Command timed out after %d seconds: %s" %(self.timeout_secs, self.cmd))
        if self.process.returncode != 0:
            raise RuntimeError("Command exited with non-zero exit code (%d): %s" \
                               %(self.process.returncode, self.cmd))


class TransferHandlerBase(object):
    """
    Base class for all transfer handlers. Derived classes should set the 
    protocol map (for example ["http->file"]) and implement the do_transfer()
    method.
    """
    
    _name = "BaseHandler"
    _protocol_map = []
    
    lock = threading.Lock()
    
    def do_transfer(self, transfer, attempt):
        """
        Handles single transfer - all derived classes should override this
        method
        """
        raise RuntimeError("do_transfer() is not implemented in " + self._name)
    
    def do_multi_transfer(self, transfers, attempt):
        """
        Handles transfers for a list of similar transfers. This is used mostly
        for efficiency when using protocols such as gsiftp. Implementing this
        method is optional for protocols
        """
        raise RuntimeError("do_multi_transfer() is not implemented in " +
                           self._name)
    
    def protocol_map_check(self, src_proto, dst_proto):
        """
        Checks to see if a src/dst protocol pair can be handled by the handler.
        This is the base for the automatic handler detection in the TransferSe
        class.
        """
        item = str(src_proto) + "->" + str(dst_proto)
        return (item in self._protocol_map)


class CpHandler(TransferHandlerBase):
    """
    Uses the system cp command to copy local file to file
    """
    
    _name = "CpHandler"
    _protocol_map = ["file->file"]

    def do_transfer(self, transfer, attempt):
        prepare_local_dir(os.path.dirname(transfer.dst_path))
        if os.path.exists(transfer.src_path) \
           and os.path.exists(transfer.dst_path):
            # make sure src and target are not the same file - have to
            # compare at the inode level as paths can differ
            src_inode = os.stat(transfer.src_path)[stat.ST_INO]
            dst_inode = os.stat(transfer.dst_path)[stat.ST_INO]
            if src_inode == dst_inode:
                logger.warning("cp: src (%s) and dst (%s) already exists"
                               % (transfer.src_path, transfer.dst_path))
                return True
        cmd = "/bin/cp -f -R -L '%s' '%s'" \
            % (transfer.src_path, transfer.dst_path)
        try:
            tc = TimedCommand(cmd)
            tc.run()
        except RuntimeError, err:
            logger.error(err)
            return False
        stats_add(transfer.dst_path)
        return True


class FDTHandler(TransferHandlerBase):
    """
    Handler for FDT (Fast Data Transfer - http://monalisa.cern.ch/FDT/)
    """

    _name = "FDTHandler"
    _protocol_map = ["file->fdt", "fdt->file"]
        
    def do_transfer(self, transfer, attempt):

        # download fdt.jar on demand - it can not be shipped with Pegasus due
        # to licensing
        if not os.path.exists("fdt.jar"):
            cmd = "wget -nv -O fdt.jar http://monalisa.cern.ch/FDT/lib/fdt.jar"
            try:
                tc = TimedCommand(cmd, timeout_secs = 10*60)
                tc.run()
            except RuntimeError, err:
                logger.error(err)
                
        cmd = "echo | java -jar fdt.jar"
        if transfer.dst_proto == "file":
            prepare_local_dir(os.path.dirname(transfer.dst_path))
            cmd += " " + transfer.src_host + ":" + transfer.src_path
            cmd += " " + transfer.dst_path
        else:
            cmd += " " + transfer.src_path
            cmd += " " + transfer.dst_host + ":" + transfer.dst_path
            stats_add(transfer.src_path)
        try:
            tc = TimedCommand(cmd)
            tc.run()
            if transfer.dst_proto == "file":
                stats_add(transfer.dst_path)    
        except RuntimeError, err:
            logger.error(err)
            return False
        return True
    


class GridFtpHandler(TransferHandlerBase):
    """
    Transfers to/from and between GridFTP servers
    """
    
    _name = "GridFtpHandler"
    _protocol_map = [
                    "file->gsiftp",
                    "gsiftp->file", 
                    "gsiftp->gsiftp",
                    "ftp->ftp",
                    "ftp->gsiftp",
                    "gsiftp->ftp",
                    "http->gsiftp",
                    "file->sshftp",
                    "sshftp->file",
                    "sshftp->sshftp"
                    ]

    def do_transfer(self, transfer, attempt):
        """
        gsiftp - use globus-url-copy for transfers
        """
        
        tools = Tools()
        if tools.find("globus-url-copy", "-version", "([0-9]+\.[0-9]+)") \
           is None:
            raise RuntimeError("Unable to do gsiftp transfers because" +
                               " globus-url-copy could not be found")
            
        third_party = (transfer.src_proto == "gsiftp" \
                       and transfer.dst_proto == "gsiftp") or \
                      (transfer.src_proto == "sshftp" \
                       and transfer.dst_proto == "sshftp")

        success = self._exec_transfers([transfer], attempt,  True, third_party)
        return success
    
    
    def do_multi_transfer(self, full_list, attempt, failed_q):
        """
        gsiftp - globus-url-copy for now, maybe uberftp in the future
        """
        if len(full_list) == 0:
            return
        
        tools = Tools()
        if tools.find("globus-url-copy", "-version", "([0-9]+\.[0-9]+)") \
           is None:
            raise RuntimeError("Unable to do gsiftp transfers because" +
                               " globus-url-copy could not be found")
        
        # create lists with similar (same src host/path, same dst host/path)
        # url pairs
        while len(full_list) > 0:
    
            similar_list = []
    
            curr = full_list.pop()
            prev = curr
            third_party = (curr.src_proto == "gsiftp" \
                           and curr.dst_proto == "gsiftp") or \
                          (curr.src_proto == "sshftp" \
                           and curr.dst_proto == "sshftp")
    
            while self._check_similar(curr, prev):
                
                similar_list.append(curr)
    
                if len(full_list) == 0:
                    break
                else:
                    prev = curr
                    curr = full_list.pop()
    
            if not self._check_similar(curr, prev):
                # the last pair is not part of the set and needs to be added
                # back to the beginning of the list
                full_list.append(curr)
    
            if len(similar_list) == 0:
                break
    
            # we now have a list of similar transfers - break up and send the
            # first one with create dir and the rest with no create dir options
            first_list = []
            first_list.append(similar_list.pop())
            mkdir_done = self._exec_transfers(first_list, attempt, 
                                              True, third_party)
    
            # first attempt get some extra tries - this is to drill down on
            # guc options
            if attempt == 1 and not mkdir_done:
                mkdir_done = self._exec_transfers(first_list, attempt, 
                                                  True, third_party)
                if not mkdir_done:
                    mkdir_done = self._exec_transfers(first_list,
                                                      attempt, True, 
                                                      third_party)
    
            if mkdir_done:
                # run the rest of the group - but limit the number of entries
                # for each pipeline
                chunks = self._split_similar(similar_list)
                for l in chunks:
                    if not self._exec_transfers(l, attempt, False, third_party):
                        for i, t in enumerate(l):
                            failed_q.put(t)
            else:
                # mkdir job failed - all subsequent jobs will fail
                failed_q.put(first_list[0])
                for i, t in enumerate(similar_list):
                    failed_q.put(t)
   
    
    def _exec_transfers(self, transfers, attempt, create_dest, third_party):
        """
        sub to gsiftp() - transfers a list of urls
        """
        global gsiftp_failures
        
        delayed_file_stat = []
    
        # create tmp file with transfer src/dst pairs
        num_pairs = 0
        try:
            tmp_fd, tmp_name = tempfile.mkstemp(prefix="pegasus-transfer-",
                                                suffix=".lst")
            tmp_file = os.fdopen(tmp_fd, "w+b")
        except:
            raise RuntimeError("Unable to create tmp file for"
                               + " globus-url-copy transfers")
            
        for i, t in enumerate(transfers):
            num_pairs += 1
            logger.debug("   adding %s %s" % (t.src_url_encoded(),
                                              t.dst_url_encoded()))
    
            # delay stating until we have finished the transfers
            if t.src_proto == "file":
                delayed_file_stat.append(t.src_path)
            elif t.dst_proto == "file":
                delayed_file_stat.append(t.dst_path)
    
            tmp_file.write("%s %s\n" % (t.src_url_encoded(),
                                        t.dst_url_encoded()))
    
        tmp_file.close()
       
        if num_pairs > 1:
            logger.info("Grouped %d similar gsiftp transfers together in"
                        " temporary file %s" %(num_pairs, tmp_name))
    
        transfer_success = False

        # for transfer to file://, run a normal mkdir command instead of
        # having g-u-c use -create-dest
        if transfers[0].dst_proto == "file":
            prepare_local_dir(os.path.dirname(transfers[0].dst_path))
            # override the create_dest flag as the directory now exists
            create_dest = False

        # build command line for globus-url-copy
        tools = Tools()
        cmd = tools.full_path('globus-url-copy')
        
        # credential handling
        cmd_creds, data_cred_req = self._guc_creds(transfers[0], third_party)
        cmd += cmd_creds
       
        # options
        cmd += self._guc_options(attempt, create_dest,
                                 third_party,
                                 data_cred_req)
        
        cmd += " -f " + tmp_name
        try:
            tc = TimedCommand(cmd)
            tc.run()
            transfer_success = True
        except Exception, err:
            logger.error(err)
            gsiftp_failures += 1
    
        if transfer_success:
            # stat the files
            for i, filename in enumerate(delayed_file_stat): 
                stats_add(filename)
    
        os.unlink(tmp_name)
    
        return transfer_success
    
    
    def _guc_creds(self, sample_transfer, third_party):
        """
        determine and set up credentials for the transfer
        """

        cmd = ""
        data_cred_req = False
        
        # gsi credentials for gsiftp transfers, ssh credentials
        # for sshftp 
        
        if sample_transfer.src_proto == "gsiftp" or \
           sample_transfer.dst_proto == "gsiftp":
            # separate credentials for src and dst? note that if one is
            # set, so must the other even if a credential is not required
            src_cred = None
            dst_cred = None
            key = "X509_USER_PROXY_" + sample_transfer.src_sitename
            if key in os.environ:
                src_cred = os.environ[key]
            key = "X509_USER_PROXY_" + sample_transfer.dst_sitename
            if key in os.environ:
                dst_cred = os.environ[key]
        
            # only set src-cred / dst-cred if at least one is specified
            if src_cred is not None or dst_cred is not None:
        
                if src_cred is None:
                    if "X509_USER_PROXY" in os.environ:
                        src_cred = os.environ["X509_USER_PROXY"]
                    else:
                        src_cred = dst_cred
        
                if dst_cred is None:
                    if "X509_USER_PROXY" in os.environ:
                        dst_cred = os.environ["X509_USER_PROXY"]
                    else:
                        dst_cred = src_cred
        
                cmd += " -src-cred " + src_cred + " -dst-cred " + dst_cred
                check_cred_fs_permissions(src_cred)
                check_cred_fs_permissions(dst_cred)
        
            # if the src and dest credentials are different, we might need
            # a data channel credential
            if third_party and src_cred != dst_cred:
                data_cred_req = True
        
        else:
            # sshftp transfers
            if sample_transfer.dst_proto == "file":
                # sshftp -> file
                key = "SSH_PRIVATE_KEY_" + sample_transfer.src_sitename
                if key in os.environ:
                    os.environ['SSH_PRIVATE_KEY'] = os.environ[key]
            else:
                # file -> sshftp
                key = "SSH_PRIVATE_KEY_" + sample_transfer.dst_sitename
                if key in os.environ:
                    os.environ['SSH_PRIVATE_KEY'] = os.environ[key]
            check_cred_fs_permissions(os.environ['SSH_PRIVATE_KEY'])

        return cmd, data_cred_req
            
   
    def _guc_options(self, attempt, create_dest, third_party, data_cred_req):
        """
        determine a set of globus-url-copy options based on how previous
        transfers went
        """
        global gsiftp_failures
    
        tools = Tools()
        options = " -r"
    
        # make output from guc match our current log level
        if logger.isEnabledFor(logging.DEBUG):
            options += " -verbose"
    
        # should we try to create directories?
        if create_dest:
            options += " -create-dest"
    
        # Only do third party transfers for gsiftp->gsiftp. For other
        # combinations, fall back to settings which will for well over for
        # example NAT
        if third_party:

            # pipeline is experimental so only allow this for the first attempt
            if gsiftp_failures == 0:
                options += " -pipeline"
    
            # parallism
            options += " -parallel 4"
    
            # -fast should be supported by all servers today
            options += " -fast"

            if data_cred_req:
                if gsiftp_failures == 0:
                    options += " -data-cred auto"
                else:
                    options += " -no-data-channel-authentication"
        else:
            # gsiftp<->file transfers
            options += " -no-third-party-transfers" \
                     + " -no-data-channel-authentication"
    
        return options
    
    def _check_similar(self, a, b):
        """
        compares two url_pairs, and determins if they are similar enough to be
        grouped together in one transfer input file
        """
        if a.src_host != b.src_host:
            return False
        if a.dst_host != b.dst_host:
            return False
        if os.path.dirname(a.src_path) != os.path.dirname(b.src_path):
            return False
        if os.path.dirname(a.dst_path) != os.path.dirname(b.dst_path):
            return False
        return True

    def _split_similar(self, full_list):
        """
        splits up a long list of similar transfers into smaller
        pieces which can easily be handled by g-u-c
        """
        chunks = []
        size = 1000
        num_chunks = int(math.ceil(len(full_list) / float(size)))
        for i in range(num_chunks):
            start = i * size
            end  = min((i + 1) * size, len(full_list))
            chunks.append(full_list[start:end])
        return chunks
  


class HTTPHandler(TransferHandlerBase):
    """
    pulls from http/https using wget
    """

    _name = "HTTPHandler"
    _protocol_map = ["http->file", "https->file"]

    def do_transfer(self, transfer, attempt):
        
        tools = Tools()
                               
        # Open Science Grid sites can inform us about local Squid proxies
        if "OSG_SQUID_LOCATION" in os.environ \
           and not "http_proxy" in os.environ:
            os.environ['http_proxy'] = os.environ['OSG_SQUID_LOCATION']

        # but only allow squid caching for the first try - after that go to
        # the source
        if attempt > 2 and "http_proxy" in os.environ:
            logger.info("Disabling HTTP proxy due to previous failures")
            del os.environ['http_proxy']

        prepare_local_dir(os.path.dirname(transfer.dst_path))
        
        # try wget first, then curl
        if tools.find("wget", "--version", "([0-9]+\.[0-9]+)") is not None:
            cmd = tools.full_path('wget')
            if logger.isEnabledFor(logging.DEBUG):
                cmd += " -v"
            else:
                cmd += " -nv"
            cmd += " --no-cookies --no-check-certificate" + \
                   " -O '" + transfer.dst_path + "'" + \
                   " '" + transfer.src_url() + "'"
        elif tools.find("curl", "--version", " ([0-9]+\.[0-9]+)") is not None:
            cmd = tools.full_path('curl')
            if not logger.isEnabledFor(logging.DEBUG):
                cmd += " -s -S"
            cmd += " --insecure" + \
                   " -o '" + transfer.dst_path + "'" + \
                   " '" + transfer.src_url() + "'"
        else:
            raise RuntimeError("Unable to do http/https transfers because" +
                               " neither curl nor wget could not be found")
            
        try:
            tc = TimedCommand(cmd)
            tc.run()
            stats_add(transfer.dst_path)
        except RuntimeError, err:
            logger.error(err)
            # wget and curl might leave 0 sized files behind after failures
            # make sure those get cleaned up
            try:
                os.unlink(transfer.dst_path)
            except:
                pass
            return False
        return True


class IRodsHandler(TransferHandlerBase):
    """
    Handler for iRods - http://www.irods.org/
    """
    
    _name = "IRodsHandler"
    _protocol_map = ["file->irods", "irods->file"]
    
    ticket = None

    def _irods_login(self, sitename):
        """
        log in to irods by using the iinit command - if the file already exists,
        we are already logged in
        """
        self.lock.acquire()
        try:
            key = "irodsEnvFile_" + sitename
            if key in os.environ:
                os.environ["irodsEnvFile"] = os.environ[key]

            if self.ticket != None:
                return
            
            # irods 4.0 rename
            os.environ['IRODS_ENVIRONMENT_FILE'] = os.environ["irodsEnvFile"]
            os.environ['IRODS_AUTHENTICATION_FILE'] = os.environ['irodsAuthFileName']

            f = os.environ['irodsAuthFileName']
            if os.path.exists(f):
                return

            # read password from env file
            if not "irodsEnvFile" in os.environ:
                raise RuntimeError("Missing irodsEnvFile - unable to do irods "
                                   + " transfers")

            check_cred_fs_permissions(os.environ["irodsEnvFile"])

            password = None
            self.ticket = None
            h = open(os.environ['irodsEnvFile'], 'r')
            for line in h:

                # flat file format
                items = line.split(" ", 2)
                if items[0].lower() == "irodspassword":
                    password = items[1].strip(" \t'\"\r\n")
                if items[0].lower() == "irodsticket":
                    self.ticket = items[1].strip(" \t'\"\r\n")

                # json - irods 4.0
                items = line.split(":", 2)
                key = items[0].lower().strip(" \t'\"\r\n")
                if key == "irodspassword":
                    password = items[1].strip(" \t'\",\r\n")

            h.close()
            if password is None and self.ticket is None:
                raise RuntimeError("No irodsTicket or irodsPassword" +
                                   " specified in irods env file")

            if password is not None:
                h = open(".irodsAc", "w")
                h.write(password + "\n")
                h.close()
                cmd = "cat .irodsAc | iinit"
                tc = TimedCommand(cmd, timeout_secs = 5*60)
                tc.run()
                os.unlink(".irodsAc")
                check_cred_fs_permissions(os.environ['irodsAuthFileName'])

        finally:
            self.lock.release()
    
    
    def do_transfer(self, transfer, attempt):
        """
        irods - use the icommands to interact with irods
        """
    
        tools = Tools()
        if tools.find("iget", "-h", "Version[ \t]+([\.0-9a-zA-Z]+)") is None:
            raise RuntimeError("Unable to do irods transfers becuase iget"
                               + " could not be found in the current path")
   
        # log in to irods
        sitename = transfer.src_sitename
        if transfer.dst_proto == "irods":
            sitename = transfer.dst_sitename
        try:
            self._irods_login(sitename)
        except Exception, loginErr:
            logger.error(loginErr)
            raise RuntimeError("Unable to log into irods")

        if transfer.dst_proto == "file":
            # irods->file
            prepare_local_dir(os.path.dirname(transfer.dst_path))
            cmd = "iget -v -f -r -T -K"
            if self.ticket is not None:
                cmd += " -t " + irods.ticket
            if transfer.src_host is not "" and attempt == 1:
                cmd += " -R " + transfer.src_host
            cmd += " '" + transfer.src_path + "'"
            cmd += " '" + transfer.dst_path + "'"
        else:
            # file->irods
            cmd = "imkdir -p '" + os.path.dirname(transfer.dst_path) + "'"
            try:
                tc = TimedCommand(cmd, timeout_secs = 60*60)
                tc.run()
            except:
                # ignore errors from the mkdir command
                pass
            cmd = "iput -v -f -r -T -K" 
            if self.ticket is not None:
                cmd += " -t " + irods.ticket
            if transfer.dst_host is not "" and attempt == 1:
                cmd += " -R " + transfer.dst_host
            cmd += " '" + transfer.src_path + "'"
            cmd += " '" +  transfer.dst_path + "'"

        cmd += " </dev/null"
        try:
            tc = TimedCommand(cmd)
            tc.run()
            # stats      
            if transfer.dst_proto == "file":
                stats_add(transfer.dst_path)
            else:
                stats_add(transfer.src_path)
        except Exception, err:
            logger.error(err)
            return False
        return True



class S3Handler(TransferHandlerBase):
    """
    Handler for S3 and S3 compatible services
    """
    
    _name = "S3Handler"
    _protocol_map = [
                    "file->s3",
                    "file->s3s",
                    "s3->file",
                    "s3s->file", 
                    "s3->s3",
                    "s3->s3s",
                    "s3s->s3",
                    "s3s->s3s"
                    ]

    def do_transfer(self, transfer, attempt):
    
        tools = Tools()
        if tools.find("pegasus-s3", "help", None) is None:
            raise RuntimeError("Unable to do S3 transfers because"
                               + " pegasus-s3 could not be found")

        local_filename = None
        
        # use cp for s3->s3 transfers, and get/put when one end is a file://
        if (transfer.src_proto == "s3" or transfer.src_proto == "s3s") and \
           (transfer.dst_proto == "s3" or transfer.dst_proto == "s3s"):
            # s3 -> s3
            key = "S3CFG_" + transfer.src_sitename
            if key in os.environ:
                os.environ["S3CFG"] = os.environ[key]
            if "S3CFG" not in os.environ:
                logger.error("At least one of the S3CFG_" + 
                             transfer.src_sitename + " or S3CFG" +
                             " environment variables has to be set")
                return False

            check_cred_fs_permissions(os.environ["S3CFG"])

            cmd = "pegasus-s3 cp -f -c '%s' '%s'" % (transfer.src_url(),
                                                     transfer.dst_url())
        elif transfer.dst_proto == "file":
            key = "S3CFG_" + transfer.src_sitename
            if key in os.environ:
                os.environ["S3CFG"] = os.environ[key]
            if "S3CFG" not in os.environ:
                logger.error("At least one of the S3CFG_" + 
                             transfer.src_sitename + " or S3CFG" +
                             " environment variables has to be set")
                return False

            check_cred_fs_permissions(os.environ["S3CFG"])

            # this is a 'get'
            local_filename = transfer.dst_path
            prepare_local_dir(os.path.dirname(transfer.dst_path))
            cmd = "pegasus-s3 get '%s' '%s'" % (transfer.src_url(),
                                                transfer.dst_path)
        else:
            # this is a 'put'
            key = "S3CFG_" + transfer.dst_sitename
            if key in os.environ:
                os.environ["S3CFG"] = os.environ[key]
            if "S3CFG" not in os.environ:
                logger.error("At least one of the S3CFG_" + 
                             transfer.dst_sitename + " or S3CFG" +
                             " environment variables has to be set")
                return False

            check_cred_fs_permissions(os.environ["S3CFG"])

            local_filename = transfer.src_path
            cmd = "pegasus-s3 put -f -b '%s' '%s'" % (transfer.src_path,
                                                      transfer.dst_url())

        try:
            tc = TimedCommand(cmd)
            tc.run()
            if local_filename is not None:
                stats_add(local_filename)
        except Exception, err:
            logger.error(err)
            return False
        return True


class GSHandler(TransferHandlerBase):
    """
    Handler for Google Storage services
    """
    
    _name = "GSHandler"
    _protocol_map = [
                    "file->gs",
                    "gs->file",
                    "gs->gs",
                    ]

    def fix_env(self, site_name):
        
        key = "BOTO_CONFIG_" + site_name
        if key in os.environ:
            os.environ["BOTO_CONFIG"] = os.environ[key]
        if "BOTO_CONFIG" not in os.environ:
            logger.error("At least one of the BOTO_CONFIG_" + 
                         site_name + " or BOTO_CONFIG" +
                         " environment variables has to be set")
            return False

        key = "GOOGLE_PKCS12_" + site_name
        if key in os.environ:
            os.environ["GOOGLE_PKCS12"] = os.environ[key]
        if "GOOGLE_PKCS12" not in os.environ:
            logger.error("At least one of the GOOGLE_PKCS12_" + 
                         site_name + " or GOOGLE_PKCS12" +
                         " environment variables has to be set")
            return False

            check_cred_fs_permissions(os.environ["BOTO_CONFIG"])
            check_cred_fs_permissions(os.environ["GOOGLE_PKCS12"])                

        return True
    

    def do_transfer(self, transfer, attempt):
    
        tools = Tools()
        if tools.find("gsutil", "version", "gsutil version: ([\.0-9a-zA-Z]+)") is None:
            raise RuntimeError("Unable to do Google Storage transfers because"
                               + " the gsutil tool could not be found")

        local_filename = None
        
        # use cp for gs->gs transfers, and get/put when one end is a file://
        if (transfer.src_proto == "gs" and transfer.dst_proto == "gs"):
            # gs -> gs
            if not self.fix_env(transfer.src_sitename):
                return False

            cmd = "gsutil -q cp '%s' '%s'" % (transfer.src_url(),
                                              transfer.dst_url())
        elif transfer.dst_proto == "file":
            # this is a 'get'
            if not self.fix_env(transfer.src_sitename):
                return False

            local_filename = transfer.dst_path
            prepare_local_dir(os.path.dirname(transfer.dst_path))
            cmd = "gsutil -q cp '%s' '%s'" % (transfer.src_url(),
                                              transfer.dst_path)
        else:
            # this is a 'put'
            if not self.fix_env(transfer.dst_sitename):
                return False

            local_filename = transfer.src_path
            cmd = "gsutil -q cp '%s' '%s'" % (transfer.src_path,
                                              transfer.dst_url())

        # we need to update the boto config file to specify the full
        # path to the PKCS12 file
        try:
            tmp_fd, tmp_name = tempfile.mkstemp(prefix="pegasus-transfer-", suffix=".lst")
            tmp_file = os.fdopen(tmp_fd, "w+b")
        except:
          raise RuntimeError("Unable to create tmp file for gs boto file")
        try:
            conf = ConfigParser.SafeConfigParser()
            conf.read(os.environ["BOTO_CONFIG"])
            conf.set("Credentials", "gs_service_key_file", os.environ["GOOGLE_PKCS12"])
            conf.write(tmp_file)
            tmp_file.close()
        except Exception, err:
            logger.error(err)
            raise RuntimeError("Unable to convert boto config file")
        
        cmd = "BOTO_CONFIG=%s %s" %(tmp_name, cmd)
        
        result = True
        try:
            tc = TimedCommand(cmd)
            tc.run()
            if local_filename is not None:
                stats_add(local_filename)
        except Exception, err:
            logger.error(err)
            result = False
        os.unlink(tmp_name)
        return result


class SRMHandler(TransferHandlerBase):
    
    _name = "SRMHandler"
    _protocol_map = ["srm->file", "file->srm", "gsiftp->srm", "srm->gsiftp"]

    def do_transfer(self, transfer, failed_q):
        """
        srm - lcg-cp is the preferred clienr, srm-copy the backup one
              Is this generic enough? Do we need to handle space tokens?
        """
        
        tools = Tools()
        if tools.find("lcg-cp",
                      "--version", 
                      "lcg_util-([\.0-9a-zA-Z]+)") is None \
           and \
           tools.find("srm-copy",
                      "-version",
                      "srm-copy[ \t]+([\.0-9a-zA-Z]+)") is None:
            raise RuntimeError("Unable to do srm transfers because" 
                               + " lcg-cp/srm-copy could not be found")
        
        if transfer.dst_proto == "file":
            key = "X509_USER_PROXY_" + transfer.src_sitename
            if key in os.environ:
                os.environ["X509_USER_PROXY"] = os.environ[key]
        else:
            key = "X509_USER_PROXY_" + transfer.dst_sitename
            if key in os.environ:
                os.environ["X509_USER_PROXY"] = os.environ[key]

        if transfer.dst_proto == "file":
            prepare_local_dir(os.path.dirname(transfer.dst_path))
            
        third_party = (transfer.src_proto == "gsiftp" 
                       or transfer.src_proto == "srm") and \
                      (transfer.dst_proto == "gsiftp"
                       or transfer.dst_proto == "srm")
      
        # prefer lcg-cp
        if tools.full_path('lcg-cp') is not None:
            cmd = "lcg-cp"
            if logger.isEnabledFor(logging.DEBUG):
                cmd = cmd + " -v"
            cmd = cmd + " -b -D srmv2 '%s' '%s'" \
                  % (transfer.src_url_srm(), transfer.dst_url_srm())
        else:
            cmd = "srm-copy '%s' '%s' -mkdir" \
                  % (transfer.src_url_srm(), transfer.dst_url_srm())
            if third_party:
                cmd = cmd + " -parallelism 4 -3partycopy"
            if not logger.isEnabledFor(logging.DEBUG):
                cmd = cmd + " >/dev/null"
            
        try:
            tc = TimedCommand(cmd)
            tc.run()
        except Exception, err:
            logger.error(err)
            return False
        return True
                

class ScpHandler(TransferHandlerBase):
    """
    Uses scp to copy to/from remote hosts
    """

    _name = "ScpHandler"
    _protocol_map = ["scp->file", "file->scp"]

    def do_transfer(self, transfer, attempt):
        global remote_dirs_created
        cmd = "/usr/bin/scp"
        cmd += " -r -B -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no"
        try:
            if transfer.dst_proto == "file":
                # scp -> file
                key = "SSH_PRIVATE_KEY_" + transfer.src_sitename
                if key in os.environ:
                    check_cred_fs_permissions(os.environ[key])
                    cmd += " -i " + os.environ[key]
                elif "SSH_PRIVATE_KEY" in os.environ:
                    check_cred_fs_permissions(os.environ['SSH_PRIVATE_KEY'])
                    cmd += " -i " + os.environ['SSH_PRIVATE_KEY']

                cmd += " -P " + self._extract_port(transfer.src_host)
                prepare_local_dir(os.path.dirname(transfer.dst_path))
                # scp wants escaped remote paths, even with quotes
                src_path = re.sub(' ', '\\ ', transfer.src_path)
                cmd += " '" + self._extract_hostname(transfer.src_host) + ":" + src_path + "'"
                cmd += " '" + transfer.dst_path + "'"
            else:
                # file -> scp
                key = "SSH_PRIVATE_KEY_" + transfer.dst_sitename
                if key in os.environ:
                    check_cred_fs_permissions(os.environ[key])
                    cmd += " -i " + os.environ[key]
                elif "SSH_PRIVATE_KEY" in os.environ:
                    check_cred_fs_permissions(os.environ['SSH_PRIVATE_KEY'])
                    cmd += " -i " + os.environ['SSH_PRIVATE_KEY']

                mkdir_key = "scp://" + self._extract_hostname(transfer.dst_host) + ":" \
                          + os.path.dirname(transfer.dst_path)
                if not mkdir_key in remote_dirs_created:
                    self._prepare_scp_dir(transfer.dst_sitename,
                                          transfer.dst_host,
                                          os.path.dirname(transfer.dst_path))
                    remote_dirs_created[mkdir_key] = True
                
                cmd += " -P " + self._extract_port(transfer.dst_host)
                # scp wants escaped remote paths, even with quotes
                dst_path = re.sub(' ', '\\ ', transfer.dst_path)
                
                cmd += " '" + transfer.src_path + "'"
                cmd += " '" + self._extract_hostname(transfer.dst_host) + ":" + dst_path + "'"
                stats_add(transfer.src_path)

            tc = TimedCommand(cmd)
            tc.run()
            if transfer.dst_proto == "file":
                stats_add(transfer.dst_path)    

        except RuntimeError, err:
            logger.error(err)
            return False
        return True


    def _prepare_scp_dir(self, rsite, rhost, rdir):
        """
        makes sure a local path exists before putting files into it
        """
        cmd = "/usr/bin/ssh"
        key = "SSH_PRIVATE_KEY_" + rsite
        if key in os.environ:
            check_cred_fs_permissions(os.environ[key])
            cmd += " -i " + os.environ[key]
        elif "SSH_PRIVATE_KEY" in os.environ:
            check_cred_fs_permissions(os.environ['SSH_PRIVATE_KEY'])
            cmd += " -i " + os.environ['SSH_PRIVATE_KEY']
        cmd += " -o UserKnownHostsFile=/dev/null"
        cmd += " -o StrictHostKeyChecking=no"
        cmd += " -p " + self._extract_port(rhost)
        cmd += " " + self._extract_hostname(rhost) + " '/bin/mkdir -p " + rdir + "'"
        tc = TimedCommand(cmd)
        tc.run()


    def _extract_hostname(self, host):
        """
        returns only the host component (strips :...)
        """
        return re.sub(':.*', '', host)

    
    def _extract_port(self, host):
        """
        returns the port to connect to, defaults to 22
        """
        port = "22"
        r = re.search(':([0-9]+)', host)
        if r:
            port = r.group(1)
        return str(port)


class StashHandler(TransferHandlerBase):
    """
    Uses the OSG stashcp command to trasfer from stash to file
    """
    
    _name = "StashHandler"
    _protocol_map = ["stash->file"]

    def do_transfer(self, transfer, attempt):
        
        tools = Tools()
        if tools.find("stashcp") is None:
            raise RuntimeError("Unable to do Stash transfers because" 
                               + " stashcp could not be found")

        # stashcp wants just the path with no leading slash
        src_path = transfer.src_url()
        src_path = re.sub('^stash:', '', src_path)
        src_path = re.sub('^/+', '', src_path)

        local_dir = os.path.dirname(transfer.dst_path)
        prepare_local_dir(local_dir)
        cmd = "%s -s '%s' -l '%s'" \
            % (tools.full_path("stashcp"), src_path, local_dir)
        remote_fname = os.path.basename(transfer.src_path)
        local_fname = os.path.basename(transfer.dst_path)
        if remote_fname != local_fname:
            cmd += " && mv %s %s" %(remote_fname, local_fname)
        try:
            tc = TimedCommand(cmd)
            tc.run()
        except RuntimeError, err:
            logger.error(err)
            return False
        stats_add(transfer.dst_path)
        return True


class SymlinkHandler(TransferHandlerBase):
    """
    Sets up symlinks - this is often used when data is local, but needs a
    reference in cwd
    """
    
    _name = "SymlinkHandler"
    _protocol_map = ["file->symlink", "symlink->symlink"]

    def do_transfer(self, transfer, attempt):

        prepare_local_dir(os.path.dirname(transfer.dst_path))

        # we do not allow dangling symlinks
        if not os.path.exists(transfer.src_path):
            logger.warning("Symlink source (%s) does not exist"
                           % (transfer.src_path))
            failed_q.put(transfer)
            return True

        if os.path.exists(transfer.src_path) \
           and os.path.exists(transfer.dst_path):
            # make sure src and target are not the same file - have to
            # compare at the inode level as paths can differ
            src_inode = os.stat(transfer.src_path)[stat.ST_INO]
            dst_inode = os.stat(transfer.dst_path)[stat.ST_INO]
            if src_inode == dst_inode:
                logger.warning("symlink: src (%s) and dst (%s) already exists"
                               % (transfer.src_path, transfer.dst_path))
                return True

        cmd = "ln -f -s '%s' '%s'" % (transfer.src_path, transfer.dst_path)
        try:
            tc = TimedCommand(cmd, timeout_secs = 60)
            tc.run()
        except RuntimeError, err:
            logger.error(err)
            return False
        return True


class TransferSet:
    """
    A transfer set is a set of similar transfers, similar in the sense
    that all the transfers have the same source and destination protocols
    """

    _transfers = None
    _available_handlers = []
    _primary_handler = None
    _secondary_handler = None
    _tmp_file = None

    def __init__(self, transfers_l):

        self._transfers = transfers_l

        # load all the handlers - does the order matter?
        self._available_handlers.append( CpHandler() )
        self._available_handlers.append( FDTHandler() )
        self._available_handlers.append( GridFtpHandler() )
        self._available_handlers.append( HTTPHandler() )
        self._available_handlers.append( IRodsHandler() )
        self._available_handlers.append( S3Handler() )
        self._available_handlers.append( GSHandler() )
        self._available_handlers.append( SRMHandler() )
        self._available_handlers.append( ScpHandler() )
        self._available_handlers.append( StashHandler() )
        self._available_handlers.append( SymlinkHandler() )

        src_proto = transfers_l[0].src_proto
        dst_proto = transfers_l[0].dst_proto

        # can we find one handler which can handle both source
        # and destination protocols directly?
        for h in self._available_handlers:
            if h.protocol_map_check(src_proto, dst_proto):
                self._primary_handler = h
                logger.debug("Selected %s for handling these transfers" 
                             %(h._name))
                return

        # we need to split the transfer from src to local file,
        # and then transfer the local file to the dst
        for h in self._available_handlers:
            if h.protocol_map_check(src_proto, "file"):
                self._primary_handler = h
                break
        for h in self._available_handlers:
            if h.protocol_map_check("file", dst_proto):
                self._secondary_handler = h
                break
        if self._primary_handler is None or self._secondary_handler is None:
            raise RuntimeError("Unable to find handlers for '%s' to '%s'"
                               %(src_proto, dst_proto))

        logger.debug("Selected %s and %s for handling these transfers"
                     %(self._primary_handler._name,
                       self._secondary_handler._name))


    def do_transfers(self, failed_q, attempt):
        """
        given a list of transfers, figure out what handlers are needed
        and then execute the transfers
        """

        self._tmp_name = None
        if self._secondary_handler is not None:
            # we have a two stage transfer to deal with and we need a temp file
            tmp_fd, self._tmp_name = \
                tempfile.mkstemp(prefix="pegasus-transfer-two-stage-",
                                 suffix=".data")
            # the temp file is only used in shell callouts, so let's close
            # the fd to make sure we are not leaving any open fds around
            try:
                os.close(tmp_fd)
            except:
                pass
            # open the permission up to make sure files downstream
            # get sane permissions to inherit
            os.chmod(self._tmp_name, 0644)
            logger.debug("Using temporary file %s for transfers" 
                         %(self._tmp_name))

        # gsiftp multi transfer case
        if (self._secondary_handler is None) and \
           (self._transfers[0].src_proto == "gsiftp" \
            or self._transfers[0].dst_proto == "gsiftp"):
            success = False
            try:
                success =  self._primary_handler.do_multi_transfer(
                                                    list(self._transfers),
                                                    attempt,
                                                    failed_q)
            except RuntimeError, rte:
                raise
            except Exception, e:
                if logger.isEnabledFor(logging.DEBUG):
                    logger.exception("Exception while doing transfer:")
                else:
                    logger.error(e)
                # add all the transfers back to the failed queue
                for i, t in enumerate(self._transfers):
                    failed_q.put(t)
            return success
        
        # standard src->dst single transfer case
        for i, transfer in enumerate(self._transfers):
            
            # We are being extra careful to detect failures here. We are 
            # considering both a False being returned or an exception being
            # thrown as a failed transfer
            success = False

            if self._secondary_handler is None:
                # one handler to rule them all!
                try:
                    success = self._primary_handler.do_transfer(transfer,
                                                                attempt)
                except Exception, e:
                    if logger.isEnabledFor(logging.DEBUG):
                        logger.exception("Exception while doing transfer:")
                    else:
                        logger.error(e)
            else:
                # break up the transfer into two, but keep a handle to the main
                # transfer as that is the one which will have to go back to the
                # failed queue in case of failure
                t_one = Transfer(transfer.pair_id)
                t_one.set_src(transfer.src_url())
                t_one.set_src_sitename(transfer.src_sitename)
                t_one.set_dst("file://" + self._tmp_name)
                t_two =  Transfer(transfer.pair_id)
                t_two.set_src("file://" + self._tmp_name)
                t_two.set_dst(transfer.dst_url())
                t_two.set_dst_sitename(transfer.dst_sitename)
                try:
                    success = self._primary_handler.do_transfer(t_one, attempt)
                    if success:
                        os.chmod(self._tmp_name, 0644)
                        success = self._secondary_handler.do_transfer(t_two, attempt)
                except Exception, e:
                    if logger.isEnabledFor(logging.DEBUG):
                        logger.exception("Exception while doing transfer:")
                    else:
                        logger.error(e)

            if success == False:
                failed_q.put(transfer)
        
        # remove temp file
        if self._tmp_name is not None:
            logger.debug("Removing temporary file %s" %(self._tmp_name))
            try:
                os.unlink(self._tmp_name)
            except:
                pass


    def has_gridftp_transfers(self):
        """
        Check if this transfer set has gridftp transfers
        """
        if (self._transfers[0].src_proto == "gsiftp" \
            or self._transfers[0].dst_proto == "gsiftp"):
            return True
        return False
        
        

class WorkThread(threading.Thread):
    """
    A thread which processes TransferSets
    """
    def __init__(self, thread_id, queue, current_attempt, failed_queue):
        threading.Thread.__init__(self)
        self.thread_id = thread_id
        self.queue = queue
        self.current_attempt = current_attempt
        self.failed_queue = failed_queue
        self.exception = None
        self.daemon = True

    def run(self):
        # give the threads a slow start
        time.sleep(self.thread_id * 2)
        logger.debug("Started new WorkThread with id " + str(self.thread_id))
        try:
            # Just keep grabbing TransferSets and executing them until
            # there are no more to process, then exit
            while True:
                ts = self.queue.get(False)
                logger.debug("Thread " + str(self.thread_id) +
                             " is executing transfer " + str(ts))
                ts.do_transfers(self.failed_queue, self.current_attempt)
        except Queue.Empty:
            return
        except Exception, e:
            self.exception = e


class Alarm(Exception):
    pass


# --- global variables ----------------------------------------------------------------

prog_dir  = os.path.realpath(os.path.join(os.path.dirname(sys.argv[0])))
prog_base = os.path.split(sys.argv[0])[1]   # Name of this program

logger = logging.getLogger("my_logger")

# threads we have currently running
threads = []

# track remote directories created so that don't have to
# try to create them over and over again
remote_dirs_created = {}

# gsiftp failure count - used to provide sane globus-url-copy options
gsiftp_failures = 0

# stats
stats_start = 0
stats_end = 0
stats_total_bytes = 0

# singleton - but should we make it a global instead?
tools = Tools()


# --- functions ----------------------------------------------------------------


def setup_logger(debug_flag):
    
    # log to the console
    console = logging.StreamHandler()
    
    # default log level - make logger/console match
    logger.setLevel(logging.INFO)
    console.setLevel(logging.INFO)

    # debug - from command line
    if debug_flag:
        logger.setLevel(logging.DEBUG)
        console.setLevel(logging.DEBUG)

    # formatter
    formatter = logging.Formatter("%(asctime)s %(levelname)7s:  %(message)s")
    console.setFormatter(formatter)
    logger.addHandler(console)
    logger.debug("Logger has been configured")


def prog_sigint_handler(signum, frame):
    logger.warn("Exiting due to signal %d" % (signum))
    myexit(1)


def alarm_handler(signum, frame):
    raise Alarm


def expand_env_vars(s):
    re_env_var = re.compile(r'\${?([a-zA-Z][a-zA-Z0-9_]+)}?')
    s = re.sub(re_env_var, get_env_var, s)
    return s


def get_env_var(match):
    name = match.group(1)
    value = ""
    logger.debug("Looking up " + name + " environment variable")
    if name in os.environ:
        value = os.environ[name]
    return value


def backticks(cmd_line):
    """
    what would a python program be without some perl love?
    """
    return subprocess.Popen(cmd_line, shell=True,
                            stdout=subprocess.PIPE).communicate()[0]


def env_setup():
    
    # PATH setup
    path = "/usr/bin:/bin"
    if "PATH" in os.environ:
        path = os.environ['PATH']
    path_entries = path.split(':')
    
    # is /usr/bin in the path?
    if not("/usr/bin" in path_entries):
        path_entries.append("/usr/bin")
        path_entries.append("/bin")

    # fink on macos x
    if os.path.exists("/sw/bin") and not("/sw/bin" in path_entries):
        path_entries.append("/sw/bin")
    
    # prepend the path pegasus-transfer was launched from - this ensures
    # a consistent environment where all the Pegasus tools are picked
    # up from the same location
    try:
        path_entries.remove(prog_dir)
    except Exception:
        pass
    path_entries.insert(0, prog_dir)
    
    # if PEGASUS_HOME is set, prepend it to the PATH (we want it early to
    # override other cruft)
    if "PEGASUS_HOME" in os.environ:
        try:
            path_entries.remove(os.environ['PEGASUS_HOME'] + "/bin")
        except Exception:
            pass
        path_entries.insert(0, os.environ['PEGASUS_HOME'] + "/bin")
    
    # need LD_LIBRARY_PATH for Globus tools
    ld_library_path = ""
    if "LD_LIBRARY_PATH" in os.environ:
        ld_library_path = os.environ['LD_LIBRARY_PATH']
    ld_library_path_entries = ld_library_path.split(':')

    # if GLOBUS_LOCATION is set, prepend it to the PATH and LD_LIBRARY_PATH 
    # (we want it early to override other cruft)
    if "GLOBUS_LOCATION" in os.environ:
        try:
            path_entries.remove(os.environ['GLOBUS_LOCATION'] + "/bin")
        except Exception:
            pass
        path_entries.insert(0, os.environ['GLOBUS_LOCATION'] + "/bin")
        try:
            ld_library_path_entries.remove(
                os.environ['GLOBUS_LOCATION'] + "/lib")
        except Exception:
            pass
        ld_library_path_entries.insert(0, 
                                       os.environ['GLOBUS_LOCATION'] + "/lib")

    os.environ['PATH'] = ":".join(path_entries)
    os.environ['LD_LIBRARY_PATH'] = ":".join(ld_library_path_entries)
    os.environ['DYLD_LIBRARY_PATH'] = ":".join(ld_library_path_entries)
    logger.info("PATH=" + os.environ['PATH'])
    logger.info("LD_LIBRARY_PATH=" + os.environ['LD_LIBRARY_PATH'])
    
    # irods requires a password hash file
    os.environ['irodsAuthFileName'] = os.getcwd() + "/.irodsA"


def check_cred_fs_permissions(path):
    """
    Checks to make sure a given credential is protected by the file system
    permissions. If left too open (for example after a transfer over GASS,
    chmod it to be readable only by us.
    """
    if not os.path.exists(path):
        raise Exception("Credential file %s does not exist" %(path))
    if oct(os.stat(path).st_mode & 0777) != '0600':
        logger.warning("%s found to have weak permissions. chmod to 0600."
                       %(path))
        os.chmod(path, 0600)


def prepare_local_dir(path):
    """
    makes sure a local path exists before putting files into it
    """
    if not(os.path.exists(path)):
        logger.debug("Creating local directory " + path)
        try:
            os.makedirs(path, 0755)
        except os.error, err:
            # if dir already exists, ignore the error
            if not(os.path.isdir(path)):
                raise RuntimeError(err)


def transfers_groupable(a, b):
    """
    compares two url_pairs, and determins if they are similar enough to be
    grouped together for one tool
    """
    if not a.groupable() or not b.groupable():
        return False
    if a.src_proto != b.src_proto:
        return False
    if a.dst_proto != b.dst_proto:
        return False
    return True


def stats_add(filename):
    global stats_total_bytes
    try:
        s = os.stat(filename)
        stats_total_bytes = stats_total_bytes + s[stat.ST_SIZE]
    except Exception, err:
        pass # ignore


def stats_summarize():
    if stats_total_bytes == 0:
        logger.info("Stats: no local files in the transfer set")
        return

    total_secs = stats_end - stats_start
    Bps = stats_total_bytes / total_secs

    logger.info("Stats: %sB transferred in %.0f seconds. Rate: %sB/s (%sb/s)" \
                % (iso_prefix_formatted(stats_total_bytes), total_secs, 
                   iso_prefix_formatted(Bps), iso_prefix_formatted(Bps*8)))
    logger.info("NOTE: stats do not include third party gsiftp/srm transfers")


def iso_prefix_formatted(n):
    prefix = ""
    n = float(n)
    if n > (1024*1024*1024*1024):
        prefix = "T"
        n = n / (1024*1024*1024*1024)
    elif n > (1024*1024*1024):
        prefix = "G"
        n = n / (1024*1024*1024)
    elif n > (1024*1024):
        prefix = "M"
        n = n / (1024*1024)
    elif n > (1024):
        prefix = "K"
        n = n / (1024)
    return "%.1f %s" % (n, prefix)


def myexit(rc):
    """
    system exit without a stack trace
    """
    try:
        sys.exit(rc)
    except SystemExit:
        sys.exit(rc)


# --- main ----------------------------------------------------------------------------

def main():
    global threads
    global stats_start
    global stats_end
    
    # dup stderr onto stdout
    sys.stderr = sys.stdout
    
    # Configure command line option parser
    prog_usage = "usage: %s [options]" % (prog_base)
    parser = optparse.OptionParser(usage=prog_usage)
    
    parser.add_option("-f", "--file", action = "store", dest = "file",
                      help = "File containing URL pairs to be transferred." +
                             " If not given, list is read from stdin.")
    parser.add_option("-m", "--max-attempts", action = "store", type="int",
                      dest = "max_attempts", default = 3,
                      help = "Number of attempts allowed for each transfer." +
                             " Default is 3.")
    parser.add_option("-n", "--threads", action = "store", type="int",
                      dest = "threads", default = 0,
                      help = "Number of threads to process transfers." +
                             " Default is 8. This option can also be set" +
                             " via the PEGASUS_TRANSFER_THREADS environment" +
                             " variable. The command line option takes" +
                             " precedence over the environment variable.")
    parser.add_option("-d", "--debug", action = "store_true", dest = "debug",
                      help = "Enables debugging ouput.")
    
    # Parse command line options
    (options, args) = parser.parse_args()
    setup_logger(options.debug)
    
    # Die nicely when asked to (Ctrl+C, system shutdown)
    signal.signal(signal.SIGINT, prog_sigint_handler)
    
    attempts_max = options.max_attempts
    
    if options.threads is None or options.threads == 0:
        if "PEGASUS_TRANSFER_THREADS" in os.environ:
            options.threads = int(os.environ["PEGASUS_TRANSFER_THREADS"])
        else:
            options.threads = 8
    
    # stdin or file input?
    if options.file is None:
        logger.info("Reading URL pairs from stdin")
        input_file = sys.stdin
    else:
        logger.info("Reading URL pairs from %s" % (options.file))
        try:
            input_file = open(options.file, 'r')
        except Exception, err:
            logger.critical('Error reading url pair list: %s' % (err))
            myexit(1)
    
    # check environment
    try:
        env_setup()
    except Exception, err:
        logger.critical(err)
        myexit(1)
    
    # queues to track the work
    transfer_q = Queue.Queue()
    failed_q = Queue.Queue()
    work_q = Queue.Queue()
    
    # fill the transfer queue with user provided entries
    line_nr = 0
    pair_nr = 0
    inputs = []
    line_state = 3  # 0=SrcComment, 1=SrcUrl, 2=DstComment, 3=DstUrl
    url_pair = None
    try:
        for line in input_file.readlines():
            line_nr += 1
            if len(line) > 4:
                line = line.rstrip('\n')

                # src comment
                if line_state == 3 and line[0] == '#':
                    line_state = 0
                    if url_pair is None:
                        pair_nr += 1
                        url_pair = Transfer(pair_nr)
                    r = re_parse_comment.search(line)
                    if r:
                        url_pair.set_src_sitename(r.group(1))
                    else:
                        logger.critical('Unable to parse comment on line %d' 
                                        %(line_nr))
                        myexit(1)
                
                # src url
                elif line_state == 0 or line_state == 3:
                    line_state = 1
                    if url_pair is None:
                        pair_nr += 1
                        url_pair = Transfer(pair_nr)
                    url_pair.set_src(line)

                # dst comment
                elif line_state == 1 and line[0] == '#':
                    line_state = 2
                    r = re_parse_comment.search(line)
                    if r:
                        url_pair.set_dst_sitename(r.group(1))
                    else:
                        logger.critical('Unable to parse comment on line %d'
                                        %(line_nr))
                        myexit(1)
                
                # dst url
                elif line_state == 2 or line_state == 1:
                    line_state = 3
                    url_pair.set_dst(line)
                    inputs.append(url_pair)
                    url_pair = None

    except Exception, err:
        logger.critical('Error reading url list: %s' % (err))
        myexit(1)
    total_transfers = len(inputs)
    logger.info("%d transfers loaded" % (total_transfers))

    # we will now sort the list as some tools (gridftp) can optimize when
    # given a group of similar transfers
    logger.info("Sorting the tranfers based on transfer type and" +
                " source/destination")
    inputs.sort()
    
    for t in inputs:
        transfer_q.put(t)
    
    # start the stats time
    stats_start = time.time()
    
    # Attempt transfers until the queue is empty. We create TransferSets
    # of the transfers, and then hand then of to our worker threads. But
    # note that we are only doing the threads for one attempt at a time.
    # After failures, transfers might be regrouped, and then handed of to the
    # thread pool again.
    done = False
    attempt_current = 0
    approx_transfer_per_thread = total_transfers / (float)(options.threads)
    while not done:
    
        attempt_current = attempt_current + 1
        logger.info('-' * 80)
        logger.info("Starting transfers - attempt %d" % (attempt_current))

        # organize the transfers
        while not transfer_q.empty():
            t_main = transfer_q.get()
            
            # create a list of transfers to pass to underlying tool
            t_list = []
            t_list.append(t_main)
    
            try:
                t_next = transfer_q.get(False)
                while t_next is not None:
                    if len(t_list) <  approx_transfer_per_thread and \
                       transfers_groupable(t_main, t_next):
                        t_list.append(t_next)
                        t_next = transfer_q.get(False)
                    else:
                        # done, put the last transfer back
                        transfer_q.put(t_next)
                        t_next = None
            except Queue.Empty:
                pass
            
            # magic!
            ts = TransferSet(t_list)
            work_q.put(ts)

        # pool of worker threads
        t_id = 0
        num_threads = min(options.threads, work_q.qsize())
        if attempt_current > 2:
            num_threads = 1
        logger.info("Using %d threads for this round of transfers"
                    %(num_threads))
        for i in range(num_threads):
            t_id += 1
            t = WorkThread(t_id, work_q, attempt_current, failed_q)
            threads.append(t)
            t.start()
    
        # wait for the threads to finish all the transfers
        for t in threads:
            t.join()
            # do we need to do any better error handling here?
            if t.exception is not None:
                if type(t.exception) is RuntimeError:
                    logger.critical(t.exception)
                    myexit(2)
                raise t.exception
        threads = []
    
        logger.debug("%d items in failed_q" %(failed_q.qsize()))
        
        # are we done?
        if attempt_current == attempts_max or failed_q.empty():
            done = True
            break
        
        # retry failed transfers with a random delay - useful when
        # large workflows overwhelm data services
        if not failed_q.empty() and attempt_current < attempts_max:
            d = random.randint(30, 180)
            logger.debug("Sleeping for the random amount of %d seconds before the next attempt" %(d))
            time.sleep(d)
        while not failed_q.empty():
            t = failed_q.get()
            t.allow_grouping = False # only allow grouping on the first try
            transfer_q.put(t)
    
    logger.info('-' * 80)
    
    # end the stats timer and show summary
    stats_end = time.time()
    stats_summarize()
    
    if not failed_q.empty():
        logger.critical("Some transfers failed! See above," +
                        " and possibly stderr.")
        myexit(1)
    
    logger.info("All transfers completed successfully.")
    
    myexit(0)


if __name__ == "__main__":
    main()
    

