createOSXinstallPkg

#!/usr/bin/env python
# encoding: utf-8
#
# Copyright 2012-2017 Greg Neagle.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
createOSXinstallPkg

Created by Greg Neagle on 2012-07-16.
Modified June 2013 for WWDC 2013
Modified June 2014 for WWDC 2014
Modified June 2015 for WWDC 2015
Modified June 2016 for WWDC 2016
Modified Aug 2016 for installs under SIP
Modified Apr 2017 to deal with fallout from 10.12.4
"""

import sys
import os

import optparse
import plistlib
import shutil
import subprocess
import tempfile

from distutils import version

from xml.dom import minidom
from xml.parsers.expat import ExpatError

# we use lots of camelCase-style names. Deal with it.
# pylint: disable=C0103

DEBUG = False
DEFAULT_INSTALLKBYTES = 8 * 1024 * 1024

def cleanUp():
    '''Cleanup our TMPDIR'''
    if TMPDIR:
        shutil.rmtree(TMPDIR, ignore_errors=True)


def fail(errmsg=''):
    '''Print any error message to stderr,
    clean up install data, and exit'''
    if errmsg:
        print >> sys.stderr, errmsg
    cleanUp()
    # exit
    exit(1)


def cleanupOutput(output_path):
    '''Attempt to clean up our output package/fake app
    if we fail during creation'''
    if not DEBUG:
        try:
            shutil.rmtree(output_path)
        except EnvironmentError:
            pass


# dmg helpers
def mountdmg(dmgpath, use_shadow=False):
    """
    Attempts to mount the dmg at dmgpath
    and returns a list of mountpoints
    If use_shadow is true, mount image with shadow file
    """
    mountpoints = []
    dmgname = os.path.basename(dmgpath)
    cmd = ['/usr/bin/hdiutil', 'attach', dmgpath,
           '-mountRandom', TMPDIR, '-nobrowse', '-plist',
           '-owners', 'on']
    if use_shadow:
        shadowname = dmgname + '.shadow'
        shadowpath = os.path.join(TMPDIR, shadowname)
        cmd.extend(['-shadow', shadowpath])
    else:
        shadowpath = None
    proc = subprocess.Popen(cmd, bufsize=-1,
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (pliststr, err) = proc.communicate()
    if proc.returncode:
        print >> sys.stderr, 'Error: "%s" while mounting %s.' % (err, dmgname)
    if pliststr:
        plist = plistlib.readPlistFromString(pliststr)
        for entity in plist['system-entities']:
            if 'mount-point' in entity:
                mountpoints.append(entity['mount-point'])

    return mountpoints, shadowpath


def unmountdmg(mountpoint):
    """
    Unmounts the dmg at mountpoint
    """
    proc = subprocess.Popen(['/usr/bin/hdiutil', 'detach', mountpoint],
                            bufsize=-1, stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    (dummy_output, err) = proc.communicate()
    if proc.returncode:
        print >> sys.stderr, 'Polite unmount failed: %s' % err
        print >> sys.stderr, 'Attempting to force unmount %s' % mountpoint
        # try forcing the unmount
        retcode = subprocess.call(['/usr/bin/hdiutil', 'detach', mountpoint,
                                   '-force'])
        if retcode:
            print >> sys.stderr, 'Failed to unmount %s' % mountpoint


def makePkgDirs(pkgpath):
    '''Makes the package directories for a package with
    full pathname "pkgpath"'''
    if os.path.exists(pkgpath):
        fail('Package exists at %s' % pkgpath)
    try:
        os.makedirs(os.path.join(pkgpath, 'Contents/Resources/en.lproj'))
        os.makedirs(os.path.join(pkgpath,
                                 'Contents/Resources/OS X Install Data'))
    except OSError, err:
        fail('Error creating package directories at %s: %s' % (pkgpath, err))


def makeDescriptionPlist(output_pkg_path, os_version='10.7', build_number=None):
    '''Writes a Resources/en.lproj/Description.plist
    for the OS to be installed'''
    pkg_description = {}
    major_version = '.'.join(os_version.split('.')[:2])
    titles = {'10.7':  'Mac OS X Lion',
              '10.8':  'OS X Mountain Lion',
              '10.9':  'OS X Mavericks',
              '10.10': 'OS X Yosemite',
              '10.11': 'OS X El Capitan',
              '10.12': 'macOS Sierra'}
    title = titles.get(major_version, major_version)
    description = ('Unattended custom install of %s version %s'
                   % (title, os_version))
    if build_number:
        description += ' build %s' % build_number
    pkg_description['IFPkgDescriptionDescription'] = description
    pkg_description['IFPkgDescriptionTitle'] = title
    output_file = os.path.join(output_pkg_path,
                               'Contents/Resources/en.lproj/Description.plist')
    try:
        plistlib.writePlist(pkg_description, output_file)
    except (OSError, ExpatError), err:
        cleanupOutput(output_pkg_path)
        fail('Error creating file at %s: %s' % (output_file, err))


def makeInfoPlist(output_pkg_path, os_version='10.7', build_number=None,
                  pkg_id=None, installKBytes=DEFAULT_INSTALLKBYTES):
    '''Creates Contents/Info.plist for package'''
    if not pkg_id:
        pkg_id = 'com.googlecode.munki.installosx.pkg'
    info = {
        'CFBundleIdentifier': pkg_id,
        'CFBundleShortVersionString': str(os_version),
        'IFMajorVersion': 1,
        'IFMinorVersion': 0,
        'IFPkgFlagDefaultLocation': '/tmp',
        'IFPkgFlagFollowLinks': True,
        'IFPkgFlagAuthorizationAction': 'RootAuthorization',
        'IFPkgFlagInstallFat': False,
        'IFPkgFlagInstalledSize': int(installKBytes),
        'IFPkgFlagIsRequired': False,
        'IFPkgFlagRestartAction': 'RequiredRestart',
        'IFPkgFlagRootVolumeOnly': False,
        'IFPkgFormatVersion': 0.10000000149011612
    }
    if build_number:
        info['CFBundleGetInfoString'] = (
            '%s Build %s' % (os_version, build_number))
    output_file = os.path.join(output_pkg_path, 'Contents/Info.plist')
    try:
        plistlib.writePlist(info, output_file)
    except (OSError, ExpatError), err:
        cleanupOutput(output_pkg_path)
        fail('Error creating file at %s: %s' % (output_file, err))


def writefile(stringdata, path):
    '''Writes string data to path.'''
    fileobject = open(path, mode='w', buffering=1)
    print >> fileobject, stringdata
    fileobject.close()


def writePkgInfo(output_pkg_path):
    '''Creates Contents/PkgInfo file'''
    output_file = os.path.join(output_pkg_path, 'Contents/PkgInfo')
    try:
        writefile('pkmkrpkg1', output_file)
    except (OSError, IOError), err:
        cleanupOutput(output_pkg_path)
        fail('Error creating file at %s: %s' % (output_file, err))


def write_package_version(output_pkg_path):
    '''Creates Contents/Resources/package_version file'''
    output_file = os.path.join(output_pkg_path,
                               'Contents/Resources/package_version')
    try:
        writefile('major: 1\nminor: 0', output_file)
    except (OSError, IOError), err:
        cleanupOutput(output_pkg_path)
        fail('Error creating file at %s: %s' % (output_file, err))


def makeArchiveAndBom(output_pkg_path):
    '''Creates an empty Archive.pax.gz and Archive.bom file'''
    emptydir = os.path.join(TMPDIR, 'EmptyDir')
    if os.path.exists(emptydir):
        try:
            os.rmdir(emptydir)
        except OSError, err:
            print >> sys.stderr, ('Existing dir at %s' % emptydir)
            exit(1)
    try:
        os.mkdir(emptydir)
    except OSError, err:
        cleanupOutput(output_pkg_path)
        fail('Can\'t create dir at %s' % emptydir)

    # Make an Archive.pax.gz of the contents of the empty directory
    archiveName = os.path.join(output_pkg_path, 'Contents/Archive.pax')
    # record our current working dir
    cwd = os.getcwd()
    # change into our EmptyDir so we can use pax to archive the
    # (non-existent) contents
    os.chdir(emptydir)
    try:
        subprocess.check_call(
            ['/bin/pax', '-w', '-x', 'cpio', '-f', archiveName, '.'])
    except subprocess.CalledProcessError, err:
        cleanupOutput(output_pkg_path)
        fail('Can\'t create archive at %s: %s' % (archiveName, err))
    # change working dir back
    os.chdir(cwd)
    try:
        subprocess.check_call(['/usr/bin/gzip', archiveName])
    except subprocess.CalledProcessError, err:
        cleanupOutput(output_pkg_path)
        fail('Can\'t gzip archive at %s: %s' % (archiveName, err))
    # now make a BOM file
    bomName = os.path.join(output_pkg_path, 'Contents/Archive.bom')
    try:
        subprocess.check_call(['/usr/bin/mkbom', emptydir, bomName])
    except subprocess.CalledProcessError, err:
        cleanupOutput(output_pkg_path)
        fail('Can\'t make BOM file at %s: %s' % (bomName, err))
    try:
        os.rmdir(emptydir)
    except OSError:
        pass


def getOSversionInfoFromDist(distfile):
    '''Gets osVersion and osBuildVersion if present in
    dist file for OSXInstall.mpkg'''
    try:
        dom = minidom.parse(distfile)
    except ExpatError, err:
        print >> sys.stderr, 'Error parsing %s: %s' % (distfile, err)
        return None, None
    osVersion = None
    osBuildVersion = None
    elements = dom.getElementsByTagName('options')
    if len(elements):
        options = elements[0]
        if 'osVersion' in options.attributes.keys():
            osVersion = options.attributes['osVersion'].value
        if 'osBuildVersion' in options.attributes.keys():
            osBuildVersion = options.attributes['osBuildVersion'].value
    return osVersion, osBuildVersion


def getItemsFromDist(filename):
    '''Gets the title, script, installation-check and volume-check elements
    from an OSXInstall.mpkg distribution file'''
    try:
        dom = minidom.parse(filename)
    except ExpatError:
        print >> sys.stderr, 'Error parsing %s' % filename
        return None
    item_dict = {'title': '',
                 'script': '',
                 'installation_check': '',
                 'volume_check': ''}
    title_elements = dom.getElementsByTagName('title')
    if len(title_elements):
        item_dict['title'] = title_elements[0].firstChild.wholeText
    script_elements = dom.getElementsByTagName('script')
    if len(script_elements):
        for script_element in script_elements:
            if script_element.toprettyxml() != '<script/>':
                item_dict['script'] = script_element.toprettyxml()
        #item_dict['script'] = script_elements[1:].toprettyxml()
    installation_check_elements = dom.getElementsByTagName('installation-check')
    if len(installation_check_elements):
        item_dict['installation_check'] = \
            installation_check_elements[0].toprettyxml()
    volume_check_elements = dom.getElementsByTagName('volume-check')
    if len(volume_check_elements):
        item_dict['volume_check'] = volume_check_elements[0].toprettyxml()
    return item_dict


def make_distribution(output_pkg_path, source_pkg_dist,
                      installKBytes=DEFAULT_INSTALLKBYTES):
    '''Makes a distribution file for the target package based on one
    from OSInstall.mpkg in the app or InstallESD.dmg'''

    item_dict = getItemsFromDist(source_pkg_dist)

    # disable the check for command line installs
    item_dict['script'] = item_dict['script'].replace(
        'system.env.COMMAND_LINE_INSTALL',
        'system.env.COMMAND_LINE_INSTALL_DISABLED')

    dist_header = ('<?xml version="1.0" encoding="utf-8"?>\n'
                   '<installer-script minSpecVersion="1.000000">\n')

    dist_title = '    <title>%s</title>' % item_dict.get('title', 'OS X')
    dist_options = '''
    <options customize="never" allow-external-scripts="yes" rootVolumeOnly="false"/>'
    '''

    dist_choices_outline = '''
    <choices-outline>
        <line choice='manual'/>
    </choices-outline>
    '''

    dist_choice_id = '''
    <choice id='manual'>
        <pkg-ref id='manual' auth='Root'>.</pkg-ref>
    </choice>
    '''

    dist_pkg_ref = ("<pkg-ref id='manual' installKBytes='%s' "
                    "onConclusion='RequireRestart' version='1.0'/>"
                    % installKBytes)

    dist_footer = '\n</installer-script>'

    dist = dist_header + dist_title + dist_options
    dist += item_dict['script']
    dist += '\n    ' + item_dict['installation_check']
    dist += '\n    ' + item_dict['volume_check']
    dist += dist_choices_outline + dist_choice_id + dist_pkg_ref + dist_footer

    output_file = os.path.join(output_pkg_path, 'Contents/distribution.dist')
    try:
        writefile(dist, output_file)
    except (OSError, IOError), err:
        cleanupOutput(output_pkg_path)
        fail('Error creating file at %s: %s' % (output_file, err))


def copyLocalizedResources(pkgpath, source_pkg_resources):
    '''Copies Resources/English.lprog/*.strings to
    Contents/Resources/en.lproj of target package so InstallCheck
    and VolumeCheck scripts can display meaningful error messages'''

    source_dir = os.path.join(source_pkg_resources, 'English.lproj')
    dest_dir = os.path.join(pkgpath, 'Contents/Resources/en.lproj')
    if os.path.isdir(source_dir) and os.path.isdir(dest_dir):
        for item in os.listdir(source_dir):
            if item.endswith('.strings'):
                itempath = os.path.join(source_dir, item)
                try:
                    shutil.copy(itempath, dest_dir)
                except IOError:
                    # not fatal, but warn anyway.
                    print >> sys.stderr, (
                        'Could not copy %s to %s'% (itempath, dest_dir))


def copy_postflight_script(output_pkg_path):
    '''Copies the postflight script into the output package'''
    destination = os.path.join(output_pkg_path, 'Contents/Resources/postflight')
    mydir = os.path.dirname(os.path.abspath(__file__))
    postflight_script_name = 'installosxpkg_postflight'
    locations = [os.path.join(mydir, postflight_script_name),
                 os.path.join(mydir, 'Resources', postflight_script_name)]
    for location in locations:
        if os.path.exists(location):
            try:
                shutil.copy(location, destination)
                # make sure it's executable
                subprocess.check_call(
                    ['/bin/chmod', 'a+x', destination])
                return True
            except (OSError, subprocess.CalledProcessError), err:
                cleanupOutput(output_pkg_path)
                fail('Error with postflight script: %s' % err)
    # if we get here, we couldn't find the postflight script.
    cleanupOutput(output_pkg_path)
    fail('Could not find postflight script.')


def copy_brtool(brtool_path, output_pkg_path):
    '''Copies brtool from the Install OS X.app into the package resources'''
    destination = os.path.join(output_pkg_path, 'Contents/Resources/brtool')
    try:
        shutil.copy(brtool_path, destination)
        # make sure it's executable
        subprocess.check_call(
            ['/bin/chmod', 'a+x', destination])
        return True
    except (OSError, subprocess.CalledProcessError), err:
        cleanupOutput(output_pkg_path)
        fail('Error copying brtool: %s' % err)


def makePackage(output_pkg_path, expanded_osinstall_mpkg,
                os_version, build_number, pkg_id=None):
    '''Makes our output package'''
    makePkgDirs(output_pkg_path)
    makeArchiveAndBom(output_pkg_path)
    makeInfoPlist(output_pkg_path, os_version, build_number,
                  pkg_id=pkg_id)
    writePkgInfo(output_pkg_path)
    makeDescriptionPlist(output_pkg_path, os_version=os_version,
                         build_number=build_number)
    write_package_version(output_pkg_path)
    copy_postflight_script(output_pkg_path)
    # copy some items from OSInstall.mpkg
    dist_file = os.path.join(expanded_osinstall_mpkg, 'Distribution')
    make_distribution(output_pkg_path, dist_file, installKBytes=8*1024*1024)
    source_pkg_resources = os.path.join(expanded_osinstall_mpkg, 'Resources')
    copyLocalizedResources(output_pkg_path, source_pkg_resources)


def makeFakeApplicationBundle(fake_app_path, source_app_path):
    '''Makes a fake app bundle VMware Fusion can use to install OS X'''
    if os.path.exists(fake_app_path):
        fail('Something already exists at %s' % fake_app_path)
    try:
        os.makedirs(os.path.join(fake_app_path, 'Contents/SharedSupport'))
    except OSError, err:
        fail('Error creating fake app directories at %s: %s'
             % (fake_app_path, err))
    source_info_plist = os.path.join(source_app_path, 'Contents/Info.plist')
    destination = os.path.join(fake_app_path, 'Contents/Info.plist')
    try:
        shutil.copy(source_info_plist, destination)
    except OSError, err:
        cleanupOutput(fake_app_path)
        fail('Error copying Info.plist: %s' % err)


def get_dir_size(some_dir):
    '''Gets the total size of some_dir. Very helpful in determining
    the size of bundle packages.'''
    total_size = 0
    for dirpath, dummy_dirnames, filenames in os.walk(some_dir):
        for filename in filenames:
            filepath = os.path.join(dirpath, filename)
            total_size += os.path.getsize(filepath)
    return total_size


def get_size_of_all_packages(pkglist):
    '''Gets the total size of all the extra packages'''
    total_pkg_size = 0
    for item in pkglist:
        if os.path.isdir(item):
            total_pkg_size += get_dir_size(item)
        else:
            total_pkg_size += os.path.getsize(item)
    return total_pkg_size/1024


def get_available_free_space_in_dmg(some_dmg):
    '''Returns free disk space on some_dmg in Kbytes'''
    (mountpoints, dummy_shadowpath) = mountdmg(some_dmg)
    if mountpoints:
        stat = os.statvfs(mountpoints[0])
        free = stat.f_bavail * stat.f_frsize
        unmountdmg(mountpoints[0])
        return int(free/1024)
    else:
        return -1


def expandOSInstallMpkg(osinstall_mpkg):
    '''Expands the flat OSInstall.mpkg We need the Distribution file
    and some .strings files from within. Returns path to the exapnded
    package.'''
    expanded_osinstall_mpkg = os.path.join(TMPDIR, 'OSInstall_mpkg')
    cmd = ['/usr/sbin/pkgutil', '--expand', osinstall_mpkg,
           expanded_osinstall_mpkg]
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError:
        fail('Failed to expand %s' % osinstall_mpkg)
    return expanded_osinstall_mpkg


def downloadURL(URL, to_file=None):
    '''Downloads URL to the current directory or as string'''
    cmd = ['/usr/bin/curl', '--silent', '--show-error', '--url', URL]
    if to_file:
        cmd.extend(['-o', to_file])
    proc = subprocess.Popen(cmd, shell=False, bufsize=-1,
                            stdin=subprocess.PIPE,
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (output, err) = proc.communicate()
    if proc.returncode:
        print >> sys.stderr, 'Error %s retrieving %s' % (proc.returncode, URL)
        print >> sys.stderr, err
        return None
    if to_file:
        return to_file
    else:
        return output


def findIncompatibleAppListPkgURL(catalog_url, package_name):
    '''Searches SU catalog to find a download URL for
    package_name. If there's more than one, returns the
    one with the most recent PostDate.'''

    def sort_by_PostDate(a, b):
        """Internal comparison function for use with sorting"""
        return cmp(b['PostDate'], a['PostDate'])

    catalog_str = downloadURL(catalog_url)
    try:
        catalog = plistlib.readPlistFromString(catalog_str)
    except ExpatError:
        print >> sys.stderr, 'Could not parse catalog!'
        return None
    product_list = []
    if 'Products' in catalog:
        for product_key in catalog['Products'].keys():
            product = catalog['Products'][product_key]
            for package in product.get('Packages', []):
                url = package.get('URL', '')
                if url.endswith(package_name):
                    product_list.append({'PostDate': product['PostDate'],
                                         'URL': url})
        if product_list:
            product_list.sort(sort_by_PostDate)
            return product_list[0]['URL']
    return None


def getPkgAndMakeIndexSproduct(destpath, os_vers='10.7'):
    '''Gets IncompatibleAppList package and creates index.sproduct'''

    LION_PKGNAME = 'MacOS_10_7_IncompatibleAppList.pkg'
    LION_CATALOG_URL = ('http://swscan.apple.com/content/catalogs/others/'
                        'index-lion-snowleopard-leopard.merged-1.sucatalog')

    MTN_LION_PKGNAME = 'OSX_10_8_IncompatibleAppList.pkg'
    MTN_LION_CATALOG_URL = ('https://swscan.apple.com/content/catalogs/others/'
                            'index-mountainlion-lion-snowleopard-leopard'
                            '.merged-1.sucatalog')

    MAVERICKS_PKGNAME = 'OSX_10_9_IncompatibleAppList.pkg'
    MAVERICKS_CATALOG_URL = ('https://swscan.apple.com/content/catalogs/others/'
                             'index-10.9-mountainlion-lion-snowleopard-leopard'
                             '.merged-1.sucatalog')

    YOSEMITE_PKGNAME = 'OSX_10_10_IncompatibleAppList.pkg'
    YOSEMITE_CATALOG_URL = (
        'https://swscan.apple.com/content/catalogs/others/'
        'index-10.10-10.9-mountainlion-lion-snowleopard-leopard'
        '.merged-1.sucatalog')

    EL_CAPITAN_PKGNAME = 'OSX_10_11_IncompatibleAppList.pkg'
    EL_CAPITAN_CATALOG_URL = (
        'https://swscan.apple.com/content/catalogs/others/'
        'index-10.11-10.10-10.9-mountainlion-lion-snowleopard-leopard'
        '.merged-1.sucatalog')

    SIERRA_PKGNAME = 'OSX_10_12_IncompatibleAppList.pkg'
    SIERRA_CATALOG_URL = (
        'https://swscan.apple.com/content/catalogs/others/'
        'index-10.12-10.11-10.10-10.9-mountainlion-lion-snowleopard-leopard'
        '.merged-1.sucatalog')

    if os_vers.startswith('10.7'):
        catalog_url = LION_CATALOG_URL
        package_name = LION_PKGNAME
        os_vers = '10.7'
    elif os_vers.startswith('10.8'):
        catalog_url = MTN_LION_CATALOG_URL
        package_name = MTN_LION_PKGNAME
        os_vers = '10.8'
    elif os_vers.startswith('10.9'):
        catalog_url = MAVERICKS_CATALOG_URL
        package_name = MAVERICKS_PKGNAME
        os_vers = '10.9'
    elif os_vers.startswith('10.10'):
        catalog_url = YOSEMITE_CATALOG_URL
        package_name = YOSEMITE_PKGNAME
        os_vers = '10.10'
    elif os_vers.startswith('10.11'):
        catalog_url = EL_CAPITAN_CATALOG_URL
        package_name = EL_CAPITAN_PKGNAME
        os_vers = '10.11'
    elif os_vers.startswith('10.12'):
        catalog_url = SIERRA_CATALOG_URL
        package_name = SIERRA_PKGNAME
        os_vers = '10.12'
    else:
        print >> sys.stderr, 'Unsupported OS version!'
        return

    destpath = os.path.abspath(destpath)
    if not os.path.isdir(destpath):
        print >> sys.stderr, 'Directory %s doesn\'t exist!' % destpath
        return

    url = findIncompatibleAppListPkgURL(catalog_url, package_name)
    if url:
        package_path = os.path.join(destpath, package_name)
        print 'Downloading %s to %s...' % (url, package_path)
        package_path = downloadURL(url, to_file=package_path)
        if package_path and os.path.exists(package_path):
            # make index.sproduct
            pkg_info = {}
            pkg_info['Identifier'] = 'com.apple.pkg.CompatibilityUpdate'
            pkg_info['Size'] = int(os.path.getsize(package_path))
            pkg_info['URL'] = package_name
            #pkg_info['Version'] = os_vers
            # nope. Version is 10.7 even for ML (!)
            pkg_info['Version'] = '10.7'
            index_dict = {}
            index_dict['Packages'] = [pkg_info]
            plist_path = os.path.join(destpath, 'index.sproduct')
            print "Writing index.sproduct to %s..." % plist_path
            try:
                plistlib.writePlist(index_dict, plist_path)
            except OSError, err:
                print >> sys.stderr, 'Write error: %s' % err
        else:
            print >> sys.stderr, 'Couldn\'t download %s' % url
    else:
        print >> sys.stderr, 'Couldn\'t find IncompatibleAppList package.'


def makeEmptyInstallerChoiceChanges(output_pkg_path):
    '''Creates an empty MacOSXInstaller.choiceChanges file'''
    destpath = os.path.join(output_pkg_path,
                            'Contents/Resources/OS X Install Data',
                            'MacOSXInstaller.choiceChanges')
    changes = []
    try:
        plistlib.writePlist(changes, destpath)
    except OSError, err:
        print >> sys.stderr, 'Error writing %s: %s' % (destpath, err)


class AddPackageError(Exception):
    '''Errors generated by addPackagesToInstallESD'''
    pass


def addPackagesToInstallESD(installesd_dmg, packages,
                            output_dmg_path, create_minstallconfig=False):
    '''Adds additional packages to the InstallESD.dmg and creates an
    OSInstall.collection file for use by the installer. New dmg is
    created at output_dmg_path'''

    # generate OSInstall.collection pkg_array
    # array needs OSInstall.mpkg twice at the beginning
    # no idea why
    pkg_array = ['/System/Installation/Packages/OSInstall.mpkg',
                 '/System/Installation/Packages/OSInstall.mpkg']
    for pkg in packages:
        pkgname = os.path.basename(pkg)
        pkg_path = os.path.join('/System/Installation/Packages', pkgname)
        pkg_array.append(pkg_path)

    # mount InstallESD.dmg with shadow
    print 'Mounting %s...' % installesd_dmg
    mountpoints, shadowpath = mountdmg(installesd_dmg, use_shadow=True)
    if not mountpoints:
        raise AddPackageError('Nothing mounted from InstallESD.dmg')

    # copy additional packages to Packages directory
    mountpoint = mountpoints[0]
    packages_dir = os.path.join(mountpoint, 'Packages')
    print 'Copying additional packages to InstallESD/Packages/:'
    try:
        for pkg in packages:
            if os.path.isdir(pkg):
                destination = os.path.join(packages_dir, os.path.basename(pkg))
                print '    Copying bundle package %s' % pkg
                shutil.copytree(pkg, destination)
            else:
                print '    Copying flat package %s' % pkg
                shutil.copy(pkg, packages_dir)
    except IOError, err:
        unmountdmg(mountpoint)
        raise AddPackageError('Error %s copying packages to disk image' % err)

    # create OSInstall.collection in Packages directory
    osinstall_collection_path = os.path.join(
        packages_dir, 'OSInstall.collection')
    print "Creating %s" % osinstall_collection_path
    try:
        plistlib.writePlist(pkg_array, osinstall_collection_path)
    except ExpatError:
        unmountdmg(mountpoint)
        raise AddPackageError('Error %s creating OSInstall.collection' % err)

    if create_minstallconfig:
        minstallconfig = {
            'InstallType': 'automated',
            'Language': 'en',
            'Package': '/System/Installation/Packages/OSInstall.collection',
            'Target': '/Volumes/Macintosh HD',
            'TargetName': 'Macintosh HD'
        }
        extras_dir = os.path.join(packages_dir, 'Extras')
        if not os.path.exists(extras_dir):
            try:
                os.mkdir(extras_dir)
            except (OSError, IOError), err:
                unmountdmg(mountpoint)
                raise AddPackageError(
                    'Error %s creating System/Installation/Packages/Extras on '
                    'disk image' % err)
        minstallconfig_path = os.path.join(extras_dir, "minstallconfig.xml")
        print "Creating %s" % minstallconfig_path
        try:
            plistlib.writePlist(minstallconfig, minstallconfig_path)
        except ExpatError:
            unmountdmg(mountpoint)
            raise AddPackageError('Error %s creating minstallconfig.xml' % err)

    # unmount InstallESD.dmg
    print 'Unmounting %s...' % installesd_dmg
    unmountdmg(mountpoint)

    # convert InstallESD.dmg + shadow to UDZO image
    print 'Creating disk image at %s...' % output_dmg_path
    cmd = ['/usr/bin/hdiutil', 'convert', '-format', 'UDZO',
           '-o', output_dmg_path, installesd_dmg, '-shadow', shadowpath]
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError, err:
        raise AddPackageError(
            'Failed to create %s at: %s' % (output_dmg_path, err))


def is_flat_distribution(pkg):
    '''Returns True if pkg appears to be a flat distribution-style package,
    False otherwise'''
    if os.path.isdir(pkg):
        # it's a bundle-style pkg
        return False
    cmd = ['/usr/bin/xar', '-tf', pkg]
    proc = subprocess.Popen(cmd, shell=False, bufsize=-1,
                            stdin=subprocess.PIPE,
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    output = proc.communicate()[0]
    if proc.returncode:
        # some xar error; invalid pkg?
        print >> sys.stderr, 'xar error getting TOC for %s' % pkg
        return False
    output_lines = output.splitlines()
    if 'Distribution' in output_lines or 'distribution' in output_lines:
        return True
    return False


TMPDIR = None
def main():
    '''Builds a custom package that installs OS X. You may specify additional
    packages to install after the OS is installed'''
    global TMPDIR

    usage = (
        'Usage: %prog --source InstallOSX.app\n'
        '                           [--pkg path/to/additional.pkg]\n'
        '                           [--output path/to/InstallOSX.pkg]\n'
        '                           [--identifier com.example.installosx.pkg]\n'
        '                           [--plist path/to/config.plist]\n\n'
        '    %prog creates a customized OS X\n'
        '    installation package containing the contents of the original\n'
        '    InstallESD.dmg plus any additional packages provided. Additional\n'
        '    packages will be installed in the order you provide them at the\n'
        '    command-line.')

    parser = optparse.OptionParser(usage=usage)
    parser.add_option(
        '--source', '-s',
        help='Required unless specified via plist. Path to Install Mac '
        'OS X Lion.app or Install OS X Foo.app ')
    parser.add_option(
        '--pkg', '-p', action="append", dest='packages',
        metavar='PACKAGE',
        help='Optional. An addtional package to include for installation. '
        'May be specified more than once. Not supported past macOS 10.12.3.')
    parser.add_option('--output', '-o', help='Optional. Path for output pkg. '
                      'Defaults to current working directory.')
    parser.add_option(
        '--identifier', '--id',
        help='Optional. Package identifier for the package. Defaults to '
        '"com.googlecode.munki.installosx.pkg"')
    parser.add_option(
        '--plist', help='Optional. Path to an XML plist file '
        'containing key/value pairs for Source, Output, Packages, and '
        'Identifier.')
    parser.add_option(
        '--make-fake-app', action='store_true', help='Optional. '
        'Instead of creating an installer pkg, creates a fake installer '
        'application for use with VMware Fusion to install a customized OS X. '
        'Experimental.')
    options, arguments = parser.parse_args()

    if arguments:
        print >> sys.stderr, 'Arguments found with no option flags!'
        parser.print_help()
        exit(-1)

    # check to see if we're root
    # need to be root to copy things into the DMG with the right
    # ownership and permissions
    if os.geteuid() != 0:
        print >> sys.stderr, 'You must run this as root, or via sudo!'
        exit(-1)

    plist_options = {}
    if options.plist:
        try:
            plist_options = plistlib.readPlist(options.plist)
        except (ExpatError, IOError), err:
            fail('Could not read %s: %s' % (options.plist, err))

    if not options.source and 'Source' not in plist_options:
        print >> sys.stderr, ('ERROR: Must have --source option!')
        parser.print_usage()
        exit(1)

    TMPDIR = tempfile.mkdtemp(dir='/tmp')
    source = options.source or plist_options.get('Source')
    source = source.rstrip('/')
    if source.endswith('.app'):
        if not os.path.isdir(source):
            fail('%s doesn\'t exist or isn\'t an app!' % source)
        installesd_dmg = os.path.join(
            source, 'Contents/SharedSupport/InstallESD.dmg')
        if not os.path.exists(installesd_dmg):
            fail('Can\'t find InstallESD.dmg at %s' % installesd_dmg)

    elif source.endswith('.dmg'):
        fail('This tool no longer supports using an InstallESD.dmg as a source.'
             '/nProvide a path to an OS X/macOS install application instead.')

    else:
        fail('Unknown/unsupported source: %s' % source)

    additional_packages = options.packages or plist_options.get('Packages', [])
    if options.make_fake_app and not additional_packages:
        fail('No additional packages specified -- nothing to add to the fake '
             'install app!')

    # get some needed info from the disk image
    print 'Examining and verifying source...'
    mountpoints, dummy_shadowpath = mountdmg(installesd_dmg)
    if not mountpoints:
        fail('Could not mount diskimage %s' % installesd_dmg)
    mountpoint = mountpoints[0]

    # get info from BaseSystem.dmg
    basesystem_dmg = os.path.join(mountpoint, 'BaseSystem.dmg')
    if not os.path.isfile(basesystem_dmg):
        unmountdmg(mountpoint)
        fail('Missing BaseSystem.dmg in %s'% source)

    basesystemmountpoints, dummy_shadowpath = mountdmg(basesystem_dmg)
    basesystemmountpoint = basesystemmountpoints[0]
    system_version_plist = os.path.join(
        basesystemmountpoint,
        'System/Library/CoreServices/SystemVersion.plist')
    try:
        version_info = plistlib.readPlist(system_version_plist)
    except (ExpatError, IOError), err:
        unmountdmg(basesystemmountpoint)
        unmountdmg(mountpoint)
        fail('Could not read %s: %s' % (system_version_plist, err))
    else:
        unmountdmg(basesystemmountpoint)

    # get info from Packages/OSInstall.mpkg (we use the Distribution file)
    osinstall_mpkg = os.path.join(mountpoint, 'Packages/OSInstall.mpkg')
    if not os.path.exists(osinstall_mpkg):
        unmountdmg(mountpoint)
        fail('Missing OSInstall.mpkg in %s'% source)

    expanded_osinstall_mpkg = expandOSInstallMpkg(osinstall_mpkg)
    distfile = os.path.join(expanded_osinstall_mpkg, 'Distribution')

    unmountdmg(mountpoint)
    os_version = version_info.get('ProductUserVisibleVersion')
    build_number = version_info.get('ProductBuildVersion')
    if os_version is None or build_number is None:
        fail('Missing OS version or build info in %s' % system_version_plist)

    # Things we have now that we need:
    # installesd_dmg: path to the InstallESD.dmg file
    # os_version: string like '10.7.4'
    # build_number: string like '11E53'
    # expanded_osinstall_mpkg: path to unflattened OSInstall.mpkg
    # distfile: path to Distribution file in OSInstall.mpkg

    print '----------------------------------------------------------------'
    print 'InstallESD.dmg: %s' % installesd_dmg
    print 'OS Version: %s' % os_version
    print 'OS Build: %s' % build_number

    if DEBUG:
        print 'expanded_osinstall_mpkg: %s' % expanded_osinstall_mpkg
        print 'distfile: %s' % distfile
    print '----------------------------------------------------------------'

    # now that we know the target OS versions, we need to do a
    # version-specific check
    if additional_packages:
        if version.LooseVersion(os_version) > version.LooseVersion('10.12.3'):
            fail('ERROR: Additional packages are not supported for macOS '
                 'installers higher than 10.12.3 due to package verification '
                 'changes in macOS installers post-10.12.3.')

    # Figure out where we will be writing this...
    custom_tag = ''

    if additional_packages:
        custom_tag = '_custom'
        # get rid of trailing slashes which often result from dragging
        # and dropping from the Finder into the Terminal
        additional_packages = [item.rstrip('/') for item in additional_packages]
    if options.make_fake_app:
        outputname_base = 'FakeInstallOSX_%s_%s%s.app'
    else:
        outputname_base = 'InstallOSX_%s_%s%s.pkg'
    outputname = outputname_base % (os_version, build_number, custom_tag)
    output_path = os.path.abspath(os.path.join('.', outputname))

    output = options.output or plist_options.get('Output')
    if output:
        if ((options.make_fake_app and output.endswith('.app')) or
                output.endswith('.pkg')):
            # we've been given a full path including the app/dmg name
            output_path = os.path.abspath(output)
        else:
            # it better be a pre-existing directory
            if not os.path.isdir(output):
                fail('Directory %s not found!' % output)
            else:
                output_path = os.path.abspath(
                    os.path.join(output, outputname))

    if os.path.exists(output_path):
        fail('%s already exists!' % output_path)

    # now we have an output path
    print 'Output path: %s' % output_path

    if additional_packages:
        # make sure they all exist and look like packages
        print 'Additional packages:'
        print '----------------------------------------------------------------'
        for pkg in additional_packages:
            if not pkg.endswith('.pkg') and not pkg.endswith('.mpkg'):
                fail('%s doesn\'t look like a package!' % pkg)
            if not os.path.exists(pkg):
                fail('Package %s not found!' % pkg)
            minor_version = int(os_version.split('.')[1])
            if minor_version >= 10 and not is_flat_distribution(pkg):
                fail('%s is not a flat distribution package. '
                     'This will cause a 10.10+ '
                     'install to fail.' % pkg)
            print os.path.basename(pkg)
        print '----------------------------------------------------------------'
        total_package_size = get_size_of_all_packages(additional_packages)
        print 'Total additional package size: %s Kbytes' % total_package_size
        print '----------------------------------------------------------------'
        # make sure we have enough space on DMG
        print 'Checking available space on %s...' % installesd_dmg
        available_disk_space = get_available_free_space_in_dmg(installesd_dmg)
        if available_disk_space == -1:
            fail('Could not mount disk image %s' % installesd_dmg)

        if total_package_size + 100 > available_disk_space:
            print >> sys.stderr, (
                'Not enough space to copy all packages to InstallESD.dmg!')
            print >> sys.stderr, (
                'Available space: %s KBytes' % available_disk_space)
            fail()

    pkg_id = (options.identifier or plist_options.get('Identifier'))
    # we have everything we need.
    need_to_print_brtool_warning = False
    if not options.make_fake_app:
        output_pkg_path = output_path
        print 'Creating package wrapper...'
        makePackage(output_pkg_path, expanded_osinstall_mpkg,
                    os_version, build_number, pkg_id=pkg_id)

        # if we have any failures after this point, we should be sure to
        # clean up the broken package as well as our TMPDIR

        if version.LooseVersion(os_version) > version.LooseVersion('10.12.3'):
            need_to_print_brtool_warning = True
        else:
            brtool_path = os.path.join(source, 'Contents/Resources/brtool')
            if not os.path.exists(brtool_path):
                fail('Can\'t find brtool at %s' % brtool_path)
            print 'Copying brtool into package resources...'
            copy_brtool(brtool_path, output_pkg_path)

        print 'Creating MacOSXInstaller.choiceChanges...'
        makeEmptyInstallerChoiceChanges(output_pkg_path)

        print '----------------------------------------------------------------'
        print 'Downloading and adding IncompatibleAppList pkg...'
        destpath = os.path.join(output_pkg_path,
                                'Contents/Resources/OS X Install Data')
        getPkgAndMakeIndexSproduct(destpath, os_vers=os_version)

        print '----------------------------------------------------------------'
        print 'Copying InstallESD into package...'
        output_dmg_path = os.path.join(
            output_path, 'Contents/Resources/InstallESD.dmg')
    else:
        print '----------------------------------------------------------------'
        print 'Creating custom fake Install.app...'
        makeFakeApplicationBundle(output_path, source)
        output_dmg_path = os.path.join(
            output_path, 'Contents/SharedSupport/InstallESD.dmg')

    if additional_packages or options.make_fake_app:
        create_minstallconfig = False
        if options.make_fake_app:
            create_minstallconfig = True
        try:
            addPackagesToInstallESD(
                installesd_dmg, additional_packages, output_dmg_path,
                create_minstallconfig=create_minstallconfig)
        except AddPackageError, err:
            cleanupOutput(output_path)
            fail(err)
    else:
        try:
            shutil.copy(installesd_dmg, output_dmg_path)
        except OSError, err:
            cleanupOutput(output_pkg_path)
            fail('Copy error: %e' % err)

    print '----------------------------------------------------------------'
    if not options.make_fake_app:
        print 'Done! Completed package at: %s' % output_path
        if need_to_print_brtool_warning:
            print ('***********************************************************'
                   '\nWARNING: the (nearly) completed package is missing '
                   'Contents/Resources/brtool.\nThe brtool included with '
                   '10.12.4 and later has changed functionality.\nYou\'ll need '
                   'to copy an older brtool from a 10.12-10.12.3 Install macOS '
                   'Sierra.app or from a createOSXinstallPkg-created '
                   '10.12-10.12.3 package into:\n%s\n'
                   '***********************************************************'
                   % os.path.join(output_path, 'Contents/Resources/'))
    else:
        print 'Done! Completed fake app image at: %s' % output_path

    cleanUp()


if __name__ == '__main__':
    main()