1#!/usr/bin/env python
2#
3# Copyright (C) 2014 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""
18Given a target-files zipfile that does not contain images (ie, does
19not have an IMAGES/ top-level subdirectory), produce the images and
20add them to the zipfile.
21
22Usage:  add_img_to_target_files [flag] target_files
23
24  -a  (--add_missing)
25      Build and add missing images to "IMAGES/". If this option is
26      not specified, this script will simply exit when "IMAGES/"
27      directory exists in the target file.
28
29  -r  (--rebuild_recovery)
30      Rebuild the recovery patch and write it to the system image. Only
31      meaningful when system image needs to be rebuilt and there're separate
32      boot / recovery images.
33
34  --replace_verity_private_key
35      Replace the private key used for verity signing. (same as the option
36      in sign_target_files_apks)
37
38  --replace_verity_public_key
39       Replace the certificate (public key) used for verity verification. (same
40       as the option in sign_target_files_apks)
41
42  --is_signing
43      Skip building & adding the images for "userdata" and "cache" if we
44      are signing the target files.
45
46  --avb-resolve-rollback-index-location-conflict
47      If provided, resolve the conflict AVB rollback index location when
48      necessary.
49"""
50
51from __future__ import print_function
52
53import avbtool
54import datetime
55import logging
56import os
57import shlex
58import shutil
59import stat
60import sys
61import uuid
62import tempfile
63import zipfile
64
65import build_image
66import build_super_image
67import common
68import verity_utils
69import ota_metadata_pb2
70import rangelib
71import sparse_img
72from concurrent.futures import ThreadPoolExecutor
73from apex_utils import GetApexInfoFromTargetFiles
74from common import ZipDelete, PARTITIONS_WITH_CARE_MAP, ExternalError, RunAndCheckOutput, IsSparseImage, MakeTempFile, ZipWrite
75from build_image import FIXED_FILE_TIMESTAMP
76
77if sys.hexversion < 0x02070000:
78  print("Python 2.7 or newer is required.", file=sys.stderr)
79  sys.exit(1)
80
81logger = logging.getLogger(__name__)
82
83OPTIONS = common.OPTIONS
84OPTIONS.add_missing = False
85OPTIONS.rebuild_recovery = False
86OPTIONS.replace_updated_files_list = []
87OPTIONS.is_signing = False
88OPTIONS.avb_resolve_rollback_index_location_conflict = False
89
90
91def ParseAvbFooter(img_path) -> avbtool.AvbFooter:
92  with open(img_path, 'rb') as fp:
93    fp.seek(-avbtool.AvbFooter.SIZE, os.SEEK_END)
94    data = fp.read(avbtool.AvbFooter.SIZE)
95    return avbtool.AvbFooter(data)
96
97
98def GetCareMap(which, imgname):
99  """Returns the care_map string for the given partition.
100
101  Args:
102    which: The partition name, must be listed in PARTITIONS_WITH_CARE_MAP.
103    imgname: The filename of the image.
104
105  Returns:
106    (which, care_map_ranges): care_map_ranges is the raw string of the care_map
107    RangeSet; or None.
108  """
109  assert which in PARTITIONS_WITH_CARE_MAP
110
111  is_sparse_img = IsSparseImage(imgname)
112  unsparsed_image_size = os.path.getsize(imgname)
113
114  # A verified image contains original image + hash tree data + FEC data
115  # + AVB footer, all concatenated together. The caremap specifies a range
116  # of blocks that update_verifier should read on top of dm-verity device
117  # to verify correctness of OTA updates. When reading off of dm-verity device,
118  # the hashtree and FEC part of image isn't available. So caremap should
119  # only contain the original image blocks.
120  try:
121    avbfooter = None
122    if is_sparse_img:
123      with tempfile.NamedTemporaryFile() as tmpfile:
124        img = sparse_img.SparseImage(imgname)
125        unsparsed_image_size = img.total_blocks * img.blocksize
126        for data in img.ReadBlocks(img.total_blocks - 1, 1):
127          tmpfile.write(data)
128        tmpfile.flush()
129        avbfooter = ParseAvbFooter(tmpfile.name)
130    else:
131      avbfooter = ParseAvbFooter(imgname)
132  except LookupError as e:
133    logger.warning(
134        "Failed to parse avbfooter for partition %s image %s, %s", which, imgname, e)
135    return None
136
137  image_size = avbfooter.original_image_size
138  assert image_size < unsparsed_image_size, f"AVB footer's original image size {image_size} is larger than or equal to image size on disk {unsparsed_image_size}, this can't happen because a verified image = original image + hash tree data + FEC data + avbfooter."
139  assert image_size > 0
140
141  image_blocks = int(image_size) // 4096 - 1
142  # It's OK for image_blocks to be 0, because care map ranges are inclusive.
143  # So 0-0 means "just block 0", which is valid.
144  assert image_blocks >= 0, "blocks for {} must be non-negative, image size: {}".format(
145      which, image_size)
146
147  # For sparse images, we will only check the blocks that are listed in the care
148  # map, i.e. the ones with meaningful data.
149  if is_sparse_img:
150    simg = sparse_img.SparseImage(imgname)
151    care_map_ranges = simg.care_map.intersect(
152        rangelib.RangeSet("0-{}".format(image_blocks)))
153
154  # Otherwise for non-sparse images, we read all the blocks in the filesystem
155  # image.
156  else:
157    care_map_ranges = rangelib.RangeSet("0-{}".format(image_blocks))
158
159  return [which, care_map_ranges.to_string_raw()]
160
161
162def AddCareMapForAbOta(output_file, ab_partitions, image_paths):
163  """Generates and adds care_map.pb for a/b partition that has care_map.
164
165  Args:
166    output_file: The output zip file (needs to be already open),
167        or file path to write care_map.pb.
168    ab_partitions: The list of A/B partitions.
169    image_paths: A map from the partition name to the image path.
170  """
171  if not output_file:
172    raise ExternalError('Expected output_file for AddCareMapForAbOta')
173
174  care_map_list = []
175  for partition in ab_partitions:
176    partition = partition.strip()
177    if partition not in PARTITIONS_WITH_CARE_MAP:
178      continue
179
180    verity_block_device = "{}_verity_block_device".format(partition)
181    avb_hashtree_enable = "avb_{}_hashtree_enable".format(partition)
182    if (verity_block_device in OPTIONS.info_dict or
183            OPTIONS.info_dict.get(avb_hashtree_enable) == "true"):
184      if partition not in image_paths:
185        logger.warning('Potential partition with care_map missing from images: %s',
186                       partition)
187        continue
188      image_path = image_paths[partition]
189      if not os.path.exists(image_path):
190        raise ExternalError('Expected image at path {}'.format(image_path))
191
192      care_map = GetCareMap(partition, image_path)
193      if not care_map:
194        continue
195      care_map_list += care_map
196
197      # adds fingerprint field to the care_map
198      # TODO(xunchang) revisit the fingerprint calculation for care_map.
199      partition_props = OPTIONS.info_dict.get(partition + ".build.prop")
200      prop_name_list = ["ro.{}.build.fingerprint".format(partition),
201                        "ro.{}.build.thumbprint".format(partition)]
202
203      present_props = [x for x in prop_name_list if
204                       partition_props and partition_props.GetProp(x)]
205      if not present_props:
206        logger.warning(
207            "fingerprint is not present for partition %s", partition)
208        property_id, fingerprint = "unknown", "unknown"
209      else:
210        property_id = present_props[0]
211        fingerprint = partition_props.GetProp(property_id)
212      care_map_list += [property_id, fingerprint]
213
214  if not care_map_list:
215    return
216
217  # Converts the list into proto buf message by calling care_map_generator; and
218  # writes the result to a temp file.
219  temp_care_map_text = MakeTempFile(prefix="caremap_text-",
220                                           suffix=".txt")
221  with open(temp_care_map_text, 'w') as text_file:
222    text_file.write('\n'.join(care_map_list))
223
224  temp_care_map = MakeTempFile(prefix="caremap-", suffix=".pb")
225  care_map_gen_cmd = ["care_map_generator", temp_care_map_text, temp_care_map]
226  RunAndCheckOutput(care_map_gen_cmd)
227
228  if not isinstance(output_file, zipfile.ZipFile):
229    shutil.copy(temp_care_map, output_file)
230    return
231  # output_file is a zip file
232  care_map_path = "META/care_map.pb"
233  if care_map_path in output_file.namelist():
234    # Copy the temp file into the OPTIONS.input_tmp dir and update the
235    # replace_updated_files_list used by add_img_to_target_files
236    if not OPTIONS.replace_updated_files_list:
237      OPTIONS.replace_updated_files_list = []
238    shutil.copy(temp_care_map, os.path.join(OPTIONS.input_tmp, care_map_path))
239    OPTIONS.replace_updated_files_list.append(care_map_path)
240  else:
241    ZipWrite(output_file, temp_care_map, arcname=care_map_path)
242
243
244class OutputFile(object):
245  """A helper class to write a generated file to the given dir or zip.
246
247  When generating images, we want the outputs to go into the given zip file, or
248  the given dir.
249
250  Attributes:
251    name: The name of the output file, regardless of the final destination.
252  """
253
254  def __init__(self, output_zip, input_dir, *args):
255    # We write the intermediate output file under the given input_dir, even if
256    # the final destination is a zip archive.
257    self.name = os.path.join(input_dir, *args)
258    self._output_zip = output_zip
259    if self._output_zip:
260      self._zip_name = os.path.join(*args)
261
262  def Write(self, compress_type=None):
263    if self._output_zip:
264      common.ZipWrite(self._output_zip, self.name,
265                      self._zip_name, compress_type=compress_type)
266
267
268def AddSystem(output_zip, recovery_img=None, boot_img=None):
269  """Turn the contents of SYSTEM into a system image and store it in
270  output_zip. Returns the name of the system image file."""
271
272  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "system.img")
273  if os.path.exists(img.name):
274    logger.info("system.img already exists; no need to rebuild...")
275    return img.name
276
277  def output_sink(fn, data):
278    output_file = os.path.join(OPTIONS.input_tmp, "SYSTEM", fn)
279    with open(output_file, "wb") as ofile:
280      ofile.write(data)
281
282    if output_zip:
283      arc_name = "SYSTEM/" + fn
284      if arc_name in output_zip.namelist():
285        OPTIONS.replace_updated_files_list.append(arc_name)
286      else:
287        common.ZipWrite(output_zip, output_file, arc_name)
288
289  board_uses_vendorimage = OPTIONS.info_dict.get(
290      "board_uses_vendorimage") == "true"
291
292  if (OPTIONS.rebuild_recovery and not board_uses_vendorimage and
293          recovery_img is not None and boot_img is not None):
294    logger.info("Building new recovery patch on system at system/vendor")
295    common.MakeRecoveryPatch(OPTIONS.input_tmp, output_sink, recovery_img,
296                             boot_img, info_dict=OPTIONS.info_dict)
297
298  block_list = OutputFile(output_zip, OPTIONS.input_tmp,
299                          "IMAGES", "system.map")
300  CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "system", img,
301              block_list=block_list)
302  return img.name
303
304
305def AddSystemOther(output_zip):
306  """Turn the contents of SYSTEM_OTHER into a system_other image
307  and store it in output_zip."""
308
309  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "system_other.img")
310  if os.path.exists(img.name):
311    logger.info("system_other.img already exists; no need to rebuild...")
312    return
313
314  CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "system_other", img)
315
316
317def AddVendor(output_zip, recovery_img=None, boot_img=None):
318  """Turn the contents of VENDOR into a vendor image and store in it
319  output_zip."""
320
321  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "vendor.img")
322  if os.path.exists(img.name):
323    logger.info("vendor.img already exists; no need to rebuild...")
324    return img.name
325
326  def output_sink(fn, data):
327    output_file = os.path.join(OPTIONS.input_tmp, "VENDOR", fn)
328    with open(output_file, "wb") as ofile:
329      ofile.write(data)
330
331    if output_zip:
332      arc_name = "VENDOR/" + fn
333      if arc_name in output_zip.namelist():
334        OPTIONS.replace_updated_files_list.append(arc_name)
335      else:
336        common.ZipWrite(output_zip, output_file, arc_name)
337
338  board_uses_vendorimage = OPTIONS.info_dict.get(
339      "board_uses_vendorimage") == "true"
340
341  if (OPTIONS.rebuild_recovery and board_uses_vendorimage and
342          recovery_img is not None and boot_img is not None):
343    logger.info("Building new recovery patch on vendor")
344    common.MakeRecoveryPatch(OPTIONS.input_tmp, output_sink, recovery_img,
345                             boot_img, info_dict=OPTIONS.info_dict)
346
347  block_list = OutputFile(output_zip, OPTIONS.input_tmp,
348                          "IMAGES", "vendor.map")
349  CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "vendor", img,
350              block_list=block_list)
351  return img.name
352
353
354def AddProduct(output_zip):
355  """Turn the contents of PRODUCT into a product image and store it in
356  output_zip."""
357
358  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "product.img")
359  if os.path.exists(img.name):
360    logger.info("product.img already exists; no need to rebuild...")
361    return img.name
362
363  block_list = OutputFile(
364      output_zip, OPTIONS.input_tmp, "IMAGES", "product.map")
365  CreateImage(
366      OPTIONS.input_tmp, OPTIONS.info_dict, "product", img,
367      block_list=block_list)
368  return img.name
369
370
371def AddSystemExt(output_zip):
372  """Turn the contents of SYSTEM_EXT into a system_ext image and store it in
373  output_zip."""
374
375  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES",
376                   "system_ext.img")
377  if os.path.exists(img.name):
378    logger.info("system_ext.img already exists; no need to rebuild...")
379    return img.name
380
381  block_list = OutputFile(
382      output_zip, OPTIONS.input_tmp, "IMAGES", "system_ext.map")
383  CreateImage(
384      OPTIONS.input_tmp, OPTIONS.info_dict, "system_ext", img,
385      block_list=block_list)
386  return img.name
387
388
389def AddOdm(output_zip):
390  """Turn the contents of ODM into an odm image and store it in output_zip."""
391
392  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "odm.img")
393  if os.path.exists(img.name):
394    logger.info("odm.img already exists; no need to rebuild...")
395    return img.name
396
397  block_list = OutputFile(
398      output_zip, OPTIONS.input_tmp, "IMAGES", "odm.map")
399  CreateImage(
400      OPTIONS.input_tmp, OPTIONS.info_dict, "odm", img,
401      block_list=block_list)
402  return img.name
403
404
405def AddVendorDlkm(output_zip):
406  """Turn the contents of VENDOR_DLKM into an vendor_dlkm image and store it in output_zip."""
407
408  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "vendor_dlkm.img")
409  if os.path.exists(img.name):
410    logger.info("vendor_dlkm.img already exists; no need to rebuild...")
411    return img.name
412
413  block_list = OutputFile(
414      output_zip, OPTIONS.input_tmp, "IMAGES", "vendor_dlkm.map")
415  CreateImage(
416      OPTIONS.input_tmp, OPTIONS.info_dict, "vendor_dlkm", img,
417      block_list=block_list)
418  return img.name
419
420
421def AddOdmDlkm(output_zip):
422  """Turn the contents of OdmDlkm into an odm_dlkm image and store it in output_zip."""
423
424  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "odm_dlkm.img")
425  if os.path.exists(img.name):
426    logger.info("odm_dlkm.img already exists; no need to rebuild...")
427    return img.name
428
429  block_list = OutputFile(
430      output_zip, OPTIONS.input_tmp, "IMAGES", "odm_dlkm.map")
431  CreateImage(
432      OPTIONS.input_tmp, OPTIONS.info_dict, "odm_dlkm", img,
433      block_list=block_list)
434  return img.name
435
436
437def AddSystemDlkm(output_zip):
438  """Turn the contents of SystemDlkm into an system_dlkm image and store it in output_zip."""
439
440  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "system_dlkm.img")
441  if os.path.exists(img.name):
442    logger.info("system_dlkm.img already exists; no need to rebuild...")
443    return img.name
444
445  block_list = OutputFile(
446      output_zip, OPTIONS.input_tmp, "IMAGES", "system_dlkm.map")
447  CreateImage(
448      OPTIONS.input_tmp, OPTIONS.info_dict, "system_dlkm", img,
449      block_list=block_list)
450  return img.name
451
452
453def AddDtbo(output_zip):
454  """Adds the DTBO image.
455
456  Uses the image under IMAGES/ if it already exists. Otherwise looks for the
457  image under PREBUILT_IMAGES/, signs it as needed, and returns the image name.
458  """
459  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "dtbo.img")
460  if os.path.exists(img.name):
461    logger.info("dtbo.img already exists; no need to rebuild...")
462    return img.name
463
464  dtbo_prebuilt_path = os.path.join(
465      OPTIONS.input_tmp, "PREBUILT_IMAGES", "dtbo.img")
466  assert os.path.exists(dtbo_prebuilt_path)
467  shutil.copy(dtbo_prebuilt_path, img.name)
468
469  # AVB-sign the image as needed.
470  if OPTIONS.info_dict.get("avb_enable") == "true":
471    # Signing requires +w
472    os.chmod(img.name, os.stat(img.name).st_mode | stat.S_IWUSR)
473
474    avbtool = OPTIONS.info_dict["avb_avbtool"]
475    part_size = OPTIONS.info_dict["dtbo_size"]
476    # The AVB hash footer will be replaced if already present.
477    cmd = [avbtool, "add_hash_footer", "--image", img.name,
478           "--partition_size", str(part_size), "--partition_name", "dtbo"]
479    common.AppendAVBSigningArgs(cmd, "dtbo")
480    args = OPTIONS.info_dict.get("avb_dtbo_add_hash_footer_args")
481    if args and args.strip():
482      cmd.extend(shlex.split(args))
483    common.RunAndCheckOutput(cmd)
484
485  img.Write()
486  return img.name
487
488
489def AddPvmfw(output_zip):
490  """Adds the pvmfw image.
491
492  Uses the image under IMAGES/ if it already exists. Otherwise looks for the
493  image under PREBUILT_IMAGES/, signs it as needed, and returns the image name.
494  """
495  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "pvmfw.img")
496  if os.path.exists(img.name):
497    logger.info("pvmfw.img already exists; no need to rebuild...")
498    return img.name
499
500  pvmfw_prebuilt_path = os.path.join(
501      OPTIONS.input_tmp, "PREBUILT_IMAGES", "pvmfw.img")
502  assert os.path.exists(pvmfw_prebuilt_path)
503  shutil.copy(pvmfw_prebuilt_path, img.name)
504
505  # AVB-sign the image as needed.
506  if OPTIONS.info_dict.get("avb_enable") == "true":
507    # Signing requires +w
508    os.chmod(img.name, os.stat(img.name).st_mode | stat.S_IWUSR)
509
510    avbtool = OPTIONS.info_dict["avb_avbtool"]
511    part_size = OPTIONS.info_dict["pvmfw_size"]
512    # The AVB hash footer will be replaced if already present.
513    cmd = [avbtool, "add_hash_footer", "--image", img.name,
514           "--partition_size", str(part_size), "--partition_name", "pvmfw"]
515    common.AppendAVBSigningArgs(cmd, "pvmfw")
516    args = OPTIONS.info_dict.get("avb_pvmfw_add_hash_footer_args")
517    if args and args.strip():
518      cmd.extend(shlex.split(args))
519    common.RunAndCheckOutput(cmd)
520
521  img.Write()
522  return img.name
523
524
525def AddCustomImages(output_zip, partition_name, image_list):
526  """Adds and signs avb custom images as needed in IMAGES/.
527
528  Args:
529    output_zip: The output zip file (needs to be already open), or None to
530        write images to OPTIONS.input_tmp/.
531    partition_name: The custom image partition name.
532    image_list: The image list of the custom image partition.
533
534  Uses the image under IMAGES/ if it already exists. Otherwise looks for the
535  image under PREBUILT_IMAGES/, signs it as needed, and returns the image name.
536
537  Raises:
538    AssertionError: If image can't be found.
539  """
540
541  builder = None
542  key_path = OPTIONS.info_dict.get("avb_{}_key_path".format(partition_name))
543  if key_path is not None:
544    algorithm = OPTIONS.info_dict.get("avb_{}_algorithm".format(partition_name))
545    extra_args = OPTIONS.info_dict.get(
546        "avb_{}_add_hashtree_footer_args".format(partition_name))
547    partition_size = OPTIONS.info_dict.get(
548        "avb_{}_partition_size".format(partition_name))
549
550    builder = verity_utils.CreateCustomImageBuilder(
551        OPTIONS.info_dict, partition_name, partition_size,
552        key_path, algorithm, extra_args)
553
554  for img_name in image_list:
555    custom_image = OutputFile(
556        output_zip, OPTIONS.input_tmp, "IMAGES", img_name)
557    if os.path.exists(custom_image.name):
558      continue
559
560    custom_image_prebuilt_path = os.path.join(
561        OPTIONS.input_tmp, "PREBUILT_IMAGES", img_name)
562    assert os.path.exists(custom_image_prebuilt_path), \
563        "Failed to find %s at %s" % (img_name, custom_image_prebuilt_path)
564
565    shutil.copy(custom_image_prebuilt_path, custom_image.name)
566
567    if builder is not None:
568      builder.Build(custom_image.name)
569
570    custom_image.Write()
571
572  default = os.path.join(OPTIONS.input_tmp, "IMAGES", partition_name + ".img")
573  assert os.path.exists(default), \
574      "There should be one %s.img" % (partition_name)
575  return default
576
577
578def CreateImage(input_dir, info_dict, what, output_file, block_list=None):
579  logger.info("creating %s.img...", what)
580
581  image_props = build_image.ImagePropFromGlobalDict(info_dict, what)
582  image_props["timestamp"] = FIXED_FILE_TIMESTAMP
583
584  if what == "system":
585    fs_config_prefix = ""
586  else:
587    fs_config_prefix = what + "_"
588
589  fs_config = os.path.join(
590      input_dir, "META/" + fs_config_prefix + "filesystem_config.txt")
591  if not os.path.exists(fs_config):
592    fs_config = None
593
594  # Override values loaded from info_dict.
595  if fs_config:
596    image_props["fs_config"] = fs_config
597  if block_list:
598    image_props["block_list"] = block_list.name
599
600  build_image.BuildImage(
601      os.path.join(input_dir, what.upper()), image_props, output_file.name)
602
603  output_file.Write()
604  if block_list:
605    block_list.Write()
606
607  # Set the '_image_size' for given image size.
608  is_verity_partition = "verity_block_device" in image_props
609  verity_supported = (image_props.get("avb_enable") == "true")
610  is_avb_enable = image_props.get("avb_hashtree_enable") == "true"
611  if verity_supported and (is_verity_partition or is_avb_enable):
612    image_size = image_props.get("image_size")
613    if image_size:
614      image_size_key = what + "_image_size"
615      info_dict[image_size_key] = int(image_size)
616
617  use_dynamic_size = (
618      info_dict.get("use_dynamic_partition_size") == "true" and
619      what in shlex.split(info_dict.get("dynamic_partition_list", "").strip()))
620  if use_dynamic_size:
621    info_dict.update(build_image.GlobalDictFromImageProp(image_props, what))
622
623
624def AddUserdata(output_zip):
625  """Create a userdata image and store it in output_zip.
626
627  In most case we just create and store an empty userdata.img;
628  But the invoker can also request to create userdata.img with real
629  data from the target files, by setting "userdata_img_with_data=true"
630  in OPTIONS.info_dict.
631  """
632
633  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "userdata.img")
634  if os.path.exists(img.name):
635    logger.info("userdata.img already exists; no need to rebuild...")
636    return
637
638  # Skip userdata.img if no size.
639  image_props = build_image.ImagePropFromGlobalDict(OPTIONS.info_dict, "data")
640  if not image_props.get("partition_size"):
641    return
642
643  logger.info("creating userdata.img...")
644
645  image_props["timestamp"] = FIXED_FILE_TIMESTAMP
646
647  if OPTIONS.info_dict.get("userdata_img_with_data") == "true":
648    user_dir = os.path.join(OPTIONS.input_tmp, "DATA")
649  else:
650    user_dir = common.MakeTempDir()
651
652  build_image.BuildImage(user_dir, image_props, img.name)
653
654  common.CheckSize(img.name, "userdata.img", OPTIONS.info_dict)
655  # Always use compression for useradata image.
656  # As it's likely huge and consist of lots of 0s.
657  img.Write(zipfile.ZIP_DEFLATED)
658
659
660def AddVBMeta(output_zip, partitions, name, needed_partitions):
661  """Creates a VBMeta image and stores it in output_zip.
662
663  It generates the requested VBMeta image. The requested image could be for
664  top-level or chained VBMeta image, which is determined based on the name.
665
666  Args:
667    output_zip: The output zip file, which needs to be already open.
668    partitions: A dict that's keyed by partition names with image paths as
669        values. Only valid partition names are accepted, as partitions listed
670        in common.AVB_PARTITIONS and custom partitions listed in
671        OPTIONS.info_dict.get("avb_custom_images_partition_list")
672    name: Name of the VBMeta partition, e.g. 'vbmeta', 'vbmeta_system'.
673    needed_partitions: Partitions whose descriptors should be included into the
674        generated VBMeta image.
675
676  Returns:
677    Path to the created image.
678
679  Raises:
680    AssertionError: On invalid input args.
681  """
682  assert needed_partitions, "Needed partitions must be specified"
683
684  img = OutputFile(
685      output_zip, OPTIONS.input_tmp, "IMAGES", "{}.img".format(name))
686  if os.path.exists(img.name):
687    logger.info("%s.img already exists; not rebuilding...", name)
688    return img.name
689
690  common.BuildVBMeta(img.name, partitions, name, needed_partitions,
691                     OPTIONS.avb_resolve_rollback_index_location_conflict)
692  img.Write()
693  return img.name
694
695
696def AddCache(output_zip):
697  """Create an empty cache image and store it in output_zip."""
698
699  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "cache.img")
700  if os.path.exists(img.name):
701    logger.info("cache.img already exists; no need to rebuild...")
702    return
703
704  image_props = build_image.ImagePropFromGlobalDict(OPTIONS.info_dict, "cache")
705  # The build system has to explicitly request for cache.img.
706  if "fs_type" not in image_props:
707    return
708
709  logger.info("creating cache.img...")
710
711  image_props["timestamp"] = FIXED_FILE_TIMESTAMP
712
713  user_dir = common.MakeTempDir()
714  build_image.BuildImage(user_dir, image_props, img.name)
715
716  common.CheckSize(img.name, "cache.img", OPTIONS.info_dict)
717  img.Write()
718
719
720def CheckAbOtaImages(output_zip, ab_partitions):
721  """Checks that all the listed A/B partitions have their images available.
722
723  The images need to be available under IMAGES/ or RADIO/, with the former takes
724  a priority.
725
726  Args:
727    output_zip: The output zip file (needs to be already open), or None to
728        find images in OPTIONS.input_tmp/.
729    ab_partitions: The list of A/B partitions.
730
731  Raises:
732    AssertionError: If it can't find an image.
733  """
734  for partition in ab_partitions:
735    img_name = partition + ".img"
736
737    # Assert that the image is present under IMAGES/ now.
738    if output_zip:
739      # Zip spec says: All slashes MUST be forward slashes.
740      images_path = "IMAGES/" + img_name
741      radio_path = "RADIO/" + img_name
742      available = (images_path in output_zip.namelist() or
743                   radio_path in output_zip.namelist())
744    else:
745      images_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
746      radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
747      available = os.path.exists(images_path) or os.path.exists(radio_path)
748
749    assert available, "Failed to find " + img_name
750
751
752def AddPackRadioImages(output_zip, images):
753  """Copies images listed in META/pack_radioimages.txt from RADIO/ to IMAGES/.
754
755  Args:
756    output_zip: The output zip file (needs to be already open), or None to
757        write images to OPTIONS.input_tmp/.
758    images: A list of image names.
759
760  Raises:
761    AssertionError: If a listed image can't be found.
762  """
763  for image in images:
764    img_name = image.strip()
765    _, ext = os.path.splitext(img_name)
766    if not ext:
767      img_name += ".img"
768
769    prebuilt_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
770    if os.path.exists(prebuilt_path):
771      logger.info("%s already exists, no need to overwrite...", img_name)
772      continue
773
774    img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
775    assert os.path.exists(img_radio_path), \
776        "Failed to find %s at %s" % (img_name, img_radio_path)
777
778    if output_zip:
779      common.ZipWrite(output_zip, img_radio_path, "IMAGES/" + img_name)
780    else:
781      shutil.copy(img_radio_path, prebuilt_path)
782
783
784def AddSuperEmpty(output_zip):
785  """Create a super_empty.img and store it in output_zip."""
786
787  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "super_empty.img")
788  if os.path.exists(img.name):
789    logger.info("super_empty.img already exists; no need to rebuild...")
790    return
791  build_super_image.BuildSuperImage(OPTIONS.info_dict, img.name)
792  img.Write()
793
794
795def AddSuperSplit(output_zip):
796  """Create split super_*.img and store it in output_zip."""
797
798  outdir = os.path.join(OPTIONS.input_tmp, "OTA")
799  built = build_super_image.BuildSuperImage(OPTIONS.input_tmp, outdir)
800
801  if built:
802    for dev in OPTIONS.info_dict['super_block_devices'].strip().split():
803      img = OutputFile(output_zip, OPTIONS.input_tmp, "OTA",
804                       "super_" + dev + ".img")
805      img.Write()
806
807
808def ReplaceUpdatedFiles(zip_filename, files_list):
809  """Updates all the ZIP entries listed in files_list.
810
811  For now the list includes META/care_map.pb, and the related files under
812  SYSTEM/ after rebuilding recovery.
813  """
814  common.ZipDelete(zip_filename, files_list)
815  output_zip = zipfile.ZipFile(zip_filename, "a",
816                               compression=zipfile.ZIP_DEFLATED,
817                               allowZip64=True)
818  for item in files_list:
819    file_path = os.path.join(OPTIONS.input_tmp, item)
820    assert os.path.exists(file_path)
821    common.ZipWrite(output_zip, file_path, arcname=item)
822  common.ZipClose(output_zip)
823
824
825def HasPartition(partition_name):
826  """Determines if the target files archive should build a given partition."""
827
828  return ((os.path.isdir(
829      os.path.join(OPTIONS.input_tmp, partition_name.upper())) and
830      OPTIONS.info_dict.get(
831      "building_{}_image".format(partition_name)) == "true") or
832      os.path.exists(
833      os.path.join(OPTIONS.input_tmp, "IMAGES",
834                   "{}.img".format(partition_name))))
835
836
837def AddApexInfo(output_zip):
838  apex_infos = GetApexInfoFromTargetFiles(OPTIONS.input_tmp)
839  apex_metadata_proto = ota_metadata_pb2.ApexMetadata()
840  apex_metadata_proto.apex_info.extend(apex_infos)
841  apex_info_bytes = apex_metadata_proto.SerializeToString()
842
843  output_file = os.path.join(OPTIONS.input_tmp, "META", "apex_info.pb")
844  with open(output_file, "wb") as ofile:
845    ofile.write(apex_info_bytes)
846  if output_zip:
847    arc_name = "META/apex_info.pb"
848    if arc_name in output_zip.namelist():
849      OPTIONS.replace_updated_files_list.append(arc_name)
850    else:
851      common.ZipWrite(output_zip, output_file, arc_name)
852
853
854def AddVbmetaDigest(output_zip):
855  """Write the vbmeta digest to the output dir and zipfile."""
856
857  # Calculate the vbmeta digest and put the result in to META/
858  boot_images = OPTIONS.info_dict.get("boot_images")
859  # Disable the digest calculation if the target_file is used as a container
860  # for boot images. A boot container might contain boot-5.4.img, boot-5.10.img
861  # etc., instead of just a boot.img and will fail in vbmeta digest calculation.
862  boot_container = boot_images and (
863      len(boot_images.split()) >= 2 or boot_images.split()[0] != 'boot.img')
864  if (OPTIONS.info_dict.get("avb_enable") == "true" and not boot_container and
865          OPTIONS.info_dict.get("avb_building_vbmeta_image") == "true"):
866    avbtool = OPTIONS.info_dict["avb_avbtool"]
867    digest = verity_utils.CalculateVbmetaDigest(OPTIONS.input_tmp, avbtool)
868    vbmeta_digest_txt = os.path.join(OPTIONS.input_tmp, "META",
869                                     "vbmeta_digest.txt")
870    with open(vbmeta_digest_txt, 'w') as f:
871      f.write(digest)
872    # writes to the output zipfile
873    if output_zip:
874      arc_name = "META/vbmeta_digest.txt"
875      if arc_name in output_zip.namelist():
876        OPTIONS.replace_updated_files_list.append(arc_name)
877      else:
878        common.ZipWriteStr(output_zip, arc_name, digest)
879
880
881def AddImagesToTargetFiles(filename):
882  """Creates and adds images (boot/recovery/system/...) to a target_files.zip.
883
884  It works with either a zip file (zip mode), or a directory that contains the
885  files to be packed into a target_files.zip (dir mode). The latter is used when
886  being called from build/make/core/Makefile.
887
888  The images will be created under IMAGES/ in the input target_files.zip.
889
890  Args:
891    filename: the target_files.zip, or the zip root directory.
892  """
893  if os.path.isdir(filename):
894    OPTIONS.input_tmp = os.path.abspath(filename)
895  else:
896    OPTIONS.input_tmp = common.UnzipTemp(filename)
897
898  if not OPTIONS.add_missing:
899    if os.path.isdir(os.path.join(OPTIONS.input_tmp, "IMAGES")):
900      logger.warning("target_files appears to already contain images.")
901      sys.exit(1)
902
903  OPTIONS.info_dict = common.LoadInfoDict(OPTIONS.input_tmp, repacking=True)
904
905  has_recovery = OPTIONS.info_dict.get("no_recovery") != "true"
906  has_boot = OPTIONS.info_dict.get("no_boot") != "true"
907  has_init_boot = OPTIONS.info_dict.get("init_boot") == "true"
908  has_vendor_boot = OPTIONS.info_dict.get("vendor_boot") == "true"
909  has_vendor_kernel_boot = OPTIONS.info_dict.get(
910      "vendor_kernel_boot") == "true"
911
912  # {vendor,odm,product,system_ext,vendor_dlkm,odm_dlkm, system_dlkm, system, system_other}.img
913  # can be built from source, or  dropped into target_files.zip as a prebuilt blob.
914  has_vendor = HasPartition("vendor")
915  has_odm = HasPartition("odm")
916  has_vendor_dlkm = HasPartition("vendor_dlkm")
917  has_odm_dlkm = HasPartition("odm_dlkm")
918  has_system_dlkm = HasPartition("system_dlkm")
919  has_product = HasPartition("product")
920  has_system_ext = HasPartition("system_ext")
921  has_system = HasPartition("system")
922  has_system_other = HasPartition("system_other")
923  has_userdata = OPTIONS.info_dict.get("building_userdata_image") == "true"
924  has_cache = OPTIONS.info_dict.get("building_cache_image") == "true"
925
926  # Set up the output destination. It writes to the given directory for dir
927  # mode; otherwise appends to the given ZIP.
928  if os.path.isdir(filename):
929    output_zip = None
930  else:
931    output_zip = zipfile.ZipFile(filename, "a",
932                                 compression=zipfile.ZIP_DEFLATED,
933                                 allowZip64=True)
934
935  # Always make input_tmp/IMAGES available, since we may stage boot / recovery
936  # images there even under zip mode. The directory will be cleaned up as part
937  # of OPTIONS.input_tmp.
938  images_dir = os.path.join(OPTIONS.input_tmp, "IMAGES")
939  if not os.path.isdir(images_dir):
940    os.makedirs(images_dir)
941
942  # A map between partition names and their paths, which could be used when
943  # generating AVB vbmeta image.
944  partitions = {}
945
946  def banner(s):
947    logger.info("\n\n++++ %s  ++++\n\n", s)
948
949  boot_image = None
950  if has_boot:
951    banner("boot")
952    boot_images = OPTIONS.info_dict.get("boot_images")
953    if boot_images is None:
954      boot_images = "boot.img"
955    for index, b in enumerate(boot_images.split()):
956      # common.GetBootableImage() returns the image directly if present.
957      boot_image = common.GetBootableImage(
958          "IMAGES/" + b, b, OPTIONS.input_tmp, "BOOT")
959      # boot.img may be unavailable in some targets (e.g. aosp_arm64).
960      if boot_image:
961        boot_image_path = os.path.join(OPTIONS.input_tmp, "IMAGES", b)
962        # Although multiple boot images can be generated, include the image
963        # descriptor of only the first boot image in vbmeta
964        if index == 0:
965          partitions['boot'] = boot_image_path
966        if not os.path.exists(boot_image_path):
967          boot_image.WriteToDir(OPTIONS.input_tmp)
968          if output_zip:
969            boot_image.AddToZip(output_zip)
970
971  if has_init_boot:
972    banner("init_boot")
973    init_boot_image = common.GetBootableImage(
974        "IMAGES/init_boot.img", "init_boot.img", OPTIONS.input_tmp, "INIT_BOOT",
975        dev_nodes=True)
976    if init_boot_image:
977      partitions['init_boot'] = os.path.join(
978          OPTIONS.input_tmp, "IMAGES", "init_boot.img")
979      if not os.path.exists(partitions['init_boot']):
980        init_boot_image.WriteToDir(OPTIONS.input_tmp)
981        if output_zip:
982          init_boot_image.AddToZip(output_zip)
983
984  if has_vendor_boot:
985    banner("vendor_boot")
986    vendor_boot_image = common.GetVendorBootImage(
987        "IMAGES/vendor_boot.img", "vendor_boot.img", OPTIONS.input_tmp,
988        "VENDOR_BOOT")
989    if vendor_boot_image:
990      partitions['vendor_boot'] = os.path.join(OPTIONS.input_tmp, "IMAGES",
991                                               "vendor_boot.img")
992      if not os.path.exists(partitions['vendor_boot']):
993        vendor_boot_image.WriteToDir(OPTIONS.input_tmp)
994        if output_zip:
995          vendor_boot_image.AddToZip(output_zip)
996
997  if has_vendor_kernel_boot:
998    banner("vendor_kernel_boot")
999    vendor_kernel_boot_image = common.GetVendorKernelBootImage(
1000        "IMAGES/vendor_kernel_boot.img", "vendor_kernel_boot.img", OPTIONS.input_tmp,
1001        "VENDOR_KERNEL_BOOT")
1002    if vendor_kernel_boot_image:
1003      partitions['vendor_kernel_boot'] = os.path.join(OPTIONS.input_tmp, "IMAGES",
1004                                                      "vendor_kernel_boot.img")
1005      if not os.path.exists(partitions['vendor_kernel_boot']):
1006        vendor_kernel_boot_image.WriteToDir(OPTIONS.input_tmp)
1007        if output_zip:
1008          vendor_kernel_boot_image.AddToZip(output_zip)
1009
1010  recovery_image = None
1011  if has_recovery:
1012    banner("recovery")
1013    recovery_image = common.GetBootableImage(
1014        "IMAGES/recovery.img", "recovery.img", OPTIONS.input_tmp, "RECOVERY")
1015    assert recovery_image, "Failed to create recovery.img."
1016    partitions['recovery'] = os.path.join(
1017        OPTIONS.input_tmp, "IMAGES", "recovery.img")
1018    if not os.path.exists(partitions['recovery']):
1019      recovery_image.WriteToDir(OPTIONS.input_tmp)
1020      if output_zip:
1021        recovery_image.AddToZip(output_zip)
1022
1023      banner("recovery (two-step image)")
1024      # The special recovery.img for two-step package use.
1025      recovery_two_step_image = common.GetBootableImage(
1026          "OTA/recovery-two-step.img", "recovery-two-step.img",
1027          OPTIONS.input_tmp, "RECOVERY", two_step_image=True)
1028      assert recovery_two_step_image, "Failed to create recovery-two-step.img."
1029      recovery_two_step_image_path = os.path.join(
1030          OPTIONS.input_tmp, "OTA", "recovery-two-step.img")
1031      if not os.path.exists(recovery_two_step_image_path):
1032        recovery_two_step_image.WriteToDir(OPTIONS.input_tmp)
1033        if output_zip:
1034          recovery_two_step_image.AddToZip(output_zip)
1035
1036  def add_partition(partition, has_partition, add_func, add_args):
1037    if has_partition:
1038      banner(partition)
1039      partitions[partition] = add_func(output_zip, *add_args)
1040
1041  add_partition_calls = (
1042      ("system", has_system, AddSystem, [recovery_image, boot_image]),
1043      ("vendor", has_vendor, AddVendor, [recovery_image, boot_image]),
1044      ("product", has_product, AddProduct, []),
1045      ("system_ext", has_system_ext, AddSystemExt, []),
1046      ("odm", has_odm, AddOdm, []),
1047      ("vendor_dlkm", has_vendor_dlkm, AddVendorDlkm, []),
1048      ("odm_dlkm", has_odm_dlkm, AddOdmDlkm, []),
1049      ("system_dlkm", has_system_dlkm, AddSystemDlkm, []),
1050      ("system_other", has_system_other, AddSystemOther, []),
1051  )
1052  # If output_zip exists, each add_partition_calls writes bytes to the same output_zip,
1053  # which is not thread-safe. So, run them in serial if output_zip exists.
1054  if output_zip:
1055    for call in add_partition_calls:
1056      add_partition(*call)
1057  else:
1058    with ThreadPoolExecutor(max_workers=len(add_partition_calls)) as executor:
1059      for future in [executor.submit(add_partition, *call) for call in add_partition_calls]:
1060        future.result()
1061
1062  AddApexInfo(output_zip)
1063
1064  if not OPTIONS.is_signing:
1065    banner("userdata")
1066    AddUserdata(output_zip)
1067    banner("cache")
1068    AddCache(output_zip)
1069
1070  add_partition("dtbo",
1071                OPTIONS.info_dict.get("has_dtbo") == "true", AddDtbo, [])
1072  add_partition("pvmfw",
1073                OPTIONS.info_dict.get("has_pvmfw") == "true", AddPvmfw, [])
1074
1075  # Custom images.
1076  custom_partitions = OPTIONS.info_dict.get(
1077      "custom_images_partition_list", "").strip().split()
1078  for partition_name in custom_partitions:
1079    partition_name = partition_name.strip()
1080    banner("custom images for " + partition_name)
1081    image_list = OPTIONS.info_dict.get(
1082          "{}_image_list".format(partition_name)).split()
1083    partitions[partition_name] = AddCustomImages(output_zip, partition_name, image_list)
1084
1085  avb_custom_partitions = OPTIONS.info_dict.get(
1086      "avb_custom_images_partition_list", "").strip().split()
1087  for partition_name in avb_custom_partitions:
1088    partition_name = partition_name.strip()
1089    banner("avb custom images for " + partition_name)
1090    image_list = OPTIONS.info_dict.get(
1091          "avb_{}_image_list".format(partition_name)).split()
1092    partitions[partition_name] = AddCustomImages(output_zip, partition_name, image_list)
1093
1094  if OPTIONS.info_dict.get("avb_enable") == "true":
1095    # vbmeta_partitions includes the partitions that should be included into
1096    # top-level vbmeta.img, which are the ones that are not included in any
1097    # chained VBMeta image plus the chained VBMeta images themselves.
1098    # Currently avb_custom_partitions are all chained to VBMeta image.
1099    vbmeta_partitions = common.AVB_PARTITIONS[:] + tuple(avb_custom_partitions)
1100
1101    vbmeta_system = OPTIONS.info_dict.get("avb_vbmeta_system", "").strip()
1102    if vbmeta_system:
1103      banner("vbmeta_system")
1104      partitions["vbmeta_system"] = AddVBMeta(
1105          output_zip, partitions, "vbmeta_system", vbmeta_system.split())
1106      vbmeta_partitions = [
1107          item for item in vbmeta_partitions
1108          if item not in vbmeta_system.split()]
1109      vbmeta_partitions.append("vbmeta_system")
1110
1111    vbmeta_vendor = OPTIONS.info_dict.get("avb_vbmeta_vendor", "").strip()
1112    if vbmeta_vendor:
1113      banner("vbmeta_vendor")
1114      partitions["vbmeta_vendor"] = AddVBMeta(
1115          output_zip, partitions, "vbmeta_vendor", vbmeta_vendor.split())
1116      vbmeta_partitions = [
1117          item for item in vbmeta_partitions
1118          if item not in vbmeta_vendor.split()]
1119      vbmeta_partitions.append("vbmeta_vendor")
1120    custom_avb_partitions = OPTIONS.info_dict.get(
1121        "avb_custom_vbmeta_images_partition_list", "").strip().split()
1122    if custom_avb_partitions:
1123      for avb_part in custom_avb_partitions:
1124        partition_name = "vbmeta_" + avb_part
1125        included_partitions = OPTIONS.info_dict.get(
1126            "avb_vbmeta_{}".format(avb_part), "").strip().split()
1127        assert included_partitions, "Custom vbmeta partition {0} missing avb_vbmeta_{0} prop".format(
1128            avb_part)
1129        banner(partition_name)
1130        logger.info("VBMeta partition {} needs {}".format(
1131            partition_name, included_partitions))
1132        partitions[partition_name] = AddVBMeta(
1133            output_zip, partitions, partition_name, included_partitions)
1134        vbmeta_partitions = [
1135            item for item in vbmeta_partitions
1136            if item not in included_partitions]
1137        vbmeta_partitions.append(partition_name)
1138
1139    if OPTIONS.info_dict.get("avb_building_vbmeta_image") == "true":
1140      banner("vbmeta")
1141      AddVBMeta(output_zip, partitions, "vbmeta", vbmeta_partitions)
1142
1143  if OPTIONS.info_dict.get("use_dynamic_partitions") == "true":
1144    if OPTIONS.info_dict.get("build_super_empty_partition") == "true":
1145      banner("super_empty")
1146      AddSuperEmpty(output_zip)
1147
1148  if OPTIONS.info_dict.get("build_super_partition") == "true":
1149    if OPTIONS.info_dict.get(
1150            "build_retrofit_dynamic_partitions_ota_package") == "true":
1151      banner("super split images")
1152      AddSuperSplit(output_zip)
1153
1154  banner("radio")
1155  ab_partitions_txt = os.path.join(OPTIONS.input_tmp, "META",
1156                                   "ab_partitions.txt")
1157  if os.path.exists(ab_partitions_txt):
1158    with open(ab_partitions_txt) as f:
1159      ab_partitions = f.read().splitlines()
1160
1161    # For devices using A/B update, make sure we have all the needed images
1162    # ready under IMAGES/ or RADIO/.
1163    CheckAbOtaImages(output_zip, ab_partitions)
1164
1165    # Generate care_map.pb for ab_partitions, then write this file to
1166    # target_files package.
1167    output_care_map = os.path.join(OPTIONS.input_tmp, "META", "care_map.pb")
1168    AddCareMapForAbOta(output_zip if output_zip else output_care_map,
1169                       ab_partitions, partitions)
1170
1171  # Radio images that need to be packed into IMAGES/, and product-img.zip.
1172  pack_radioimages_txt = os.path.join(
1173      OPTIONS.input_tmp, "META", "pack_radioimages.txt")
1174  if os.path.exists(pack_radioimages_txt):
1175    with open(pack_radioimages_txt) as f:
1176      AddPackRadioImages(output_zip, f.readlines())
1177
1178  AddVbmetaDigest(output_zip)
1179
1180  if output_zip:
1181    common.ZipClose(output_zip)
1182    if OPTIONS.replace_updated_files_list:
1183      ReplaceUpdatedFiles(output_zip.filename,
1184                          OPTIONS.replace_updated_files_list)
1185
1186
1187def OptimizeCompressedEntries(zipfile_path):
1188  """Convert files that do not compress well to uncompressed storage
1189
1190  EROFS images tend to be compressed already, so compressing them again
1191  yields little space savings. Leaving them uncompressed will make
1192  downstream tooling's job easier, and save compute time.
1193  """
1194  if not zipfile.is_zipfile(zipfile_path):
1195    return
1196  entries_to_store = []
1197  with tempfile.TemporaryDirectory() as tmpdir:
1198    with zipfile.ZipFile(zipfile_path, "r", allowZip64=True) as zfp:
1199      for zinfo in zfp.filelist:
1200        if not zinfo.filename.startswith("IMAGES/") and not zinfo.filename.startswith("META"):
1201          continue
1202        # Don't try to store userdata.img uncompressed, it's usually huge.
1203        if zinfo.filename.endswith("userdata.img"):
1204          continue
1205        if zinfo.compress_size > zinfo.file_size * 0.80 and zinfo.compress_type != zipfile.ZIP_STORED:
1206          entries_to_store.append(zinfo)
1207          zfp.extract(zinfo, tmpdir)
1208    if len(entries_to_store) == 0:
1209      return
1210    # Remove these entries, then re-add them as ZIP_STORED
1211    ZipDelete(zipfile_path, [entry.filename for entry in entries_to_store])
1212    with zipfile.ZipFile(zipfile_path, "a", allowZip64=True) as zfp:
1213      for entry in entries_to_store:
1214        zfp.write(os.path.join(tmpdir, entry.filename),
1215                  entry.filename, compress_type=zipfile.ZIP_STORED)
1216
1217
1218def main(argv):
1219  def option_handler(o, a):
1220    if o in ("-a", "--add_missing"):
1221      OPTIONS.add_missing = True
1222    elif o in ("-r", "--rebuild_recovery",):
1223      OPTIONS.rebuild_recovery = True
1224    elif o == "--replace_verity_private_key":
1225      raise ValueError("--replace_verity_private_key is no longer supported,"
1226                       " please switch to AVB")
1227    elif o == "--replace_verity_public_key":
1228      raise ValueError("--replace_verity_public_key is no longer supported,"
1229                       " please switch to AVB")
1230    elif o == "--is_signing":
1231      OPTIONS.is_signing = True
1232    elif o == "--avb_resolve_rollback_index_location_conflict":
1233      OPTIONS.avb_resolve_rollback_index_location_conflict = True
1234    else:
1235      return False
1236    return True
1237
1238  args = common.ParseOptions(
1239      argv, __doc__, extra_opts="ar",
1240      extra_long_opts=["add_missing", "rebuild_recovery",
1241                       "replace_verity_public_key=",
1242                       "replace_verity_private_key=",
1243                       "is_signing",
1244                       "avb_resolve_rollback_index_location_conflict"],
1245      extra_option_handler=option_handler)
1246
1247  if len(args) != 1:
1248    common.Usage(__doc__)
1249    sys.exit(1)
1250
1251  common.InitLogging()
1252
1253  AddImagesToTargetFiles(args[0])
1254  OptimizeCompressedEntries(args[0])
1255  logger.info("done.")
1256
1257
1258if __name__ == '__main__':
1259  try:
1260    common.CloseInheritedPipes()
1261    main(sys.argv[1:])
1262  finally:
1263    common.Cleanup()
1264