path: root/ur_upgrade.py

#!/usr/bin/env python
from __future__ import division

import argparse
import os
import shutil
import sys
import subprocess
import tempfile
import urllib2
import signal
import tarfile
from distutils.version import StrictVersion
from collections import namedtuple
from string import Template
from itertools import chain

def which(cmd, mode=os.F_OK | os.X_OK, path=None):
    """Given a command, mode, and a PATH string, return the path which
    conforms to the given mode on the PATH, or None if there is no such
    file.

    `mode` defaults to os.F_OK | os.X_OK. `path` defaults to the result
    of os.environ.get("PATH"), or can be overridden with a custom search
    path.

    """
    # Check that a given file can be accessed with the correct mode.
    # Additionally check that `file` is not a directory, as on Windows
    # directories pass the os.access check.
    def _access_check(fn, mode):
        return (os.path.exists(fn) and os.access(fn, mode)
                and not os.path.isdir(fn))

    # If we're given a path with a directory part, look it up directly rather
    # than referring to PATH directories. This includes checking relative to the
    # current directory, e.g. ./script
    if os.path.dirname(cmd):
        if _access_check(cmd, mode):
            return cmd
        return None

    if path is None:
        path = os.environ.get("PATH", os.defpath)
    if not path:
        return None
    path = path.split(os.pathsep)

    if sys.platform == "win32":
        # The current directory takes precedence on Windows.
        if not os.curdir in path:
            path.insert(0, os.curdir)

        # PATHEXT is necessary to check on Windows.
        pathext = os.environ.get("PATHEXT", "").split(os.pathsep)
        # See if the given file matches any of the expected path extensions.
        # This will allow us to short circuit when given "python.exe".
        # If it does match, only test that one, otherwise we have to try
        # others.
        if any(cmd.lower().endswith(ext.lower()) for ext in pathext):
            files = [cmd]
        else:
            files = [cmd + ext for ext in pathext]
    else:
        # On other platforms you don't have things like PATHEXT to tell you
        # what file suffixes are executable, so just pass on cmd as-is.
        files = [cmd]

    seen = set()
    for di in path:
        normdir = os.path.normcase(di)
        if not normdir in seen:
            seen.add(normdir)
            for thefile in files:
                name = os.path.join(di, thefile)
                if _access_check(name, mode):
                    return name
    return None

class Ureka(object):
    ''' 
    '''
    def __init__(self, basepath):
        self.path = basepath
        self.path_data = os.path.abspath(os.path.join(self.path, 'misc'))
        self._files = ['os', 'bits', 'version', 'name']
        if os.path.exists(self.path_data):
            self.info = self._info()
        else:
            self.info = {}
        

    def _info(self):
        data = []
        for key in self._files:
            path = os.path.join(self.path_data, key)
            item = None
            try:
                item = open(path, 'r').readline().strip()
            except:
                item = ''
            data.append(item)
            
        ureka_info = namedtuple('ureka_info', self._files)
        return ureka_info._make(data)
    
    def __getitem__(self, key):
        info = self.info._asdict()
        if key not in info:
            return None
        return info[key]
    
    def __iter__(self):
        for item in self.info._asdict().iteritems():
            yield item

def ur_getenv(ur_dir):
    ''' Evaluate environment variables produced by ur-setup-real
    '''
    path = os.path.join(ur_dir, 'bin')
    command = '{}/ur-setup-real -sh'.format(path).split()
    output = subprocess.check_output(command, shell=True, stderr=open(os.path.devnull))
    output = output.split(os.linesep)
    output_env = []
    
    # Generate environment keypairs
    for line in output:
        if not line:
            continue
        if line.startswith('export'):
            continue
        line = line.strip()
        line = line.replace(' ', '')
        line = line.replace(';', '')
        line = line.replace('"', '')
        output_env.append(line.partition("=")[::2])
    
    output_env_temp = dict(output_env)
    
    # Perform shell expansion of Ureka's environment variables
    for k, v in output_env_temp.items():
        template = Template(v)
        v = template.safe_substitute(output_env_temp)
        output_env_temp[k] = v
    
    # Merge Ureka's environment with existing system environment
    output_env_temp = dict(chain(os.environ.items(), output_env_temp.items()))
    
    # Assign expanded variables
    output_env = output_env_temp
    
    return output_env

def ur_check_version(urobj, vers):
    if StrictVersion(urobj['version']) > vers:
        return False
    return True

def cleanup(sig, stack):
    print('')
    print('Received signal {}!'.format(sig))
    print('Cleaning up...')
    shutil.rmtree(UPGRADE_TEMP)
    exit(sig)
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Ureka Upgrade Utility')
    parser.add_argument('UR_DIR', action='store', type=str, help='Absolute path to Ureka installation')
    parser.add_argument('--latest', action='store_true', default=False, help='Upgrade to the latest (stable) version')
    parser.add_argument('--request', type=str, help='Upgrade to a specific version')
    parser.add_argument('--mirror', type=str, help='Use a Ureka download mirror')
    args = parser.parse_args()
    
    if args.latest and args.request:
        print("--latest and --request are mutually exclusive options.")
        exit(1)
    
    if args.latest:
        print("--latest is not implemented yet.")
        exit(1)
        
    if not args.request:
        print("No version requested. Use --request (e.g. 1.4.1)")
        exit(1)
    
    UPGRADE_TEMP = tempfile.mkdtemp(prefix="upgrade")
    ORIGINAL = args.UR_DIR
    
    # Register signal handlers AFTER temporary directory is created
    signal.signal(signal.SIGINT, cleanup)
    signal.signal(signal.SIGTERM, cleanup)
             
    if not which('rsync'):
        print('rsync not found in PATH. Please install it.')
        exit(1)
    
    u = Ureka(ORIGINAL)
    
    if not ur_check_version(u, args.request):
        print("Refusing to downgrade from {} to {}".format(u['version'], args.request))
        exit(1)
    
    INSTALL_PATH = os.path.join(UPGRADE_TEMP, args.request)
    UREKA_TARBALL = "Ureka_{}_{}_{}.tar.gz".format(u['os'], u['bits'], args.request)
    ''' http://ssb.stsci.edu/ureka/1.0/Ureka_linux-rhe5_64_1.4.1.tar.gz '''

    #url = "http://ssb.stsci.edu/ureka/{}/{}".format(args.version, UREKA_TARBALL)
    url = "http://localhost/ureka/{}/{}".format(args.request, UREKA_TARBALL)
    
    if args.mirror:
        url = "{}/{}/{}".format(args.mirror, args.request, UREKA_TARBALL)
    
    print("Downloading archive... {}".format(url))
    try:
        req_data = urllib2.urlopen(url)
    except Exception as ex:
        print("Error {}, {}: {}".format(ex.code, ex.msg, url))
        exit(1)
        
    remote_size = float(req_data.headers['content-length'])
    total = 0
    with open(os.path.join(UPGRADE_TEMP, UREKA_TARBALL), 'wb') as installer:
        for chunk in iter(lambda: req_data.read(16*1024), ''):
            print("\r{:.2f}% [{:.2f}/{:.2f} MB]".format((total / remote_size * 100), (total / 1024 ** 2), (remote_size / 1024 ** 2))),
            total += float(len(chunk))
            installer.write(chunk)
    print("")
    
    os.mkdir(INSTALL_PATH)
    tarball = tarfile.open(os.path.join(UPGRADE_TEMP, UREKA_TARBALL), 'r')
    print("Preparing to unpack upgrade... (please wait)")
    files_total = len(tarball.getmembers()) 
    
    print("Unpacking upgrade...")
    for index, member in enumerate(tarball, start=1):
        print("\r{:.2f}% [{}/{}]".format((index / files_total * 100), index, files_total)),
        tarball.extract(member, path=INSTALL_PATH)
    print("")
    
    open(os.path.join(INSTALL_PATH, 'Ureka', 'misc', 'name'), 'w+').write(u['name'] + os.linesep)
    n = Ureka(os.path.join(INSTALL_PATH, 'Ureka'))
    
    print("Upgrading distribution {} to {}...".format(u['version'], n['version']))
    proc = subprocess.Popen('rsync -a -u {} {}'.format(n.path, os.path.dirname(u.path)).split())
    proc.wait()
    
    # Regenerate upgraded ureka installation object
    u = Ureka(ORIGINAL)

    print("Performing upgrade maintenance tasks...")
    ur_env = ur_getenv(u.path)
    
    proc = subprocess.Popen('ur_normalize -n -i -x'.split(), env=ur_env)
    proc.wait()
    
    print("Purging temporary data...")
    shutil.rmtree(UPGRADE_TEMP)
    
    print("Upgrade complete!")