Merge "releasetools: Allow building AVB-enabled recovery.img."

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

def AddVBMeta(output_zip, boot_img_path, system_img_path, vendor_img_path,

              dtbo_img_path, prefix="IMAGES/"):

  """Create a VBMeta image and store it in output_zip."""

def AddVBMeta(output_zip, partitions, prefix="IMAGES/"):

  """Creates a VBMeta image and store it in output_zip.

  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'.

  """

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

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

    print("vbmeta.img already exists in %s; not rebuilding..." % (prefix,))

  public_key_dir = tempfile.mkdtemp(prefix="avbpubkey-")

  OPTIONS.tempfiles.append(public_key_dir)

  AppendVBMetaArgsForPartition(cmd, "boot", boot_img_path, public_key_dir)

  AppendVBMetaArgsForPartition(cmd, "system", system_img_path, public_key_dir)

  AppendVBMetaArgsForPartition(cmd, "vendor", vendor_img_path, public_key_dir)

  AppendVBMetaArgsForPartition(cmd, "dtbo", dtbo_img_path, public_key_dir)

  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)

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

  if args and args.strip():

    if fp:

      OPTIONS.info_dict["avb_salt"] = hashlib.sha256(fp).hexdigest()

  # A map between partition names and their paths, which could be used when

  # generating AVB vbmeta image.

  partitions = dict()

  def banner(s):

    print("\n\n++++ " + s + " ++++\n\n")

      "IMAGES/boot.img", "boot.img", OPTIONS.input_tmp, "BOOT")

  # boot.img may be unavailable in some targets (e.g. aosp_arm64).

  if boot_image:

    boot_img_path = os.path.join(OPTIONS.input_tmp, "IMAGES", "boot.img")

    if not os.path.exists(boot_img_path):

    partitions['boot'] = os.path.join(OPTIONS.input_tmp, "IMAGES", "boot.img")

    if not os.path.exists(partitions['boot']):

      boot_image.WriteToDir(OPTIONS.input_tmp)

      if output_zip:

        boot_image.AddToZip(output_zip)

    recovery_image = common.GetBootableImage(

        "IMAGES/recovery.img", "recovery.img", OPTIONS.input_tmp, "RECOVERY")

    assert recovery_image, "Failed to create recovery.img."

    recovery_img_path = os.path.join(

    partitions['recovery'] = os.path.join(

        OPTIONS.input_tmp, "IMAGES", "recovery.img")

    if not os.path.exists(recovery_img_path):

    if not os.path.exists(partitions['recovery']):

      recovery_image.WriteToDir(OPTIONS.input_tmp)

      if output_zip:

        recovery_image.AddToZip(output_zip)

          recovery_two_step_image.AddToZip(output_zip)

  banner("system")

  system_img_path = AddSystem(

  partitions['system'] = system_img_path = AddSystem(

      output_zip, recovery_img=recovery_image, boot_img=boot_image)

  vendor_img_path = None

  if has_vendor:

    banner("vendor")

    vendor_img_path = AddVendor(output_zip)

    partitions['vendor'] = vendor_img_path = AddVendor(output_zip)

  if has_system_other:

    banner("system_other")

    AddSystemOther(output_zip)

  if not OPTIONS.is_signing:

    banner("userdata")

    AddUserdata(output_zip)

    banner("partition-table")

    AddPartitionTable(output_zip)

  dtbo_img_path = None

  if OPTIONS.info_dict.get("has_dtbo") == "true":

    banner("dtbo")

    dtbo_img_path = AddDtbo(output_zip)

    partitions['dtbo'] = AddDtbo(output_zip)

  if OPTIONS.info_dict.get("avb_enable") == "true":

    banner("vbmeta")

    AddVBMeta(output_zip, boot_img_path, system_img_path, vendor_img_path,

              dtbo_img_path)

    AddVBMeta(output_zip, partitions)

  # For devices using A/B update, copy over images from RADIO/ and/or

  # VENDOR_IMAGES/ to IMAGES/ and make sure we have all the needed

  else:

    cmd.extend(["--output", img.name])

  # "boot" or "recovery", without extension.

  partition_name = os.path.basename(sourcedir).lower()

  p = Run(cmd, stdout=subprocess.PIPE)

  p.communicate()

  assert p.returncode == 0, "mkbootimg of %s image failed" % (

      os.path.basename(sourcedir),)

  assert p.returncode == 0, "mkbootimg of %s image failed" % (partition_name,)

  if (info_dict.get("boot_signer", None) == "true" and

      info_dict.get("verity_key", None)):

    if two_step_image:

      path = "/boot"

    else:

      path = "/" + os.path.basename(sourcedir).lower()

      path = "/" + partition_name

    cmd = [OPTIONS.boot_signer_path]

    cmd.extend(OPTIONS.boot_signer_args)

    cmd.extend([path, img.name,

  # Sign the image if vboot is non-empty.

  elif info_dict.get("vboot", None):

    path = "/" + os.path.basename(sourcedir).lower()

    path = "/" + partition_name

    img_keyblock = tempfile.NamedTemporaryFile()

    # We have switched from the prebuilt futility binary to using the tool

    # (futility-host) built from the source. Override the setting in the old

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

    part_size = info_dict["boot_size"]

    cmd = [avbtool, "add_hash_footer", "--image", img.name,

           "--partition_size", str(part_size), "--partition_name", "boot"]

    AppendAVBSigningArgs(cmd, "boot")

           "--partition_size", str(part_size), "--partition_name",

           partition_name]

    AppendAVBSigningArgs(cmd, partition_name)

    args = info_dict.get("avb_boot_add_hash_footer_args")

    if args and args.strip():

      cmd.extend(shlex.split(args))

    p = Run(cmd, stdout=subprocess.PIPE)

    p.communicate()

    assert p.returncode == 0, "avbtool add_hash_footer of %s failed" % (

        os.path.basename(OPTIONS.input_tmp))

        partition_name,)

  img.seek(os.SEEK_SET, 0)

  data = img.read()