path: root/ur_upgrade.py

#!/usr/bin/env python
# Copyright (c) 2014, Joseph Hunkeler <jhunkeler at gmail.com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from __future__ import division

import argparse
import os
import shutil
import subprocess
import tempfile
import urllib2
import signal
import tarfile
import time
from distutils.version import StrictVersion
from distutils.spawn import find_executable as which
from collections import namedtuple
from string import Template
from itertools import chain

DEFAULT_MIRROR = "http://ssb.stsci.edu/ureka"
DEFAULT_PUBLIC_RELEASE = "{}/public_releases.txt".format(DEFAULT_MIRROR)

def retype(t):
    for cast in (int, float):
        try:
            return cast(t)
        except ValueError:
            pass
    return t

def get_vfs_data(path):
    path = os.path.abspath(path)
    if not os.path.exists(path):
        return False

    headers = ['device', 'type', 'size', 'used', 'available', 'percent', 'mountpoint']
    command = 'df -Tk {}'.format(path).split()
    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=open(os.devnull, 'w'))
    stdout, _ = proc.communicate()

    data = stdout.split('\n')[1].split()
    for index, record in enumerate(data):
        data[index] = retype(record)

    vfs_data = zip(headers, data)
    return dict(vfs_data)

def ur_storage_strategy(required_size, storage_map):
    vfs = []
    for d in storage_map:
        data = get_vfs_data(d)
        data['path'] = d
        if (data['available'] - required_size) > 0:
            vfs.append(data)

    try:
        best = max([dev['size'] for dev in vfs])
        if best in [dev['size'] for dev in vfs]:
            return dev['path']
    except ValueError:
        return None

def ur_getenv(ur_dir):
    ''' Evaluates 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']) > StrictVersion(vers):
        return False
    elif StrictVersion(urobj['version']) == StrictVersion(vers):
        return False
    return True

def ur_get_public_releases(m):
    data = []
    try:
        req_data = urllib2.urlopen(m)
        data = req_data.readlines()
        data = [ x.strip() for x in data ]
    except:
        return []
    return data


class Ureka(object):
    def __init__(self, basepath):
        self.path = os.path.abspath(basepath)
        self.path_data = os.path.join(self.path, 'misc')
        self._files = ['os', 'bits', 'version', 'name']

        if not os.path.exists(self.path):
            print('{} does not exist.'.format(self.path))
            exit(1)

        try:
            self.info = self._info()
        except:
            print('{} does not contain a valid Ureka installation.'.format(self.path))
            exit(1)

    def _info(self):
        data = []
        for key in self._files:
            path = os.path.join(self.path_data, key)
            item = None
            item = open(path, 'r').readline().strip()
            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


class Upgrade(object):
    def __init__(self, ur_dir, to_version, mirror=DEFAULT_MIRROR, **kwargs):
        self.mirror = mirror
        self.ureka = Ureka(ur_dir)
        self.ureka_next = None
        self.to_version = to_version
        _tmp = ur_storage_strategy(storage_requires, storage_areas)
        self.tmp = tempfile.mkdtemp(dir=_tmp, prefix="upgrade")
        if _tmp is None:
            print("++ failed to determine best upgrade storage strategy!")
            print("++ defaulting to {}".format(self.tmp))
        self.tmp_dist = os.path.join(self.tmp, self.to_version)
        self.force = False
        self.backup = True
        self.backup_path = os.path.abspath(os.curdir)
        self.archive_ext = '.tar.gz'

        if 'archive_ext' in kwargs:
            self.archive_ext = kwargs['archive_ext']

        self.archive = "Ureka_{}_{}_{}{}".format(self.ureka['os'],
                                                 self.ureka['bits'],
                                                 self.to_version,
                                                 self.archive_ext)
        self.archive_url = "{}/{}/{}".format(self.mirror,
                                             self.to_version,
                                             self.archive)
        self.archive_path = os.path.abspath(os.path.join(self.tmp, self.archive))

        signal.signal(signal.SIGINT, self._cleanup_on_signal)
        signal.signal(signal.SIGTERM, self._cleanup_on_signal)

        if not which('rsync'):
            self._cleanup()
            print('++ rsync not found in PATH. Please install it.')
            exit(1)

        if self.ureka.path in os.path.abspath(os.path.curdir):
            self._cleanup()
            print("Impossible. Please change to a directory above {}.".format(self.ureka.path))
            exit(1)

    def run(self):
        if not self.force:
            if not ur_check_version(self.ureka, self.to_version):
                self._cleanup()
                print("Refusing upgrade from {} to {}. Use --force to override.".format(self.ureka['version'], self.to_version))
                exit(1)

        self._get_archive()

        ureka_next_path = self._unpack_archive()

        if self.backup:
            self._backup()

        self._pre()

        # Populate temporary Ureka upgrade object
        self.ureka_next = Ureka(ureka_next_path)

        # Sync data
        if not self._upgrade(self.ureka_next, self.ureka):
            self._cleanup()
            print("++ Upgrade failed!")
            return 1

        # Regenerate original Ureka object
        self.ureka = Ureka(self.ureka.path)

        if not self._post(self.ureka):
            self._cleanup()
            print("++ Post-installation failed!")
            return 1

        self._cleanup()

    def _cleanup(self):
        shutil.rmtree(self.tmp)

    def _cleanup_on_signal(self, sig, stack):
        print('')
        print('++ Received signal {}...'.format(sig))
        self._cleanup()
        exit(sig)

    def _get_archive(self):
        print("+ Downloading archive... {}".format(self.archive_url))
        try:
            req_data = urllib2.urlopen(self.archive_url)
        except Exception as ex:
            print("Error {}, {}: {}".format(ex.code, ex.msg, self.archive_url))
            exit(1)

        remote_size = float(req_data.headers['content-length'])

        total = 0
        with open(self.archive_path, '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("")

    def _unpack_archive(self):
        print('+ Preparing to unpack (please wait)...'),
        installer = tarfile.open(self.archive_path, 'r')
        files_total = len(installer.getmembers())
        print('done')
        print("+ Unpacking...")
        for index, member in enumerate(installer, start=1):
            print("\r{:.2f}% [{}/{}]".format((index / files_total * 100),
                                             index,
                                             files_total)),
            installer.extract(member, path=self.tmp_dist)
        print("")

        return os.path.join(self.tmp_dist, 'Ureka')

    def _backup(self):
        filename = '{}-{}-{}-backup.tar'.format(self.ureka['name'],
                                            self.ureka['version'],
                                            str(time.time()))
        backup_path = os.path.join(self.backup_path, filename)
        print('+ Generating backup... {}'.format(filename))
        with tarfile.open(backup_path, 'w') as tar:
            tar.add(self.ureka.path, arcname=os.path.basename(self.ureka.path))

    def _pre(self):
        print("+ Executing pre-upgrade tasks...")
        misc = os.path.join(self.tmp_dist, 'Ureka', 'misc', 'name')
        with open(misc, 'w+') as fp:
            fp.write(self.ureka['name'] + os.linesep)

    def _upgrade(self, src, dest):
        print("+ Upgrade in progress ({} to {})...".format(dest['version'],
                                                               src['version']))
        command = 'rsync -a -u {} {}'.format(src.path,
                                             os.path.dirname(dest.path)).split()
        proc = subprocess.Popen(command)
        proc.wait()

        if proc.returncode:
            return False
        return True

    def _post(self, ur):
        print("+ Executing post-upgrade tasks...")
        ur_env = ur_getenv(ur.path)

        command = 'ur_normalize -n -i -x'.split()
        proc = subprocess.Popen(command, env=ur_env)
        proc.wait()

        if proc.returncode:
            return False
        return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Ureka Upgrade Utility')
    parser.add_argument('--list-available', action='store_true', help='List available releases')
    parser.add_argument('--latest', action='store_true', 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')
    parser.add_argument('--no-backup', action='store_true', help='Do not backup existing installation')
    parser.add_argument('--backup-dir', action='store_true', help='Alternative backup storage location')
    parser.add_argument('--force', action='store_true', help='Ignore version checking')

    if 'UR_DIR' in os.environ:
        args = parser.parse_args()
        args.UR_DIR = os.environ['UR_DIR']
    else:
        parser.add_argument('UR_DIR',
                            action='store',
                            help='Absolute path to Ureka installation')
    args = parser.parse_args()

    storage_requires = (1024 ** 2) * 7  # 7GB
    storage_areas = [tempfile.gettempdir(),
                     os.path.abspath(os.curdir),
                     os.environ['HOME']]

    mirror = DEFAULT_MIRROR
    mirror_releases = DEFAULT_PUBLIC_RELEASE

    if args.mirror:
        mirror = args.mirror
        mirror_releases = "{}/public_releases.txt".format(mirror)

    if args.list_available:
        ureka = Ureka(args.UR_DIR)
        for release in ur_get_public_releases(mirror_releases):
            if StrictVersion(ureka['version']) < StrictVersion(release):
                flag = 'Available for upgrade'
            elif StrictVersion(ureka['version']) == StrictVersion(release):
                flag = 'Currently installed'
            else:
                flag = 'Older release'
            print('{:>6s} - {}'.format(release, flag))
        exit(0)

    if args.latest and args.request:
        print('--latest and --request are mutually exclusive options.')
        exit(1)

    if args.latest and not args.request:
        try:
            request = ur_get_public_releases(DEFAULT_PUBLIC_RELEASE)[-1:][-1]
        except:
            request = ''

        if not request:
            print('Unable to retrieve the latest public release from {}.'.format(DEFAULT_PUBLIC_RELEASE))
            exit(1)

    if not args.latest:
        if not args.request:
            print("No version requested. Use --request (e.g. 1.4.1)")
            exit(1)

    upgrade = Upgrade(args.UR_DIR, request, mirror)

    if args.force:
        upgrade.force = True

    if args.backup_dir:
        if not os.path.exists(args.backup_dir):
            print("Backup directory {} does not exist.".format(args.backup_dir))
            exit(1)

        upgrade.backup_path = os.path.abspath(args.backup_dir)

    if args.no_backup:
        upgrade.backup = False

    exit_code = upgrade.run()

    exit(exit_code)