appengine_update_sdk.py

# -*- coding: utf-8 -*-

import os
import sys
import argparse
import logging
import shutil
import tempfile
import urllib

import setuptools.archive_util
import yaml


logger = logging.getLogger(__name__)


def get_remote_version(server, runtime, local):
    return yaml.safe_load(
        server.Send('/api/updatecheck',
                    timeout=None,
                    runtime=runtime,
                    release=local['release'],
                    timestamp=local['timestamp'],
                    api_versions=local['api_versions']))


def get_local_version(root_path):
    path = os.path.join(root_path, 'google_appengine', 'VERSION')
    return yaml.safe_load(open(path).read())


def get_download_url(options, version):
    if options.runtime not in ('python', 'python27'):
        raise NotImplementedError

    basename = 'google_appengine_{release}.zip'.format(**version)
    return options.urlprefix + basename


def download(options, version, url):
    path = os.path.abspath(options.download_dir)

    if not (path and os.path.isdir(path)):
        raise ValueError("Invalid download dir %r" % path)

    handle, tmp_path = tempfile.mkstemp(prefix='gaesdk-')
    try:
        try:
            tmp_path, headers = urllib.urlretrieve(url, tmp_path)
            # TODO: support sha1sum check
        except IOError:
            os.remove(tmp_path)
            raise
        except Exception:
            os.remove(tmp_path)
            raise
    finally:
        os.close(handle)

    to_path = os.path.join(path, os.path.basename(url))
    shutil.move(tmp_path, to_path)
    return to_path


def extract_and_replace(options, version, archive_path):
    extract_dir = tempfile.mkdtemp(prefix='gaesdk-')
    try:
        setuptools.archive_util.unpack_archive(archive_path, extract_dir)
    except setuptools.archive_util.UnrecognizedFormat:
        raise RuntimeError(
            'Unable to extract the package: %s', archive_path)
    else:
        name = 'google_appengine'
        src = os.path.join(extract_dir, os.listdir(extract_dir)[0])
        if version['release'] not in src:
            destbase = os.path.basename(src) + '-' + version['release']
        else:
            destbase = os.path.basename(src)

        prefix = os.path.abspath(options.prefix)
        dest = os.path.join(prefix, destbase)
        shutil.move(src, dest)

        destlink = os.path.join(prefix, name)
        if os.path.exists(destlink):
            if os.path.islink(destlink):
                os.remove(destlink)
            else:
                local_version = get_local_version(prefix)
                backup_dest = destlink + '-' + local_version['release']
                shutil.move(destlink, backup_dest)

        os.symlink(dest, destlink)
        return dest, destlink

    finally:
        if extract_dir and extract_dir.startswith(tempfile.gettempdir()):
            shutil.rmtree(extract_dir)


def main():
    parser = argparse.ArgumentParser(description='GAE SDK updater')
    parser.add_argument(
        '-p', '--prefix', default='/usr/local',
        help='Install prefix (defaults to %(default)s).')
    parser.add_argument(
        '-d', '--download-dir', dest='download_dir', default='/usr/local/src',
        help='Download directory (defaults to %(default)s).')
    parser.add_argument(
        '-s', '--server', default='appengine.google.com',
        help='The App Engine server.')
    parser.add_argument(
        '--insecure', action='store_true', help=argparse.SUPPRESS,
        default=False)
    parser.add_argument(
        '--host', default=None,
        help='Overrides the Host header sent with all RPCs.')
    parser.add_argument(
        '--runtime', default='python27',
        help='The App Engine runtime (defaults to %(default)s).')
    parser.add_argument(
        '-u', '--url-prefix', dest='urlprefix',
        default='https://googleappengine.googlecode.com/files/')
    parser.add_argument('-v', '--verbose', action='store_true', default=False)

    args = parser.parse_args()


    logging.basicConfig(level=logging.INFO if args.verbose else logging.WARN)

    sys.path.append(os.path.join(args.prefix, 'google_appengine'))

    import dev_appserver
    dev_appserver.fix_sys_path()

    from google.appengine.tools import appcfg, appengine_rpc

    rpcserver = appengine_rpc.HttpRpcServer(
        args.server,
        lambda: ('unused_email', 'unused_password'),
        appcfg.GetUserAgent(),
        appcfg.GetSourceName(),
        host_override=args.host,
        secure=False if args.insecure else True)
    rpcserver.authenticated = True

    local_version = get_local_version(args.prefix)
    remote_version = get_remote_version(rpcserver, args.runtime, local_version)

    # TODO use sdk_update_checker ?
    if local_version['timestamp'] >= remote_version['timestamp']:
        logger.info('Already up-to-date.')
    else:
        logger.info('Found new version: %s', remote_version['release'])
        srcurl = get_download_url(args, remote_version)
        logger.info('Downloading %s ...', srcurl)
        src = download(args, remote_version, srcurl)
        res = extract_and_replace(args, remote_version, src)
        logger.info('Installed to %s (-> %s)', res[1], res[0])

    sys.exit(0)

if __name__ == '__main__':
    sys.exit(main())