Merge "Reland Support incremental dev option OTA during signing" into main

        "releasetools_common",

        "ota_metadata_proto",

        "ota_utils_lib",

        "update_payload",

    ],

}

import add_img_to_target_files

import ota_from_raw_img

import ota_utils

import apex_utils

import common

import payload_signer

import update_payload

from payload_signer import SignOtaPackage, PAYLOAD_BIN

OPTIONS.allow_gsi_debug_sepolicy = False

OPTIONS.override_apk_keys = None

OPTIONS.override_apex_keys = None

OPTIONS.input_tmp = None

AVB_FOOTER_ARGS_BY_PARTITION = {

def RegenerateKernelPartitions(input_tf_zip: zipfile.ZipFile, output_tf_zip: zipfile.ZipFile, misc_info):

  """Re-generate boot and dtbo partitions using new signing configuration"""

  files_to_unzip = [

      "PREBUILT_IMAGES/*", "BOOTABLE_IMAGES/*.img", "*/boot_16k.img", "*/dtbo_16k.img"]

  if OPTIONS.input_tmp is None:

    OPTIONS.input_tmp = common.UnzipTemp(input_tf_zip.filename, [

                                "*/boot.img", "*/dtbo.img"])

    OPTIONS.input_tmp = common.UnzipTemp(input_tf_zip.filename, files_to_unzip)

  else:

    common.UnzipToDir(input_tf_zip, OPTIONS.input_tmp, [

                                "*/boot.img", "*/dtbo.img"])

    common.UnzipToDir(input_tf_zip.filename, OPTIONS.input_tmp, files_to_unzip)

  unzip_dir = OPTIONS.input_tmp

  image_dir = os.path.join(unzip_dir, "IMAGES")

  shutil.rmtree(image_dir)

  os.makedirs(image_dir, exist_ok=True)

  boot_image = common.GetBootableImage(

      "IMAGES/boot.img", "boot.img", unzip_dir, "BOOT", misc_info)

    boot_image = os.path.join(unzip_dir, boot_image.name)

    common.ZipWrite(output_tf_zip, boot_image, "IMAGES/boot.img",

                    compress_type=zipfile.ZIP_STORED)

  if misc_info.get("has_dtbo") == "true":

    add_img_to_target_files.AddDtbo(output_tf_zip)

  return unzip_dir

def RegenerateBootOTA(input_tf_zip: zipfile.ZipFile, output_tf_zip: zipfile.ZipFile, misc_info, filename, input_ota):

  if filename not in ["VENDOR/boot_otas/boot_ota_4k.zip", "SYSTEM/boot_otas/boot_ota_4k.zip"]:

    # We only need to re-generate 4K boot OTA, for other OTA packages

    # simply copy as is

def RegenerateBootOTA(input_tf_zip: zipfile.ZipFile, filename, input_ota):

  with input_tf_zip.open(filename, "r") as in_fp:

      shutil.copyfileobj(in_fp, input_ota)

      input_ota.flush()

    payload = update_payload.Payload(in_fp)

  is_incremental = any([part.HasField('old_partition_info')

                        for part in payload.manifest.partitions])

  is_boot_ota = filename.startswith(

      "VENDOR/boot_otas/") or filename.startswith("SYSTEM/boot_otas/")

  if not is_boot_ota:

    return

  is_4k_boot_ota = filename in [

      "VENDOR/boot_otas/boot_ota_4k.zip", "SYSTEM/boot_otas/boot_ota_4k.zip"]

  # Only 4K boot image is re-generated, so if 16K boot ota isn't incremental,

  # we do not need to re-generate

  if not is_4k_boot_ota and not is_incremental:

    return

  timestamp = misc_info["build.prop"].GetProp(

      "ro.system.build.date.utc")

  unzip_dir = RegenerateKernelPartitions(

      input_tf_zip, output_tf_zip, misc_info)

  signed_boot_image = os.path.join(unzip_dir, "IMAGES/boot.img")

  signed_dtbo_image = os.path.join(unzip_dir, "IMAGES/dtbo.img")

  timestamp = str(payload.manifest.max_timestamp)

  partitions = [part.partition_name for part in payload.manifest.partitions]

  unzip_dir = OPTIONS.input_tmp

  signed_boot_image = os.path.join(unzip_dir, "IMAGES", "boot.img")

  if not os.path.exists(signed_boot_image):

    logger.warn("Need to re-generate boot OTA {} but failed to get signed boot image. 16K dev option will be impacted, after rolling back to 4K user would need to sideload/flash their device to continue receiving OTAs.")

    return

  logger.info(

      "Re-generating boot OTA {} with timestamp {}".format(filename, timestamp))

  signed_dtbo_image = os.path.join(unzip_dir, "IMAGES", "dtbo.img")

  if "dtbo" in partitions and not os.path.exists(signed_dtbo_image):

    raise ValueError(

        "Boot OTA {} has dtbo partition, but no dtbo image found in target files.".format(filename))

  if is_incremental:

    signed_16k_boot_image = os.path.join(

        unzip_dir, "IMAGES", "boot_16k.img")

    signed_16k_dtbo_image = os.path.join(

        unzip_dir, "IMAGES", "dtbo_16k.img")

    if is_4k_boot_ota:

      if os.path.exists(signed_16k_boot_image):

        signed_boot_image = signed_16k_boot_image + ":" + signed_boot_image

      if os.path.exists(signed_16k_dtbo_image):

        signed_dtbo_image = signed_16k_dtbo_image + ":" + signed_dtbo_image

    else:

      if os.path.exists(signed_16k_boot_image):

        signed_boot_image += ":" + signed_16k_boot_image

      if os.path.exists(signed_16k_dtbo_image):

        signed_dtbo_image += ":" + signed_16k_dtbo_image

  args = ["ota_from_raw_img", "--package_key", OPTIONS.package_key,

          "--max_timestamp", timestamp, "--output", input_ota.name]

  if os.path.exists(signed_dtbo_image):

  if "dtbo" in partitions:

    args.extend(["--partition_name", "boot,dtbo",

                signed_boot_image, signed_dtbo_image])

  else:

    args.extend(["--partition_name", "boot", signed_boot_image])

  logger.info(

      "Re-generating boot OTA {} using cmd {}".format(filename, args))

  ota_from_raw_img.main(args)

  if misc_info.get('avb_enable') == 'true':

    RewriteAvbProps(misc_info)

  RegenerateKernelPartitions(input_tf_zip, output_tf_zip, misc_info)

  for info in input_tf_zip.infolist():

    filename = info.filename

    if filename.startswith("IMAGES/"):

    if filename.startswith("OTA/") and filename.endswith(".img"):

      continue

    data = input_tf_zip.read(filename)

    out_info = copy.copy(info)

    (is_apk, is_compressed, should_be_skipped) = GetApkFileInfo(

        filename, compressed_extension, OPTIONS.skip_apks_with_path_prefix)

    data = input_tf_zip.read(filename)

    out_info = copy.copy(info)

    if is_apk and should_be_skipped:

      # Copy skipped APKs verbatim.

    elif filename.endswith(".zip") and IsEntryOtaPackage(input_tf_zip, filename):

      logger.info("Re-signing OTA package {}".format(filename))

      with tempfile.NamedTemporaryFile() as input_ota, tempfile.NamedTemporaryFile() as output_ota:

        RegenerateBootOTA(input_tf_zip, output_tf_zip,

                          misc_info, filename, input_ota)

        RegenerateBootOTA(input_tf_zip, filename, input_ota)

        SignOtaPackage(input_ota.name, output_ota.name)

        common.ZipWrite(output_tf_zip, output_ota.name, filename,