Merge "Resign apks contained in apex"

import os.path

import re

import shlex

import shutil

import zipfile

import common

    Exception.__init__(self, message)

class ApexApkSigner(object):

  """Class to sign the apk files in a apex payload image and repack the apex"""

  def __init__(self, apex_path, key_passwords, codename_to_api_level_map):

    self.apex_path = apex_path

    self.key_passwords = key_passwords

    self.codename_to_api_level_map = codename_to_api_level_map

  def ProcessApexFile(self, apk_keys, payload_key, payload_public_key):

    """Scans and signs the apk files and repack the apex

    Args:

      apk_keys: A dict that holds the signing keys for apk files.

      payload_key: The path to the apex payload signing key.

      payload_public_key: The path to the public key corresponding to the

       payload signing key.

    Returns:

      The repacked apex file containing the signed apk files.

    """

    list_cmd = ['deapexer', 'list', self.apex_path]

    entries_names = common.RunAndCheckOutput(list_cmd).split()

    apk_entries = [name for name in entries_names if name.endswith('.apk')]

    # No need to sign and repack, return the original apex path.

    if not apk_entries:

      logger.info('No apk file to sign in %s', self.apex_path)

      return self.apex_path

    for entry in apk_entries:

      apk_name = os.path.basename(entry)

      if apk_name not in apk_keys:

        raise ApexSigningError('Failed to find signing keys for apk file {} in'

                               ' apex {}.  Use "-e <apkname>=" to specify a key'

                               .format(entry, self.apex_path))

      if not any(dirname in entry for dirname in ['app/', 'priv-app/',

                                                  'overlay/']):

        logger.warning('Apk path does not contain the intended directory name:'

                       ' %s', entry)

    payload_dir, has_signed_apk = self.ExtractApexPayloadAndSignApks(

        apk_entries, apk_keys)

    if not has_signed_apk:

      logger.info('No apk file has been signed in %s', self.apex_path)

      return self.apex_path

    return self.RepackApexPayload(payload_dir, payload_key, payload_public_key)

  def ExtractApexPayloadAndSignApks(self, apk_entries, apk_keys):

    """Extracts the payload image and signs the containing apk files."""

    payload_dir = common.MakeTempDir()

    extract_cmd = ['deapexer', 'extract', self.apex_path, payload_dir]

    common.RunAndCheckOutput(extract_cmd)

    has_signed_apk = False

    for entry in apk_entries:

      apk_path = os.path.join(payload_dir, entry)

      assert os.path.exists(self.apex_path)

      key_name = apk_keys.get(os.path.basename(entry))

      if key_name in common.SPECIAL_CERT_STRINGS:

        logger.info('Not signing: %s due to special cert string', apk_path)

        continue

      logger.info('Signing apk file %s in apex %s', apk_path, self.apex_path)

      # Rename the unsigned apk and overwrite the original apk path with the

      # signed apk file.

      unsigned_apk = common.MakeTempFile()

      os.rename(apk_path, unsigned_apk)

      common.SignFile(unsigned_apk, apk_path, key_name, self.key_passwords,

                      codename_to_api_level_map=self.codename_to_api_level_map)

      has_signed_apk = True

    return payload_dir, has_signed_apk

  def RepackApexPayload(self, payload_dir, payload_key, payload_public_key):

    """Rebuilds the apex file with the updated payload directory."""

    apex_dir = common.MakeTempDir()

    # Extract the apex file and reuse its meta files as repack parameters.

    common.UnzipToDir(self.apex_path, apex_dir)

    android_jar_path = common.OPTIONS.android_jar_path

    if not android_jar_path:

      android_jar_path = os.path.join(os.environ.get(

          'ANDROID_BUILD_TOP'), 'prebuilts/sdk/current/public/android.jar')

      logger.warning('android_jar_path not found in options, falling back to'

                     ' use %s', android_jar_path)

    arguments_dict = {

        'manifest': os.path.join(apex_dir, 'apex_manifest.pb'),

        'build_info': os.path.join(apex_dir, 'apex_build_info.pb'),

        'assets_dir': os.path.join(apex_dir, 'assets'),

        'android_jar_path': android_jar_path,

        'key': payload_key,

        'pubkey': payload_public_key,

    }

    for filename in arguments_dict.values():

      assert os.path.exists(filename), 'file {} not found'.format(filename)

    # The repack process will add back these files later in the payload image.

    for name in ['apex_manifest.pb', 'apex_manifest.json', 'lost+found']:

      path = os.path.join(payload_dir, name)

      if os.path.isfile(path):

        os.remove(path)

      elif os.path.isdir(path):

        shutil.rmtree(path)

    repacked_apex = common.MakeTempFile(suffix='.apex')

    repack_cmd = ['apexer', '--force', '--include_build_info',

                  '--do_not_check_keyname', '--apexer_tool_path',

                  os.getenv('PATH')]

    for key, val in arguments_dict.items():

      repack_cmd.append('--' + key)

      repack_cmd.append(val)

    manifest_json = os.path.join(apex_dir, 'apex_manifest.json')

    if os.path.exists(manifest_json):

      repack_cmd.append('--manifest_json')

      repack_cmd.append(manifest_json)

    repack_cmd.append(payload_dir)

    repack_cmd.append(repacked_apex)

    common.RunAndCheckOutput(repack_cmd)

    return repacked_apex

def SignApexPayload(avbtool, payload_file, payload_key_path, payload_key_name,

                    algorithm, salt, no_hashtree, signing_args=None):

  """Signs a given payload_file with the payload key."""

def SignApex(avbtool, apex_data, payload_key, container_key, container_pw,

             codename_to_api_level_map, no_hashtree, signing_args=None):

             apk_keys, codename_to_api_level_map,

             no_hashtree, signing_args=None):

  """Signs the current APEX with the given payload/container keys.

  Args:

    payload_key: The path to payload signing key (w/ extension).

    container_key: The path to container signing key (w/o extension).

    container_pw: The matching password of the container_key, or None.

    apk_keys: A dict that holds the signing keys for apk files.

    codename_to_api_level_map: A dict that maps from codename to API level.

    no_hashtree: Don't include hashtree in the signed APEX.

    signing_args: Additional args to be passed to the payload signer.

  APEX_PAYLOAD_IMAGE = 'apex_payload.img'

  APEX_PUBKEY = 'apex_pubkey'

  # 1a. Extract and sign the APEX_PAYLOAD_IMAGE entry with the given

  # 1. Extract the apex payload image and sign the containing apk files. Repack

  # the apex file after signing.

  payload_public_key = common.ExtractAvbPublicKey(avbtool, payload_key)

  apk_signer = ApexApkSigner(apex_file, container_pw,

                             codename_to_api_level_map)

  apex_file = apk_signer.ProcessApexFile(apk_keys, payload_key,

                                         payload_public_key)

  # 2a. Extract and sign the APEX_PAYLOAD_IMAGE entry with the given

  # payload_key.

  payload_dir = common.MakeTempDir(prefix='apex-payload-')

  with zipfile.ZipFile(apex_file) as apex_fd:

      no_hashtree,

      signing_args)

  # 1b. Update the embedded payload public key.

  payload_public_key = common.ExtractAvbPublicKey(avbtool, payload_key)

  # 2b. Update the embedded payload public key.

  common.ZipDelete(apex_file, APEX_PAYLOAD_IMAGE)

  if APEX_PUBKEY in zip_items:

  common.ZipWrite(apex_zip, payload_public_key, arcname=APEX_PUBKEY)

  common.ZipClose(apex_zip)

  # 2. Align the files at page boundary (same as in apexer).

  # 3. Align the files at page boundary (same as in apexer).

  aligned_apex = common.MakeTempFile(prefix='apex-container-', suffix='.apex')

  common.RunAndCheckOutput(['zipalign', '-f', '4096', apex_file, aligned_apex])

  # 3. Sign the APEX container with container_key.

  # 4. Sign the APEX container with container_key.

  signed_apex = common.MakeTempFile(prefix='apex-container-', suffix='.apex')

  # Specify the 4K alignment when calling SignApk.

    self.extra_signapk_args = []

    self.java_path = "java"  # Use the one on the path by default.

    self.java_args = ["-Xmx2048m"]  # The default JVM args.

    self.android_jar_path = None

    self.public_key_suffix = ".x509.pem"

    self.private_key_suffix = ".pk8"

    # use otatools built boot_signer by default

        argv, "hvp:s:x:" + extra_opts,

        ["help", "verbose", "path=", "signapk_path=",

         "signapk_shared_library_path=", "extra_signapk_args=",

         "java_path=", "java_args=", "public_key_suffix=",

         "java_path=", "java_args=", "android_jar_path=", "public_key_suffix=",

         "private_key_suffix=", "boot_signer_path=", "boot_signer_args=",

         "verity_signer_path=", "verity_signer_args=", "device_specific=",

         "extra=", "logfile=", "aftl_server=", "aftl_key_path=",

      OPTIONS.java_path = a

    elif o in ("--java_args",):

      OPTIONS.java_args = shlex.split(a)

    elif o in ("--android_jar_path",):

      OPTIONS.android_jar_path = a

    elif o in ("--public_key_suffix",):

      OPTIONS.public_key_suffix = a

    elif o in ("--private_key_suffix",):

  --payload_extra_args <args>

      Optional flag that specifies any extra args to be passed to payload signer

      (e.g. --payload_extra_args="--signing_helper_with_files /path/to/helper").

  -e  (--extra_apks)  <name,name,...=key>

      Add extra APK name/key pairs. This is useful to sign the apk files in the

      apex payload image.

"""

import logging

logger = logging.getLogger(__name__)

def SignApexFile(avbtool, apex_file, payload_key, container_key,

                 no_hashtree, signing_args=None):

def SignApexFile(avbtool, apex_file, payload_key, container_key, no_hashtree,

                 apk_keys=None, signing_args=None):

  """Signs the given apex file."""

  with open(apex_file, 'rb') as input_fp:

    apex_data = input_fp.read()

      container_pw=None,

      codename_to_api_level_map=None,

      no_hashtree=no_hashtree,

      apk_keys=apk_keys,

      signing_args=signing_args)

      options['payload_key'] = a

    elif o == '--payload_extra_args':

      options['payload_extra_args'] = a

    elif o in ("-e", "--extra_apks"):

      names, key = a.split("=")

      names = names.split(",")

      for n in names:

        if 'extra_apks' not in options:

          options['extra_apks'] = {}

        options['extra_apks'].update({n: key})

    else:

      return False

    return True

  args = common.ParseOptions(

      argv, __doc__,

      extra_opts='',

      extra_opts='e:',

      extra_long_opts=[

          'avbtool=',

          'container_key=',

          'payload_extra_args=',

          'payload_key=',

          'extra_apks=',

      ],

      extra_option_handler=option_handler)

      options['payload_key'],

      options['container_key'],

      no_hashtree=False,

      apk_keys=options.get('extra_apks', {}),

      signing_args=options.get('payload_extra_args'))

  shutil.copyfile(signed_apex, args[1])

  logger.info("done.")

      Specify any additional args that are needed to AVB-sign the image

      (e.g. "--signing_helper /path/to/helper"). The args will be appended to

      the existing ones in info dict.

  --android_jar_path <path>

      Path to the android.jar to repack the apex file.

"""

from __future__ import print_function

OPTIONS.avb_keys = {}

OPTIONS.avb_algorithms = {}

OPTIONS.avb_extra_args = {}

OPTIONS.android_jar_path = None

AVB_FOOTER_ARGS_BY_PARTITION = {

            payload_key,

            container_key,

            key_passwords[container_key],

            apk_keys,

            codename_to_api_level_map,

            no_hashtree=True,

            signing_args=OPTIONS.avb_extra_args.get('apex'))

  apex_keys_info = ReadApexKeysInfo(input_zip)

  apex_keys = GetApexKeys(apex_keys_info, apk_keys)

  # TODO(xunchang) check for the apks inside the apex files, and abort early if

  # the keys are not available.

  CheckApkAndApexKeysAvailable(

      input_zip,

      set(apk_keys.keys()) | set(apex_keys.keys()),

    # The default payload signing key.

    self.payload_key = os.path.join(self.testdata_dir, 'testkey.key')

    common.OPTIONS.search_path = test_utils.get_search_path()

  @staticmethod

  def _GetTestPayload():

    payload_file = common.MakeTempFile(prefix='apex-', suffix='.img')

        payload_file,

        os.path.join(self.testdata_dir, 'testkey_with_passwd.key'),

        no_hashtree=True)

  @test_utils.SkipIfExternalToolsUnavailable()

  def test_ApexApkSigner_noApkPresent(self):

    apex_path = os.path.join(self.testdata_dir, 'foo.apex')

    signer = apex_utils.ApexApkSigner(apex_path, None, None)

    processed_apex = signer.ProcessApexFile({}, self.payload_key,

                                            None)

    self.assertEqual(apex_path, processed_apex)

  @test_utils.SkipIfExternalToolsUnavailable()

  def test_ApexApkSigner_apkKeyNotPresent(self):

    apex_path = os.path.join(self.testdata_dir, 'has_apk.apex')

    signer = apex_utils.ApexApkSigner(apex_path, None, None)

    self.assertRaises(apex_utils.ApexSigningError, signer.ProcessApexFile, {},

                      self.payload_key, None)

  @test_utils.SkipIfExternalToolsUnavailable()

  def test_ApexApkSigner_signApk(self):

    apex_path = os.path.join(self.testdata_dir, 'has_apk.apex')

    signer = apex_utils.ApexApkSigner(apex_path, None, None)

    apk_keys = {'wifi-service-resources.apk': os.path.join(

        self.testdata_dir, 'testkey')}

    self.payload_key = os.path.join(self.testdata_dir, 'testkey_RSA4096.key')

    payload_pubkey = common.ExtractAvbPublicKey('avbtool',

                                                self.payload_key)

    signer.ProcessApexFile(apk_keys, self.payload_key, payload_pubkey)