# Copyright 2016 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import print_function

import argparse
import codecs
import datetime
import fnmatch
import glob
import json
import os
import plistlib
import shutil
import subprocess
import sys
import tempfile

if sys.version_info.major < 3:
    basestring_compat = basestring
else:
    basestring_compat = str


def GetProvisioningProfilesDir():
    """Returns the location of the installed mobile provisioning profiles.

    Returns:
        The path to the directory containing the installed mobile provisioning
        profiles as a string.
    """
    return os.path.join(os.environ['HOME'], 'Library', 'MobileDevice',
                        'Provisioning Profiles')


def ReadPlistFromString(plist_bytes):
    """Parse property list from given |plist_bytes|.

    Args:
        plist_bytes: contents of property list to load. Must be bytes in python
        3.

    Returns:
        The contents of property list as a python object.
    """
    if sys.version_info.major == 2:
        return plistlib.readPlistFromString(plist_bytes)
    else:
        return plistlib.loads(plist_bytes)


def LoadPlistFile(plist_path):
    """Loads property list file at |plist_path|.

    Args:
        plist_path: path to the property list file to load.

    Returns:
        The content of the property list file as a python object.
    """
    if sys.version_info.major == 2:
        return plistlib.readPlistFromString(
            subprocess.check_output(
                ['xcrun', 'plutil', '-convert', 'xml1', '-o', '-', plist_path]))
    else:
        with open(plist_path, 'rb') as fp:
            return plistlib.load(fp)


def CreateSymlink(value, location):
    """Creates symlink with value at location if the target exists."""
    target = os.path.join(os.path.dirname(location), value)
    if os.path.exists(location):
        os.unlink(location)
    os.symlink(value, location)


class Bundle(object):
    """Wraps a bundle."""

    def __init__(self, bundle_path, platform):
        """Initializes the Bundle object with data from bundle Info.plist
        file."""
        self._path = bundle_path
        self._kind = Bundle.Kind(platform, os.path.splitext(bundle_path)[-1])
        self._data = None

    def Load(self):
        self._data = LoadPlistFile(self.info_plist_path)

    @staticmethod
    def Kind(platform, extension):
        if platform == 'appletvsimulator' or platform == 'appletvos':
            return 'tvos'
        if platform == 'iphonesimulator' or platform == 'iphoneos':
            return 'ios'
        if platform == 'macosx':
            if extension == '.framework':
                return 'mac_framework'
            return 'mac'
        raise ValueError('unknown bundle type %s for %s' %
                         (extension, platform))

    @property
    def kind(self):
        return self._kind

    @property
    def path(self):
        return self._path

    @property
    def contents_dir(self):
        if self._kind == 'mac':
            return os.path.join(self.path, 'Contents')
        if self._kind == 'mac_framework':
            return os.path.join(self.path, 'Versions/A')
        return self.path

    @property
    def executable_dir(self):
        if self._kind == 'mac':
            return os.path.join(self.contents_dir, 'MacOS')
        return self.contents_dir

    @property
    def resources_dir(self):
        if self._kind == 'mac' or self._kind == 'mac_framework':
            return os.path.join(self.contents_dir, 'Resources')
        return self.path

    @property
    def info_plist_path(self):
        if self._kind == 'mac_framework':
            return os.path.join(self.resources_dir, 'Info.plist')
        return os.path.join(self.contents_dir, 'Info.plist')

    @property
    def signature_dir(self):
        return os.path.join(self.contents_dir, '_CodeSignature')

    @property
    def identifier(self):
        return self._data['CFBundleIdentifier']

    @property
    def binary_name(self):
        return self._data['CFBundleExecutable']

    @property
    def binary_path(self):
        return os.path.join(self.executable_dir, self.binary_name)

    def Validate(self, expected_mappings):
        """Checks that keys in the bundle have the expected value.

        Args:
            expected_mappings: a dictionary of string to object, each mapping
            will be looked up in the bundle data to check it has the same value
            (missing values will be ignored)

        Returns:
            A dictionary of the key with a different value between
            expected_mappings and the content of the bundle (i.e. errors) so
            that caller can format the error message. The dictionary will be
            empty if there are no errors.
        """
        errors = {}
        for key, expected_value in expected_mappings.items():
            if key in self._data:
                value = self._data[key]
                if value != expected_value:
                    errors[key] = (value, expected_value)
        return errors


class ProvisioningProfile(object):
    """Wraps a mobile provisioning profile file."""

    def __init__(self, provisioning_profile_path):
        """Initializes the ProvisioningProfile with data from profile file."""
        self._path = provisioning_profile_path
        self._data = ReadPlistFromString(
            subprocess.check_output([
                'xcrun', 'security', 'cms', '-D', '-u', 'certUsageAnyCA', '-i',
                provisioning_profile_path
            ]))

    @property
    def path(self):
        return self._path

    @property
    def team_identifier(self):
        return self._data.get('TeamIdentifier', [''])[0]

    @property
    def name(self):
        return self._data.get('Name', '')

    @property
    def application_identifier_pattern(self):
        return self._data.get('Entitlements', {}).get('application-identifier',
                                                      '')

    @property
    def application_identifier_prefix(self):
        return self._data.get('ApplicationIdentifierPrefix', [''])[0]

    @property
    def entitlements(self):
        return self._data.get('Entitlements', {})

    @property
    def expiration_date(self):
        return self._data.get('ExpirationDate', datetime.datetime.now())

    def ValidToSignBundle(self, bundle_identifier):
        """Checks whether the provisioning profile can sign bundle_identifier.

        Args:
            bundle_identifier: the identifier of the bundle that needs to be
            signed.

        Returns:
            True if the mobile provisioning profile can be used to sign a bundle
            with the corresponding bundle_identifier, False otherwise.
        """
        return fnmatch.fnmatch(
            '%s.%s' % (self.application_identifier_prefix, bundle_identifier),
            self.application_identifier_pattern)

    def Install(self, installation_path):
        """Copies mobile provisioning profile info to |installation_path|."""
        shutil.copy2(self.path, installation_path)


class Entitlements(object):
    """Wraps an Entitlement plist file."""

    def __init__(self, entitlements_path):
        """Initializes Entitlements object from entitlement file."""
        self._path = entitlements_path
        self._data = LoadPlistFile(self._path)

    @property
    def path(self):
        return self._path

    def ExpandVariables(self, substitutions):
        self._data = self._ExpandVariables(self._data, substitutions)

    def _ExpandVariables(self, data, substitutions):
        if isinstance(data, basestring_compat):
            for key, substitution in substitutions.items():
                data = data.replace('$(%s)' % (key,), substitution)
            return data

        if isinstance(data, dict):
            for key, value in data.items():
                data[key] = self._ExpandVariables(value, substitutions)
            return data

        if isinstance(data, list):
            for i, value in enumerate(data):
                data[i] = self._ExpandVariables(value, substitutions)

        return data

    def LoadDefaults(self, defaults):
        for key, value in defaults.items():
            if key not in self._data:
                self._data[key] = value

    def WriteTo(self, target_path):
        with open(target_path, 'wb') as fp:
            if sys.version_info.major == 2:
                plistlib.writePlist(self._data, fp)
            else:
                plistlib.dump(self._data, fp)


def FindProvisioningProfile(bundle_identifier, required):
    """Finds mobile provisioning profile to use to sign bundle.

    Args:
        bundle_identifier: the identifier of the bundle to sign.

    Returns:
        The ProvisioningProfile object that can be used to sign the Bundle
        object or None if no matching provisioning profile was found.
    """
    provisioning_profile_paths = glob.glob(
        os.path.join(GetProvisioningProfilesDir(), '*.mobileprovision'))

    # Iterate over all installed mobile provisioning profiles and filter those
    # that can be used to sign the bundle, ignoring expired ones.
    now = datetime.datetime.now()
    valid_provisioning_profiles = []
    one_hour = datetime.timedelta(0, 3600)
    for provisioning_profile_path in provisioning_profile_paths:
        provisioning_profile = ProvisioningProfile(provisioning_profile_path)
        if provisioning_profile.expiration_date - now < one_hour:
            sys.stderr.write(
                'Warning: ignoring expired provisioning profile: %s.\n' %
                provisioning_profile_path)
            continue
        if provisioning_profile.ValidToSignBundle(bundle_identifier):
            valid_provisioning_profiles.append(provisioning_profile)

    if not valid_provisioning_profiles:
        if required:
            sys.stderr.write(
                'Error: no mobile provisioning profile found for "%s".\n' %
                bundle_identifier)
            sys.exit(1)
        return None

    # Select the most specific mobile provisioning profile, i.e. the one with
    # the longest application identifier pattern (prefer the one with the latest
    # expiration date as a secondary criteria).
    selected_provisioning_profile = max(
        valid_provisioning_profiles,
        key=lambda p:
        (len(p.application_identifier_pattern), p.expiration_date))

    one_week = datetime.timedelta(7)
    if selected_provisioning_profile.expiration_date - now < 2 * one_week:
        sys.stderr.write(
            'Warning: selected provisioning profile will expire soon: %s' %
            selected_provisioning_profile.path)
    return selected_provisioning_profile


def CodeSignBundle(bundle_path, identity, extra_args):
    process = subprocess.Popen([
        'xcrun', 'codesign', '--force', '--sign', identity, '--timestamp=none'
    ] + list(extra_args) + [bundle_path],
                               stderr=subprocess.PIPE,
                               universal_newlines=True)
    _, stderr = process.communicate()
    if process.returncode:
        sys.stderr.write(stderr)
        sys.exit(process.returncode)
    for line in stderr.splitlines():
        if line.endswith(': replacing existing signature'):
            # Ignore warning about replacing existing signature as this should only
            # happen when re-signing system frameworks (and then it is expected).
            continue
        sys.stderr.write(line)
        sys.stderr.write('\n')


def InstallSystemFramework(framework_path, bundle_path, args):
    """Install framework from |framework_path| to |bundle| and code-re-sign
    it."""
    installed_framework_path = os.path.join(bundle_path, 'Frameworks',
                                            os.path.basename(framework_path))

    if os.path.isfile(framework_path):
        shutil.copy(framework_path, installed_framework_path)
    elif os.path.isdir(framework_path):
        if os.path.exists(installed_framework_path):
            shutil.rmtree(installed_framework_path)
        shutil.copytree(framework_path, installed_framework_path)

    CodeSignBundle(
        installed_framework_path, args.identity,
        ['--deep', '--preserve-metadata=identifier,entitlements,flags'])


def GenerateEntitlements(path, provisioning_profile, bundle_identifier):
    """Generates an entitlements file.

    Args:
        path: path to the entitlements template file
        provisioning_profile: ProvisioningProfile object to use, may be None
        bundle_identifier: identifier of the bundle to sign.
    """
    entitlements = Entitlements(path)
    if provisioning_profile:
        entitlements.LoadDefaults(provisioning_profile.entitlements)
        app_identifier_prefix = \
          provisioning_profile.application_identifier_prefix + '.'
    else:
        app_identifier_prefix = '*.'
    entitlements.ExpandVariables({
        'CFBundleIdentifier': bundle_identifier,
        'AppIdentifierPrefix': app_identifier_prefix,
    })
    return entitlements


def GenerateBundleInfoPlist(bundle, plist_compiler, partial_plist):
    """Generates the bundle Info.plist for a list of partial .plist files.

    Args:
        bundle: a Bundle instance
        plist_compiler: string, path to the Info.plist compiler
        partial_plist: list of path to partial .plist files to merge
    """

    # Filter empty partial .plist files (this happens if an application
    # does not compile any asset catalog, in which case the partial .plist
    # file from the asset catalog compilation step is just a stamp file).
    filtered_partial_plist = []
    for plist in partial_plist:
        plist_size = os.stat(plist).st_size
        if plist_size:
            filtered_partial_plist.append(plist)

    # Invoke the plist_compiler script. It needs to be a python script.
    subprocess.check_call([
        'python3',
        plist_compiler,
        'merge',
        '-f',
        'binary1',
        '-o',
        bundle.info_plist_path,
    ] + filtered_partial_plist)


class Action(object):
    """Class implementing one action supported by the script."""

    @classmethod
    def Register(cls, subparsers):
        parser = subparsers.add_parser(cls.name, help=cls.help)
        parser.set_defaults(func=cls._Execute)
        cls._Register(parser)


class CodeSignBundleAction(Action):
    """Class implementing the code-sign-bundle action."""

    name = 'code-sign-bundle'
    help = 'perform code signature for a bundle'

    @staticmethod
    def _Register(parser):
        parser.add_argument('--entitlements',
                            '-e',
                            dest='entitlements_path',
                            help='path to the entitlements file to use')
        parser.add_argument('path', help='path to the iOS bundle to codesign')
        parser.add_argument('--identity',
                            '-i',
                            required=True,
                            help='identity to use to codesign')
        parser.add_argument('--binary',
                            '-b',
                            required=True,
                            help='path to the iOS bundle binary')
        parser.add_argument('--framework',
                            '-F',
                            action='append',
                            default=[],
                            dest='frameworks',
                            help='install and resign system framework')
        parser.add_argument('--disable-code-signature',
                            action='store_true',
                            dest='no_signature',
                            help='disable code signature')
        parser.add_argument(
            '--disable-embedded-mobileprovision',
            action='store_false',
            default=True,
            dest='embedded_mobileprovision',
            help='disable finding and embedding mobileprovision')
        parser.add_argument('--platform',
                            '-t',
                            required=True,
                            help='platform the signed bundle is targeting')
        parser.add_argument(
            '--partial-info-plist',
            '-p',
            action='append',
            default=[],
            help=
            'path to partial Info.plist to merge to create bundle Info.plist')
        parser.add_argument(
            '--plist-compiler-path',
            '-P',
            action='store',
            help='path to the plist compiler script (for --partial-info-plist)')
        parser.set_defaults(no_signature=False)

    @staticmethod
    def _Execute(args):
        if not args.identity:
            args.identity = '-'

        bundle = Bundle(args.path, args.platform)

        if args.partial_info_plist:
            GenerateBundleInfoPlist(bundle, args.plist_compiler_path,
                                    args.partial_info_plist)

        # The bundle Info.plist may have been updated by
        # GenerateBundleInfoPlist() above. Load the bundle information from
        # Info.plist after the modification have been written to disk.
        bundle.Load()

        # According to Apple documentation, the application binary must be the
        # same as the bundle name without the .app suffix. See crbug.com/740476
        # for more information on what problem this can cause.
        #
        # To prevent this class of error, fail with an error if the binary name
        # is incorrect in the Info.plist as it is not possible to update the
        # value in Info.plist at this point (the file has been copied by a
        # different target and ninja would consider the build dirty if it was
        # updated).
        #
        # Also checks that the name of the bundle is correct too (does not cause
        # the build to be considered dirty, but still terminate the script in
        # case of an incorrect bundle name).
        #
        # Apple documentation is available at:
        # https://developer.apple.com/library/content/documentation/CoreFoundation/Conceptual/CFBundles/BundleTypes/BundleTypes.html
        bundle_name = os.path.splitext(os.path.basename(bundle.path))[0]
        errors = bundle.Validate({
            'CFBundleName': bundle_name,
            'CFBundleExecutable': bundle_name,
        })
        if errors:
            for key in sorted(errors):
                value, expected_value = errors[key]
                sys.stderr.write('%s: error: %s value incorrect: %s != %s\n' %
                                 (bundle.path, key, value, expected_value))
            sys.stderr.flush()
            sys.exit(1)

        # Delete existing embedded mobile provisioning.
        embedded_provisioning_profile = os.path.join(
            bundle.path, 'embedded.mobileprovision')
        if os.path.isfile(embedded_provisioning_profile):
            os.unlink(embedded_provisioning_profile)

        # Delete existing code signature.
        if os.path.exists(bundle.signature_dir):
            shutil.rmtree(bundle.signature_dir)

        # Install system frameworks if requested.
        for framework_path in args.frameworks:
            InstallSystemFramework(framework_path, args.path, args)

        # Copy main binary into bundle.
        if not os.path.isdir(bundle.executable_dir):
            os.makedirs(bundle.executable_dir)
        shutil.copy(args.binary, bundle.binary_path)

        if bundle.kind == 'mac_framework':
            # Create Versions/Current -> Versions/A symlink
            CreateSymlink('A', os.path.join(bundle.path, 'Versions/Current'))

            # Create $binary_name -> Versions/Current/$binary_name symlink
            CreateSymlink(os.path.join('Versions/Current', bundle.binary_name),
                          os.path.join(bundle.path, bundle.binary_name))

            # Create optional symlinks.
            for name in ('Headers', 'Resources', 'Modules'):
                target = os.path.join(bundle.path, 'Versions/A', name)
                if os.path.exists(target):
                    CreateSymlink(os.path.join('Versions/Current', name),
                                  os.path.join(bundle.path, name))
                else:
                    obsolete_path = os.path.join(bundle.path, name)
                    if os.path.exists(obsolete_path):
                        os.unlink(obsolete_path)

        if args.no_signature:
            return

        codesign_extra_args = []

        if args.embedded_mobileprovision:
            # Find mobile provisioning profile and embeds it into the bundle (if
            # a code signing identify has been provided, fails if no valid
            # mobile provisioning is found).
            provisioning_profile_required = args.identity != '-'
            provisioning_profile = FindProvisioningProfile(
                bundle.identifier, provisioning_profile_required)
            if provisioning_profile and args.platform != 'iphonesimulator':
                provisioning_profile.Install(embedded_provisioning_profile)

                if args.entitlements_path is not None:
                    temporary_entitlements_file = \
                        tempfile.NamedTemporaryFile(suffix='.xcent')
                    codesign_extra_args.extend(
                        ['--entitlements', temporary_entitlements_file.name])

                    entitlements = GenerateEntitlements(args.entitlements_path,
                                                        provisioning_profile,
                                                        bundle.identifier)
                    entitlements.WriteTo(temporary_entitlements_file.name)

        CodeSignBundle(bundle.path, args.identity, codesign_extra_args)


class CodeSignFileAction(Action):
    """Class implementing code signature for a single file."""

    name = 'code-sign-file'
    help = 'code-sign a single file'

    @staticmethod
    def _Register(parser):
        parser.add_argument('path', help='path to the file to codesign')
        parser.add_argument('--identity',
                            '-i',
                            required=True,
                            help='identity to use to codesign')
        parser.add_argument(
            '--output',
            '-o',
            help='if specified copy the file to that location before signing it'
        )
        parser.set_defaults(sign=True)

    @staticmethod
    def _Execute(args):
        if not args.identity:
            args.identity = '-'

        install_path = args.path
        if args.output:

            if os.path.isfile(args.output):
                os.unlink(args.output)
            elif os.path.isdir(args.output):
                shutil.rmtree(args.output)

            if os.path.isfile(args.path):
                shutil.copy(args.path, args.output)
            elif os.path.isdir(args.path):
                shutil.copytree(args.path, args.output)

            install_path = args.output

        CodeSignBundle(
            install_path, args.identity,
            ['--deep', '--preserve-metadata=identifier,entitlements'])


class GenerateEntitlementsAction(Action):
    """Class implementing the generate-entitlements action."""

    name = 'generate-entitlements'
    help = 'generate entitlements file'

    @staticmethod
    def _Register(parser):
        parser.add_argument('--entitlements',
                            '-e',
                            dest='entitlements_path',
                            help='path to the entitlements file to use')
        parser.add_argument('path',
                            help='path to the entitlements file to generate')
        parser.add_argument('--info-plist',
                            '-p',
                            required=True,
                            help='path to the bundle Info.plist')

    @staticmethod
    def _Execute(args):
        info_plist = LoadPlistFile(args.info_plist)
        bundle_identifier = info_plist['CFBundleIdentifier']
        provisioning_profile = FindProvisioningProfile(bundle_identifier, False)
        entitlements = GenerateEntitlements(args.entitlements_path,
                                            provisioning_profile,
                                            bundle_identifier)
        entitlements.WriteTo(args.path)


class FindProvisioningProfileAction(Action):
    """Class implementing the find-codesign-identity action."""

    name = 'find-provisioning-profile'
    help = 'find provisioning profile for use by Xcode project generator'

    @staticmethod
    def _Register(parser):
        parser.add_argument('--bundle-id',
                            '-b',
                            required=True,
                            help='bundle identifier')

    @staticmethod
    def _Execute(args):
        provisioning_profile_info = {}
        provisioning_profile = FindProvisioningProfile(args.bundle_id, False)
        for key in ('team_identifier', 'name'):
            if provisioning_profile:
                provisioning_profile_info[key] = getattr(
                    provisioning_profile, key)
            else:
                provisioning_profile_info[key] = ''
        print(json.dumps(provisioning_profile_info))


def Main():
    # Cache this codec so that plistlib can find it. See
    # https://crbug.com/999461#c12 for more details.
    codecs.lookup('utf-8')

    parser = argparse.ArgumentParser('codesign iOS bundles')
    subparsers = parser.add_subparsers()

    actions = [
        CodeSignBundleAction,
        CodeSignFileAction,
        GenerateEntitlementsAction,
        FindProvisioningProfileAction,
    ]

    for action in actions:
        action.Register(subparsers)

    args = parser.parse_args()
    args.func(args)


if __name__ == '__main__':
    sys.exit(Main())
