scripts/fetch-stage3.py

#!/usr/bin/env python
import argparse
import codecs
import gpgme
import logging
import hashlib
import io
import os
import shutil
import subprocess
import sys
import time
import urllib.parse
import urllib.request


log = logging.getLogger('stage3')


XDG_CACHE_DIR = os.environ.get('XDG_CACHE_DIR', '~/.cache')

ARCH_NAMES = {
    'i386': 'x86',
    'i486': 'x86',
    'i586': 'x86',
    'i686': 'x86',
    'x86_64': 'amd64',
    'em64t': 'amd64',
}


class FetchError(Exception):
    pass


class VerifyError(Exception):
    pass


class Fetcher(object):

    DEFAULT_MIRROR = 'http://distfiles.gentoo.org/'
    LIST_CACHE_LIFE = 86400

    log = log.getChild('fetch')

    def __init__(self, cache_dir=None, mirror=None):
        if mirror is None:
            mirror = os.environ.get('GENTOO_MIRROR', self.DEFAULT_MIRROR)
        if cache_dir is None:
            self.cache_dir = os.path.join(
                os.path.expanduser(XDG_CACHE_DIR),
                'stage3s',
            )
        else:
            self.cache_dir = cache_dir
        if not os.path.isdir(self.cache_dir):
            os.makedirs(self.cache_dir)
        self.mirror = mirror
        if not self.mirror.endswith('/'):
            self.mirror += '/'

    @staticmethod
    def verify(filename):
        log.debug('Verifying PGP signature for {}'.format(filename))
        ctx = gpgme.Context()
        plaintext = io.BytesIO()
        with open(filename, 'rb') as f:
            sigs = ctx.verify(f, None, plaintext)
        for sig in sigs:
            if sig.status:
                raise VerifyError(sig.status.args[2])
            if sig.wrong_key_usage:
                raise VerifyError('wrong key usage')
        log.info('Successfully verified PGP signature')
        plaintext.seek(0)
        buf = codecs.getreader('utf-8')(plaintext)
        dirname = os.path.dirname(filename)
        for line in buf:
            if not line.lstrip().startswith('#'):
                continue
            if 'SHA512' in line:
                h = hashlib.sha512()
            else:
                continue
            line = buf.readline()
            try:
                digest, filename = line.split()
            except ValueError:
                pass
            path = os.path.join(dirname, filename)
            log.debug('Verifying checksum of {}'.format(path))
            with open(path, 'rb') as f:
                for data in iter(lambda: f.read(4096), b''):
                    h.update(data)
            if h.hexdigest() != digest.lower():
                raise VerifyError(
                    '{} checksum mismatch: {}'.format(h.name, filename))
            log.info('Verified checksum of {}'.format(filename))

    def wget(self, *uris):
        cmd = ['wget', '--continue']
        cmd += uris
        log.debug('Running command: {}'.format(' '.join(cmd)))
        try:
            p = subprocess.Popen(cmd, cwd=self.cache_dir, stderr=sys.stderr)
        except OSError as e:
            raise FetchError('Failed to run wget: {}'.format(e))
        if p.wait() != 0:
            raise FetchError('wget returned status {}'.format(p.returncode))

    def fetch_stage(self, arch=None, subtype=None):
        if not arch:
            arch = os.uname().machine
        try:
            arch = ARCH_NAMES[arch]
        except KeyError:
            pass

        want = 'stage3-{}-'.format(subtype if subtype else arch)
        with self._get_latest_list(arch) as latest_list:
            for line in latest_list:
                line = line.split('#')[0]
                if not line:
                    continue
                try:
                    path, size = line.split(None, 1)
                except ValueError:
                    log.warning('Unexpected value: {}'.format(line))
                    continue
                filename = os.path.basename(path)
                if filename.startswith(want):
                    break
            else:
                raise FetchError(
                    'No stage3 tarballs for {}'.format(subtype or arch))
        log.info('Found latest stage3 tarball: {}'.format(filename))
        full_path = 'releases/{}/autobuilds/{}'.format(arch, path)
        uri = urllib.parse.urljoin(self.mirror, full_path)
        local_path = os.path.join(self.cache_dir, filename)
        to_fetch = [
            uri,
            uri + '.CONTENTS',
            uri + '.DIGESTS.asc',
        ]
        try:
            st = os.stat(local_path)
        except OSError:
            pass
        else:
            if st.st_size == int(size):
                log.info('Cached copy of {} is complete'.format(filename))
                to_fetch.remove(uri)
        for fn in to_fetch[-2:]:
            c_fn = os.path.join(self.cache_dir, os.path.basename(fn))
            try:
                st = os.stat(c_fn)
            except OSError:
                pass
            else:
                if st.st_size > 0:
                    to_fetch.remove(fn)
        if to_fetch:
            self.wget(*to_fetch)
        return local_path

    def _get_latest_list(self, arch):
        cache_fname = os.path.join(
            self.cache_dir,
            'latest-stage3-{}.txt'.format(arch),
        )
        try:
            st = os.stat(cache_fname)
        except OSError:
            pass
        else:
            if st.st_mtime > time.time() - self.LIST_CACHE_LIFE:
                return open(cache_fname)

        path = 'releases/{}/autobuilds/latest-stage3.txt'.format(arch)
        url = urllib.parse.urljoin(self.mirror, path)
        log.debug('Fetching URL: {}'.format(url))
        try:
            response = urllib.request.urlopen(url)
        except urllib.error.HTTPError as e:
            log.error('Failed to fetch latest stage3 list: {}'.format(e))
            raise FetchError('Could not fetch latest stage3 list')
        with open(cache_fname, 'wb') as f:
            for line in response:
                f.write(line)
        return open(cache_fname)


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--arch',
        help='Architecture to use when fetching a stage3 tarball'
    )
    parser.add_argument(
        '--subtype',
        help='stage3 subtype/profile (e.g. nomultilib)'
    )
    parser.add_argument(
        'dest',
        nargs='?',
        default='.',
        help='Cache location for fetched stage3 tarballs'
    )
    return parser.parse_args()


def main():
    args = _parse_args()

    try:
        fetcher = Fetcher(args.dest)
        stagetbz = fetcher.fetch_stage(args.arch, args.subtype)
        fetcher.verify(stagetbz + '.DIGESTS.asc')
    except KeyboardInterrupt:
        print('')
        raise SystemExit(os.EX_TEMPFAIL)
    except (OSError, FetchError, VerifyError) as e:
        sys.stderr.write('Failed to fetch stage3 tarball: {}\n'.format(e))
        raise SystemExit(os.EX_UNAVAILABLE)


if __name__ == '__main__':
    main()