fetch-stage3: Rewrite without mkvm

fetch-stage3.py

@@ -1,11 +1,197 @@
 #!/usr/bin/env python
-from mkvm import stage3
 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
 
 
-def parse_args():
+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',
@@ -16,24 +202,25 @@ def parse_args():
         help='stage3 subtype/profile (e.g. nomultilib)'
     )
     parser.add_argument(
-        '--cache-dir',
+        'dest',
-        metavar='PATH',
+        nargs='?',
+        default='.',
         help='Cache location for fetched stage3 tarballs'
     )
     return parser.parse_args()
 
 
 def main():
-    args = parse_args()
+    args = _parse_args()
 
     try:
-        fetcher = stage3.Fetcher(args.cache_dir)
+        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, stage3.FetchError, stage3.VerifyError) as e:
+    except (OSError, FetchError, VerifyError) as e:
         sys.stderr.write('Failed to fetch stage3 tarball: {}\n'.format(e))
         raise SystemExit(os.EX_UNAVAILABLE)