Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit a2ff4c9b authored by Tao Bao's avatar Tao Bao
Browse files

releasetools: Clean up three functions in add_img_to_target_files.py. 

This CL makes the following changes:
 (a) cleans up the similar codes in AddCareMapTxtForAbOta() that handle
     'system' and 'vendor' partitions;
 (b) fixes an issue with the arcname in AddPackRadioImages() and
     AddRadioImagesForAbOta(), where forward slash should always be used
     in zip entry names;
 (c) refactors the branching statements in AddRadioImagesForAbOta() to
     reduce indentation levels.

Test: python -m unittest test_add_img_to_target_files
Test: `m dist` with aosp_marlin-userdebug. Check META/care_map.txt in the
      generated target_files.zip.
Change-Id: I3d6f794962d0c68390fbd18eb13c2622acab3ff5
parent bea20ac7

tools/releasetools/add_img_to_target_files.py

+34 −34
Original line number Diff line number Diff line
@@ -52,7 +52,6 @@ import shlex 
import shutil

import subprocess

import sys

import tempfile

import uuid

import zipfile
@@ -75,6 +74,10 @@ OPTIONS.replace_verity_private_key = False 
OPTIONS.is_signing = False





# Partitions that should have their care_map added to META/care_map.txt.

PARTITIONS_WITH_CARE_MAP = ('system', 'vendor')





class OutputFile(object):

  def __init__(self, output_zip, input_dir, prefix, name):

    self._output_zip = output_zip
@@ -94,13 +97,10 @@ class OutputFile(object): 




def GetCareMap(which, imgname):

  """Generate care_map of system (or vendor) partition"""



  assert which in ("system", "vendor")

  """Generates the care_map for the given partition."""

  assert which in PARTITIONS_WITH_CARE_MAP



  simg = sparse_img.SparseImage(imgname)

  care_map_list = [which]



  care_map_ranges = simg.care_map

  key = which + "_adjusted_partition_size"

  adjusted_blocks = OPTIONS.info_dict.get(key)
@@ -109,8 +109,7 @@ def GetCareMap(which, imgname): 
    care_map_ranges = care_map_ranges.intersect(rangelib.RangeSet(

        "0-%d" % (adjusted_blocks,)))



  care_map_list.append(care_map_ranges.to_string_raw())

  return care_map_list

  return [which, care_map_ranges.to_string_raw()]





def AddSystem(output_zip, prefix="IMAGES/", recovery_img=None, boot_img=None):
@@ -490,21 +489,23 @@ def AddRadioImagesForAbOta(output_zip, ab_partitions): 
    img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)

    if os.path.exists(img_radio_path):

      if output_zip:

        common.ZipWrite(output_zip, img_radio_path,

                        os.path.join("IMAGES", img_name))

        common.ZipWrite(output_zip, img_radio_path, "IMAGES/" + img_name)

      else:

        shutil.copy(img_radio_path, prebuilt_path)

    else:

      continue



    # Walk through VENDOR_IMAGES/ since files could be under subdirs.

    img_vendor_dir = os.path.join(OPTIONS.input_tmp, "VENDOR_IMAGES")

    for root, _, files in os.walk(img_vendor_dir):

      if img_name in files:

        if output_zip:

          common.ZipWrite(output_zip, os.path.join(root, img_name),

                            os.path.join("IMAGES", img_name))

                          "IMAGES/" + img_name)

        else:

          shutil.copy(os.path.join(root, img_name), prebuilt_path)

        break



    # Assert that the image is present under IMAGES/ now.

    if output_zip:

      # Zip spec says: All slashes MUST be forward slashes.

      img_path = 'IMAGES/' + img_name
@@ -526,18 +527,16 @@ def AddCareMapTxtForAbOta(output_zip, ab_partitions, image_paths): 
  care_map_list = []

  for partition in ab_partitions:

    partition = partition.strip()

    if (partition == "system" and

        ("system_verity_block_device" in OPTIONS.info_dict or

         OPTIONS.info_dict.get("avb_system_hashtree_enable") == "true")):

      system_img_path = image_paths[partition]

      assert os.path.exists(system_img_path)

      care_map_list += GetCareMap("system", system_img_path)

    if (partition == "vendor" and

        ("vendor_verity_block_device" in OPTIONS.info_dict or

         OPTIONS.info_dict.get("avb_vendor_hashtree_enable") == "true")):

      vendor_img_path = image_paths[partition]

      assert os.path.exists(vendor_img_path)

      care_map_list += GetCareMap("vendor", vendor_img_path)

    if partition not in PARTITIONS_WITH_CARE_MAP:

      continue



    verity_block_device = "{}_verity_block_device".format(partition)

    avb_hashtree_enable = "avb_{}_hashtree_enable".format(partition)

    if (verity_block_device in OPTIONS.info_dict or

        OPTIONS.info_dict.get(avb_hashtree_enable) == "true"):

      image_path = image_paths[partition]

      assert os.path.exists(image_path)

      care_map_list += GetCareMap(partition, image_path)



  if care_map_list:

    care_map_path = "META/care_map.txt"
@@ -566,6 +565,7 @@ def AddPackRadioImages(output_zip, images): 
    _, ext = os.path.splitext(img_name)

    if not ext:

      img_name += ".img"



    prebuilt_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)

    if os.path.exists(prebuilt_path):

      print("%s already exists, no need to overwrite..." % (img_name,))
@@ -574,9 +574,9 @@ def AddPackRadioImages(output_zip, images): 
    img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)

    assert os.path.exists(img_radio_path), \

        "Failed to find %s at %s" % (img_name, img_radio_path)



    if output_zip:

      common.ZipWrite(output_zip, img_radio_path,

                      os.path.join("IMAGES", img_name))

      common.ZipWrite(output_zip, img_radio_path, "IMAGES/" + img_name)

    else:

      shutil.copy(img_radio_path, prebuilt_path)