sign-release.py

#!/usr/bin/python
import tempfile, os, shutil, time, bz2, sys, re
import logging
from subprocess import *

from signing import *

log = logging.getLogger()

def signfile(filename, keydir, fake=False):
    """Sign the given file with keys in keydir.

    If fake is True, then don't actually sign anything, just sleep for a
    second to simulate signing time."""
    if fake:
        time.sleep(1)
        return
    basename = os.path.basename(filename)
    dirname = os.path.dirname(filename)
    stdout = tempfile.TemporaryFile()
    command = ['signcode',
        '-spc', '%s/MozAuthenticode.spc' % keydir,
        '-v', '%s/MozAuthenticode.pvk' % keydir,
        '-t', 'http://timestamp.verisign.com/scripts/timestamp.dll',
        '-i', 'http://www.mozilla.com',
        # Try 5 times, and wait 60 seconds between tries
        '-tr', '5',
        '-tw', '60',
        basename]
    try:
        check_call(command, cwd=dirname, stdout=stdout, stderr=STDOUT)
        stdout.seek(0)
        data = stdout.read()
        # Make sure that the command output "Succeeded".  Sometimes signcode
        # returns with 0, but doesn't output "Succeeded", which in the past has
        # meant that the file has been signed, but is missing a timestmap.
        if data.strip() != "Succeeded":
            raise ValueError("signcode didn't report success")
    except:
        stdout.seek(0)
        data = stdout.read()
        log.exception(data)
        raise

def _signLocale(obj, dstdir, files, remember=None):
    # Helper function to sign one individual locale, and keep track of some
    # stats
    nFiles, cacheHits, nSigned = 0, 0, 0
    for f in files:
        ext = os.path.splitext(f)[1]
        if ext == '.mar':
            compressed = True
            # We normally sign the .mar file after the .exe, so there's no point
            # in caching the results
            if remember is None:
                remember = False
        else:
            compressed = False
            if remember is None:
                remember = True
        results = obj.signPackage(f, dstdir, remember, compressed)
        if not results:
            return results
        nFiles += results[0]
        cacheHits += results[1]
        nSigned += results[2]
    return nFiles, cacheHits, nSigned

class Signer:
    def __init__(self, keydir, concurrency=1, keepCache=False, fake=False,
            unsignedInstallers=False):
        """
        `keydir`    - where MozAuthenticode.svc,.pvk can be found

        `concurrency` - How many processes to use for signing.  If > 1, then
            that many children processes will be spawned to sign locales in
            parallel, and the main process will remain mostly idle to manage the
            children.

        `keepCache` - If False, then our cache of signed files is deleted as
            soon as this object is created

        `fake`      - If True, then no signing is actually done.

        `unsignedInstallers` - If True, then don't sign the final installer .exe's
        """

        # What directory to use to store our cached files
        self.cacheDir = 'cache'
        self.keydir = keydir
        self.fake = fake
        self.unsignedInstallers = unsignedInstallers
        self.concurrency = concurrency

        # Blow away our cache if we're not keeping it
        if not keepCache:
            log.debug("Purging cache")
            if os.path.exists(self.cacheDir):
                shutil.rmtree(self.cacheDir)

    def rememberFile(self, hsh, filename):
        """Remember this file in our cache.

        `filename` will be copied into our cache identified by `hsh`
        
        Usually `hsh` is the sha1sum of the unsigned file, and `filename`
        points to the signed file.
        """
        dstfile = os.path.join(self.cacheDir, hsh, os.path.basename(filename))
        dstdir = os.path.dirname(dstfile)
        if not os.path.exists(dstdir):
            os.makedirs(dstdir, 0755)
        # Copying files isn't atomic, so copy to a temporary file first, and
        # then rename to the final destination when we're done copying
        copyfile(filename, dstfile+".tmp")
        os.rename(dstfile+".tmp", dstfile)

    def getFile(self, hsh, filename):
        """Returns the path to a file in our cache identified by `hsh`, and that is named `filename`

        Returns None if there is no such file in our cache
        """
        dstfile = os.path.join(self.cacheDir, hsh, os.path.basename(filename))
        if os.path.exists(dstfile):
            return dstfile
        return None

    def signPackage(self, pkgfile, dstdir, remember=False, compressed=False):
        """Sign `pkgfile`, putting the results into `dstdir`.

        If `remember` is True, then cache the newly signed files into our
        cache.

        If `compressed` is True, then the contents of pkgfile are bz2
        compressed (e.g. in a mar file), and should be decompressed before
        signing.
        """
        log.info("Processing %s", pkgfile)
        basename = os.path.basename(pkgfile)
        dstfile = convertPath(pkgfile, dstdir)

        # Keep track of our output in a list here, and we can output everything
        # when we're done This is to avoid interleaving the output from
        # multiple processes.
        logs = []
        logs.append("Repacking %s to %s" % (pkgfile, dstfile))
        parentdir = os.path.dirname(dstfile)
        if not os.path.exists(parentdir):
            os.makedirs(parentdir, 0755)

        nFiles = 0
        cacheHits = 0
        nSigned = 0
        tmpdir = tempfile.mkdtemp()
        try:
            # Unpack it
            logs.append("Unpacking %s to %s" % (pkgfile, tmpdir))
            unpackfile(pkgfile, tmpdir)
            # Swap in files we have already signed
            for f in findfiles(tmpdir):
                # Possible future optimization: only hash files we know might
                # change, like .exe's and .dll's
                h = sha1sum(f)
                basename = os.path.basename(f)
                nFiles += 1

                # Look in the cache for another file with the same original hash
                cachedFile = self.getFile(h, f)
                if cachedFile:
                    cacheHits += 1
                    assert os.path.basename(cachedFile) == basename
                    logs.append("Copying %s from %s" % (basename, cachedFile))
                    # Preserve the original file's mode; don't use the cached mode
                    # We usually process installer .exe's first, and 7z doesn't
                    # preserve the file mode, so the cached copies of the files
                    # are mode 0666.  In the mar files, executables have mode
                    # 0777, so we want to preserve that.
                    copyfile(cachedFile, f, copymode=False)
                elif shouldSign(f):
                    # We need to sign this file
                    # If this file is compressed, check if we have a cached copy that
                    # is uncompressed
                    if compressed:
                        bunzip2(f)
                        h2 = sha1sum(f)
                        cachedFile = self.getFile(h2, f)
                        if cachedFile:
                            # We have a cached copy of this file that is uncompressed.
                            # So copy it into our dstdir, and recompress it, and
                            # save it for future use.
                            cacheHits += 1
                            assert os.path.basename(cachedFile) == basename
                            logs.append("Copying %s from uncompressed %s" % (basename, cachedFile))
                            # See note above about not copying the file's mode
                            copyfile(cachedFile, f, copymode=False)
                            bzip2(f)
                            if remember:
                                logs.append("Caching compressed %s as %s" % (f, h))
                                self.rememberFile(h, f)
                            continue

                    nSigned += 1
                    logs.append("Signing %s" % f)
                    signfile(f, self.keydir, self.fake)
                    if compressed:
                        bzip2(f)
                    if remember:
                        logs.append("Caching %s as %s" % (f, h))
                        self.rememberFile(h, f)

            # Repack it
            logs.append("Packing %s" % dstfile)
            packfile(dstfile, tmpdir)
            # Sign installer
            if dstfile.endswith('.exe') and not self.unsignedInstallers:
                logs.append("Signing %s" % dstfile)
                signfile(dstfile, self.keydir, self.fake)
            return nFiles, cacheHits, nSigned
        except:
            log.exception("Error signing %s", pkgfile)
            return False
        finally:
            # Clean up after ourselves, and output our logs
            shutil.rmtree(tmpdir)
            log.info("\n  ".join(logs))

    def signFiles(self, files, dstdir, product, firstLocale='en-US', firstLocaleSigned=False):
        """Sign `files`, putting the results into `dstdir`.  If `files` has a
        single entry, and refers to a directory, then all files under that
        directory are signed.

        `product` is the product name, e.g. 'firefox', and is used to filename
        pattern matching

        `firstLocale` is the first locale that should be signed, and will have
        the results cached.  Other locales will use the cached versions of
        these signed files where appropriate.

        If `firstLocaleSigned` is True, then all `firstLocale` files must
        already be signed.  The unsigned and signed copies of `firstLocale`
        will be unpacked, and the results cached so that other locales will use
        the signed copies of `firstLocale` files where appropriate.
        """
        start = time.time()

        if len(files) == 1 and os.path.isdir(files[0]):
            files = findfiles(files[0])

        # Filter out files that we can't sign
        # Right now this means that anything that isn't a win32 .exe or .mar file gets filtered out
        for f in files[:]:
            skip = False
            try:
                info = fileInfo(f, product)
                if info['platform'] != 'win32':
                    skip = True
                if info['contents'] not in ('complete', 'installer'):
                    skip = True
            except ValueError:
                skip = True

            if skip:
                files.remove(f)
                if 'win32' in f:
                    if 'xpi' not in f:
                        log.info("Skipping %s", f)

        # Utility function for sorting files
        # Makes sure that .mar files follow their corresponding .exe files
        def fileKey(f):
            info = fileInfo(f, product)
            locale = info['locale']
            if locale == firstLocale:
                localeVal = 0
            else:
                localeVal = 1
            if info['format'] == 'exe':
                exeVal = 0
            else:
                exeVal = 1
            return (localeVal, locale, exeVal, f)

        files.sort(key=fileKey)
        nfiles = len(files)

        # Split the files up into locales
        locales  = {}
        for f in files:
            info = fileInfo(f, product)
            locale = info['locale']
            if not locale in locales:
                locales[locale] = []
            locales[locale].append(f)

        # This is required to be global because localeFinished is called via a
        # callback when child processes finish signing locales, and the
        # callback doesn't have access to this local scope
        global doneLocales
        doneLocales = 0

        # pool_start records when the pool of children has started processing
        # files.  We initialize it to start here because we haven't started the
        # pool yet!
        pool_start = start

        # This is called when we finish signing a locale.  We update some
        # stats, and report on our progress
        def localeFinished(locale):
            global doneLocales
            doneLocales += 1
            now = int(time.time())
            n = len(locales)
            t_per_locale = (now - pool_start) / doneLocales
            remaining = (n-doneLocales) * t_per_locale
            eta_s = remaining % 60
            eta_m = (remaining / 60) % 60
            eta_h = (remaining / 3600)
            percent = 100.0 * doneLocales / float(n)
            eta_time = time.strftime("%H:%M:%S", time.localtime(int(now + remaining)))
            log.info("%i/%i (%.2f%%) complete, ETA in %i:%02i:%02i at %s", doneLocales, n, percent, eta_h, eta_m, eta_s, eta_time)

        # Sign the firstLocale first.  If we don't have any firstLocale files,
        # then something is wrong!
        if firstLocale not in locales:
            log.error("No files found with locale %s", firstLocale)
            sys.exit(1)

        if not firstLocaleSigned:
            if not _signLocale(self, dstdir, locales[firstLocale], remember=True):
                log.error("Error signing %s", firstLocale)
                sys.exit(1)
        else:
            # If the first locale is already signed, we need to unpack both the
            # signed and unsigned copies, and then make sure the signed
            # versions are cached
            for uf in locales[firstLocale]:
                sf = convertPath(uf, dstdir)
                if not os.path.exists(sf):
                    log.error("Signed version of %s doesn't exist as expected at %s", uf, sf)
                    sys.exit(1)

                unsigned_dir = tempfile.mkdtemp()
                signed_dir = tempfile.mkdtemp()
                try:
                    log.info("Unpacking %s into %s", uf, unsigned_dir)
                    unpackfile(uf, unsigned_dir)
                    log.info("Unpacking %s into %s", sf, signed_dir)
                    unpackfile(sf, signed_dir)
                    for f in findfiles(unsigned_dir):
                        # Calculate the hash of the original, unsigned file
                        h = sha1sum(f)
                        sf = signed_dir + f[len(unsigned_dir):]
                        # Cache the signed version
                        log.info("Caching %s as %s" % (f, h))
                        self.rememberFile(h, sf)
                finally:
                    shutil.rmtree(unsigned_dir)
                    shutil.rmtree(signed_dir)
        localeFinished(firstLocale)

        results = [0,0,0]
        # We're going to start our pool of children for processing the other
        # locales, so reset our pool_start time to the current time.
        # We do this because the time to sign the firstLocale isn't
        # representative of the time to sign future locales, due to caching and
        # parallelization.  Using pool_start gives more accurate ETA calculations.
        pool_start = time.time()

        # Setup before signing
        # We define different addLocale functions depending on if we're signing
        # things concurrently or serially
        if self.concurrency > 1:
            # If we're doing signing in parallel, start up our pool of children
            # using the processing module
            import processing.pool
            pool = processing.pool.Pool(self.concurrency)
            result_list = []
            def addLocale(locale):
                def cb(r):
                    if not r:
                        log.error("Signing locale %s failed", locale)
                        return

                    for i in range(3):
                        results[i] += r[i]
                    localeFinished(locale)

                files = locales[locale]
                r = pool.apply_async(_signLocale, (self, dstdir, files), callback=cb)
                result_list.append(r)
        else:
            pool = None
            def addLocale(locale):
                files = locales[locale]
                r = _signLocale(self, dstdir, files)
                if not r:
                    log.error("Signing locale %s failed", locale)
                    sys.exit(1)
                localeFinished(locale)
                for i in range(3):
                    results[i] += r[i]

        # Now go through our locales and call addLocale to start processing
        # them
        for locale in sorted(locales.keys()):
            if locale == firstLocale:
                continue
            addLocale(locale)

        error = False
        if pool:
            # Clean up the pool of child processes afterwards
            for r in result_list:
                try:
                    if not r.get():
                        log.error("Signing locale %s failed", locale)
                        error = True
                        pool.terminate()
                        break
                except:
                    log.exception("Error")
                    pool.terminate()
                    sys.exit(1)
            try:
                pool.close()
                pool.join()
            except:
                log.exception("Error")
                pool.terminate()
                raise

        # Report on our stats
        end = time.time()
        nTotalFiles, cacheHits, nSigned = results
        if nTotalFiles > 0:
            hitrate = cacheHits / float(nTotalFiles) * 100.0
        else:
            hitrate = 0
        if nfiles > 0:
            time_per_file = (end-start)/float(nfiles)
        else:
            time_per_file = 0
        log.info("%.2f%% hit rate, %i files signed", hitrate, nSigned)
        log.info("%s files repacked in %.0f seconds (%.2f seconds per file)", nfiles, end-start, time_per_file)

        if error:
            sys.exit(1)

if __name__ == "__main__":
    from optparse import OptionParser
    parser = OptionParser()
    parser.set_defaults(
            dest=None,
            first_locale="en-US",
            concurrency=1,
            keep_cache=False,
            keydir="e:/2008-keys",
            fake=False,
            prev=False,
            product="firefox",
            unsigned_installers=False,
            )
    parser.add_option("-o", "--dest", dest="dest", help="destination directory")
    parser.add_option("", "--first-locale", dest="first_locale", help="first locale to sign")
    parser.add_option("-j", "--concurrency", dest="concurrency", help="concurrency", type="int")
    parser.add_option("", "--keep-cache", dest="keep_cache", action="store_true", help="don't delete cache of signed files before starting")
    parser.add_option("", "--keydir", dest="keydir", help="directory where signing keys are located")
    parser.add_option("", "--fake", dest="fake", action="store_true", help="don't actually sign anything")
    parser.add_option("-p", "--previously-signed", dest="prev", action="store_true", help="use a previously signed first locale")
    parser.add_option("", "--product", dest="product", help="product name")
    parser.add_option("", "--unsigned-installers", dest="unsigned_installers", action="store_true", help="don't sign the installer .exes")

    options, args = parser.parse_args()
    if not options.dest:
        parser.error("Must specify a destination directory")

    try:
        checkTools()
    except OSError,e:
        parser.error(e.message)

    if not os.path.exists('checkouts/stubs'):
        parser.error("No checkouts/stubs directory found")

    if options.concurrency == 1:
        logging.basicConfig(level=logging.INFO, format="%(message)s")
    else:
        logging.basicConfig(level=logging.INFO, format="[%(process)d] %(message)s")

    s = Signer(options.keydir, options.concurrency, options.keep_cache, options.fake, options.unsigned_installers)
    s.signFiles(args, options.dest, options.product, firstLocale=options.first_locale, firstLocaleSigned=options.prev)