# Copyright 2016 The Chromium Authors. All rights reserved.
# 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 == '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([
      'python',
      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())