path: root/linaro-cp.py

#!/usr/bin/env python

import argparse
import atexit
import cStringIO
import os
import mimetypes
import pycurl
import sys
import tempfile
import time
import re
import pprint

# Public artifacts BUILD-INFO.txt
build_info = 'Format-Version: 0.5\n\nFiles-Pattern: *\nLicense-Type: open\n'


class API_v1(object):
    def __init__(self, server, build_info, api_key):
        self.server = server
        self.api_base = server
        self.build_info = build_info
        self.api_key = api_key
        self.curl = pycurl.Curl()

    def __del__(self):
        self.curl.close()

    def _upload_data(self, url, data, headers=None, retry_count=3):
        response = cStringIO.StringIO()
        self.curl.setopt(pycurl.URL, url)
        self.curl.setopt(pycurl.HTTPPOST, data)
        self.curl.setopt(pycurl.WRITEFUNCTION, response.write)
        if headers:
            self.curl.setopt(pycurl.HTTPHEADER, headers)
        try:
            self.curl.perform()
        except Exception as e:
            if retry_count > 0:
                # server could be reloading or something. give it a second and
                # try again
                print('Upload failed for %s, retrying in 2 seconds' % url)
                time.sleep(2)
                return self._upload_data(url, data, headers, retry_count - 1)
            else:
                return repr(e)

        code = self.curl.getinfo(pycurl.RESPONSE_CODE)
        if code == 503 and retry_count > 0:
                print('503 failure for %s, retrying in 2 seconds' % url)
                time.sleep(2)
                return self._upload_data(url, data, headers, retry_count - 1)
        if code not in (200, 201):
            return 'HTTP_%d: %s' % (code, response.getvalue())

    def upload_file(self, url, filename):
        data = [
            ('file', (pycurl.FORM_FILE, filename)),
            ('key', (pycurl.FORM_CONTENTS, self.api_key)),
        ]
        return self._upload_data(url, data)

    def upload_transfer_queue(self, transfer_queue):
        transfer_failures = []
        for transfer_item in transfer_queue:
            failure = self.upload_file(
                transfer_item, transfer_queue[transfer_item])
            if failure:
                transfer_failures.append('%s: %s' % (transfer_item, failure))
        return transfer_failures

    def get_transfer_queue(self, src, dst, options):
        if self.api_base[-1] != '/' and dst[0] != '/':
            # one of these needs to be a slash to produce a url
            dst = '/' + dst

        transfer_queue = {}
        src_dir = os.path.abspath(src)

        if not os.path.isdir(src_dir):
            transfer_queue[self.api_base + dst + "/" + os.path.basename(src)] = src_dir
            return transfer_queue

        manifest = None
        if options.manifest:
            manifest = tempfile.NamedTemporaryFile(prefix="MANIFEST", delete=False)

        for root, dirs, files in os.walk(src_dir):
            rel_dir = root[len(src_dir) + 1:]
            for f in files:
                rel_path = os.path.join(rel_dir, f)
                # If there's at least one include pattern, skip unless
                # some pattern matches
                skip = len(options.include) > 0
                for pat in options.include:
                    if re.search(pat, rel_path):
                        skip = False
                        break
                if not skip:
                    dst_file = '%s%s/%s' % (self.api_base, dst, rel_path)
                    transfer_queue[dst_file] = os.path.join(root, f)
                    if manifest:
                        manifest.write(rel_path + "\n")
            if not options.no_build_info:
                if rel_dir and rel_dir[-1] != '/':
                    rel_dir += '/'
                build_info_file = os.path.join(root, 'BUILD-INFO.txt')
                if not os.path.exists(build_info_file):
                    dst_file = '%s%s/%s%s' % (
                        self.api_base, dst, rel_dir, 'BUILD-INFO.txt')
                    transfer_queue[dst_file] = self.build_info
        if manifest:
            manifest.close()
            dst_file = '%s%s/%s' % (
                self.api_base, dst, 'MANIFEST')
            transfer_queue[dst_file] = manifest.name
        return transfer_queue

    def upload(self, src, dst, options):
        transfer_queue = self.get_transfer_queue(src, dst, options)
        if not transfer_queue:
            print "Warning: no files to publish"
        if options.verbose:
            pprint.pprint(transfer_queue)
        if not options.dry_run:
            return self.upload_transfer_queue(transfer_queue)
        else:
            return []


class API_v2(API_v1):
    def __init__(self, server, build_info, api_key):
        super(API_v2, self).__init__(server, build_info, api_key)
        self.api_base = server + '/api/v2/publish/'

    def upload_file(self, url, filename):
        headers = ['AuthToken: ' + self.api_key]
        data = [('file', (pycurl.FORM_FILE, filename))]
        return self._upload_data(url, data, headers)

    def link_latest(self, dst):
        headers = ['AuthToken: ' + self.api_key]
        url = self.server + '/api/v2/link_latest/' + dst
        # pycurl requires data to be passed, or it will do an
        # HTTP GET even though we said to POST
        return self._upload_data(url, [('foo', 'bar')], headers)


class API_v3(API_v1):
    def __init__(self, server, build_info, api_key):
        super(API_v3, self).__init__(server, build_info, api_key)
        self.api_base = server + '/api/v3/publish/'

    def _upload_data(self, url, data, headers=None, retry_count=3):
        self.last_headers = cStringIO.StringIO()
        self.curl.setopt(pycurl.HEADERFUNCTION, self.last_headers.write)
        return super(API_v3, self)._upload_data(
            url, data, headers, retry_count)

    def _put_s3(self, url, filename, mtype, retry_count=3):
        response = cStringIO.StringIO()
        size = os.path.getsize(filename)
        headers = ['Content-Type: ' + mtype]
        c = pycurl.Curl()
        c.setopt(pycurl.URL, url)
        c.setopt(pycurl.HTTPHEADER, headers)
        c.setopt(pycurl.INFILESIZE, size)
        c.setopt(pycurl.PUT, 1)
        c.setopt(pycurl.WRITEFUNCTION, response.write)
        try:
            with open(filename, 'rb') as f:
                c.setopt(pycurl.INFILE, f)
                c.perform()
        except Exception as e:
            if retry_count > 0:
                # server could be reloading or something. give it a second and
                # try again
                print('Upload failed for %s, retrying in 2 seconds' % url)
                time.sleep(2)
                return self._put_s3(url, filename, mtype, retry_count - 1)
            else:
                return repr(e)
        code = c.getinfo(pycurl.RESPONSE_CODE)
        if code == 503 and retry_count > 0:
                print('503 failure for %s, retrying in 2 seconds' % url)
                time.sleep(2)
                return self._put_s3(url, filename, mtype, retry_count - 1)
        if code not in (200, 201):
            return response.getvalue()

    def upload_file(self, url, filename):
        # ask llp for an s3 tempurl:
        mtype = mimetypes.guess_type(filename)[0]
        if not mtype:
            mtype = 'other'

        headers = ['AuthToken: ' + self.api_key]
        code = self._upload_data(url, [('Content-Type', mtype)], headers)
        if code:
            return code

        # now find the tempurl and really publish
        for header in self.last_headers.getvalue().split('\n'):
            if header.startswith('Location:'):
                location = header[9:].strip()
                return self._put_s3(location, filename, mtype)
        raise RuntimeError('l-l-p response missing s3 location')

    def link_latest(self, dst):
        headers = ['AuthToken: ' + self.api_key]
        url = self.server + '/api/v3/link_latest/' + dst
        # pycurl requires data to be passed, or it will do an
        # HTTP GET even though we said to POST
        return self._upload_data(url, [('foo', 'bar')], headers)


def main():
    parser = argparse.ArgumentParser(
        description='Copy file(s) from source to destination')
    parser.add_argument('-k', '--key', help='key used for the copy')
    parser.add_argument('-a', '--api_version', choices=('1', '2', '3'), default='1',
                        help='API version to use. default=%(default)s')
    parser.add_argument('--server', default='http://snapshots.linaro.org/',
                        help='Publishing API server. default=%(default)s')
    parser.add_argument('-i', '--include', action='append', default=[],
                        help='Include regex for files')
    parser.add_argument('--no-build-info', action='store_true',
                        help="Don't auto-generate BUILD-INFO.txt")
    parser.add_argument('-b', '--build-info', help='Custom build-info file')
    parser.add_argument('--link-latest', action='store_true',
                        help='''Create symlink for "latest" to point to this
                             build.''')
    parser.add_argument('--split-job-owner', action='store_true',
                        help='Split Jenkins job owner in dst (owner_job -> ~owner/job)')
    parser.add_argument('--manifest', action='store_true',
                        help='Generate MANIFEST file with list of all files published')
    parser.add_argument('--verbose', action='store_true',
                        help='Verbose operation')
    parser.add_argument('--dry-run', action='store_true',
                        help="Don't actually publish files")
    g = parser.add_mutually_exclusive_group(required=True)
    g.add_argument('--make-link', action='store_true', required=False,
                   help='Don\'t publish files, just create "latest"')
    g.add_argument('src', nargs='?', help='source directory with files to publish')
    parser.add_argument('dst', help='destination to publish to')

    arguments = parser.parse_args()
    # Publish key is required. Fallback to PUBLISH_KEY environment
    # variable when it isn't passed as an argument
    if arguments.key:
        key = arguments.key
    else:
        key = os.environ.get('PUBLISH_TOKEN')
        if key:
            if arguments.api_version == '1':
                # Default to api v2 if not specified
                arguments.api_version = '2'
        else:
            key = os.environ.get('PUBLISH_KEY')
            if key is None:
                sys.exit('Error: Key is not defined.')

    if not arguments.build_info:
        fd, arguments.build_info = tempfile.mkstemp(prefix='BUILD-INFO.txt')
        atexit.register(os.unlink, arguments.build_info)
        os.write(fd, build_info)

    cls = globals()['API_v' + arguments.api_version]
    api = cls(arguments.server, arguments.build_info, key)

    if arguments.split_job_owner:
        # Rewrite job path .../owner_jobname/123 -> .../~owner/jobname/123 ,
        # as required for android-build publishing.
        arguments.dst = re.sub(r"^(.+?)/([^/]+?)_([^/]+?)/([0-9]+)/?$", r"\1/~\2/\3/\4/", arguments.dst)

    transfer_failures = []
    if not arguments.make_link:
        transfer_failures = api.upload(arguments.src, arguments.dst, arguments)

    if arguments.link_latest or arguments.make_link:
        err = api.link_latest(arguments.dst)
        if err:
            transfer_failures.append('unable to create symlink: ' + err)

    if len(transfer_failures) > 0:
        sys.exit('Error: Failed to transfer:\n  %s' % '\n  '.join(transfer_failures))


if __name__ == '__main__':
    main()