Merge "releasetools: Clean up AppendVBMetaArgsForPartition()." am: 540a71d2

  img.Write()

def AppendVBMetaArgsForPartition(cmd, partition, img_path, public_key_dir):

  if not img_path:

    return

def AppendVBMetaArgsForPartition(cmd, partition, image):

  """Appends the VBMeta arguments for partition.

  It sets up the VBMeta argument by including the partition descriptor from the

  given 'image', or by configuring the partition as a chained partition.

  Args:

    cmd: A list of command args that will be used to generate the vbmeta image.

        The argument for the partition will be appended to the list.

    partition: The name of the partition (e.g. "system").

    image: The path to the partition image.

  """

  # Check if chain partition is used.

  key_path = OPTIONS.info_dict.get("avb_" + partition + "_key_path")

  if key_path:

    # extract public key in AVB format to be included in vbmeta.img

    avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]

    public_key_path = os.path.join(public_key_dir, "%s.avbpubkey" % partition)

    p = common.Run([avbtool, "extract_public_key", "--key", key_path,

                    "--output", public_key_path],

                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    p.communicate()

    assert p.returncode == 0, \

        "avbtool extract_public_key fail for partition: %r" % partition

    pubkey_path = common.MakeTempFile(prefix="avb-", suffix=".pubkey")

    proc = common.Run(

        [avbtool, "extract_public_key", "--key", key_path, "--output",

         pubkey_path],

        stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

    stdoutdata, _ = proc.communicate()

    assert proc.returncode == 0, \

        "Failed to extract pubkey for {}:\n{}".format(

            partition, stdoutdata)

    rollback_index_location = OPTIONS.info_dict[

        "avb_" + partition + "_rollback_index_location"]

    cmd.extend(["--chain_partition", "%s:%s:%s" % (

        partition, rollback_index_location, public_key_path)])

        partition, rollback_index_location, pubkey_path)])

  else:

    cmd.extend(["--include_descriptors_from_image", img_path])

    cmd.extend(["--include_descriptors_from_image", image])

def AddVBMeta(output_zip, partitions):

  Args:

    output_zip: The output zip file, which needs to be already open.

    partitions: A dict that's keyed by partition names with image paths as

        values. Only valid partition names are accepted, which include 'boot',

        'recovery', 'system', 'vendor', 'dtbo'.

        values. Only valid partition names are accepted, as listed in

        common.AVB_PARTITIONS.

  """

  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "vbmeta.img")

  if os.path.exists(img.input_name):

  cmd = [avbtool, "make_vbmeta_image", "--output", img.name]

  common.AppendAVBSigningArgs(cmd, "vbmeta")

  public_key_dir = common.MakeTempDir(prefix="avbpubkey-")

  for partition, path in partitions.items():

    assert partition in common.AVB_PARTITIONS, 'Unknown partition: %s' % (

        partition,)

    assert os.path.exists(path), 'Failed to find %s for partition %s' % (

        path, partition)

    AppendVBMetaArgsForPartition(cmd, partition, path, public_key_dir)

    assert partition in common.AVB_PARTITIONS, \

        'Unknown partition: {}'.format(partition)

    assert os.path.exists(path), \

        'Failed to find {} for {}'.format(path, partition)

    AppendVBMetaArgsForPartition(cmd, partition, path)

  args = OPTIONS.info_dict.get("avb_vbmeta_args")

  if args and args.strip():

import common

import test_utils

from add_img_to_target_files import (

    AddCareMapTxtForAbOta, AddPackRadioImages, CheckAbOtaImages, GetCareMap)

    AddCareMapTxtForAbOta, AddPackRadioImages, AppendVBMetaArgsForPartition,

    CheckAbOtaImages, GetCareMap)

from rangelib import RangeSet

    # The existing entry should be scheduled to be replaced.

    self.assertIn('META/care_map.txt', OPTIONS.replace_updated_files_list)

  def test_AppendVBMetaArgsForPartition(self):

    OPTIONS.info_dict = {}

    cmd = []

    AppendVBMetaArgsForPartition(cmd, 'system', '/path/to/system.img')

    self.assertEqual(

        ['--include_descriptors_from_image', '/path/to/system.img'], cmd)

  def test_AppendVBMetaArgsForPartition_vendorAsChainedPartition(self):

    testdata_dir = test_utils.get_testdata_dir()

    pubkey = os.path.join(testdata_dir, 'testkey.pubkey.pem')

    OPTIONS.info_dict = {

        'avb_avbtool': 'avbtool',

        'avb_vendor_key_path': pubkey,

        'avb_vendor_rollback_index_location': 5,

    }

    cmd = []

    AppendVBMetaArgsForPartition(cmd, 'vendor', '/path/to/vendor.img')

    self.assertEqual(2, len(cmd))

    self.assertEqual('--chain_partition', cmd[0])

    chained_partition_args = cmd[1].split(':')

    self.assertEqual(3, len(chained_partition_args))

    self.assertEqual('vendor', chained_partition_args[0])

    self.assertEqual('5', chained_partition_args[1])

    self.assertTrue(os.path.exists(chained_partition_args[2]))

  def test_GetCareMap(self):

    sparse_image = test_utils.construct_sparse_image([

        (0xCAC1, 6),