Merge "Fix custom image OTA generation with extracted target files" into main

  return input_file

def GetTargetFilesZipForCustomImagesUpdates(input_file, custom_images):

def GetTargetFilesZipForCustomImagesUpdates(input_file, custom_images: dict):

  """Returns a target-files.zip for custom partitions update.

  This function modifies ab_partitions list with the desired custom partitions

  and puts the custom images into the target target-files.zip.

  Args:

    input_file: The input target-files.zip filename.

    input_file: The input target-files extracted directory

    custom_images: A map of custom partitions and custom images.

  Returns:

    The filename of a target-files.zip which has renamed the custom images in

    the IMAGES/ to their partition names.

    The extracted dir of a target-files.zip which has renamed the custom images

    in the IMAGES/ to their partition names.

  """

  for custom_image in custom_images.values():

    if not os.path.exists(os.path.join(input_file, "IMAGES", custom_image)):

      raise ValueError("Specified custom image {} not found in target files {}, available images are {}",

                       custom_image, input_file, os.listdir(os.path.join(input_file, "IMAGES")))

  # First pass: use zip2zip to copy the target files contents, excluding

  # the "custom" images that will be replaced.

  target_file = common.MakeTempFile(prefix="targetfiles-", suffix=".zip")

  cmd = ['zip2zip', '-i', input_file, '-o', target_file]

  images = {}

  for custom_partition, custom_image in custom_images.items():

    default_custom_image = '{}.img'.format(custom_partition)

    if default_custom_image != custom_image:

      src = 'IMAGES/' + custom_image

      dst = 'IMAGES/' + default_custom_image

      cmd.extend(['-x', dst])

      images[dst] = src

  common.RunAndCheckOutput(cmd)

  # Second pass: write {custom_image}.img as {custom_partition}.img.

  with zipfile.ZipFile(input_file, allowZip64=True) as input_zip:

    with zipfile.ZipFile(target_file, 'a', allowZip64=True) as output_zip:

      for dst, src in images.items():

        data = input_zip.read(src)

        logger.info("Update custom partition '%s'", dst)

        common.ZipWriteStr(output_zip, dst, data)

      output_zip.close()

      src = os.path.join(input_file, 'IMAGES', custom_image)

      dst = os.path.join(input_file, 'IMAGES', default_custom_image)

      os.rename(src, dst)

  return target_file

  return input_file

def GeneratePartitionTimestampFlags(partition_state):