add_img_to_target_files.py revision ae396d9b0cbc5b49aea4c2d0d8146b77160354d4 
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.
32
33  --replace_verity_private_key
34      Replace the private key used for verity signing. (same as the option
35      in sign_target_files_apks)
36
37  --replace_verity_public_key
38       Replace the certificate (public key) used for verity verification. (same
39       as the option in sign_target_files_apks)
40
41  --is_signing
42      Skip building & adding the images for "userdata" and "cache" if we
43      are signing the target files.
44"""
45
46from __future__ import print_function
47
48import sys
49
50if sys.hexversion < 0x02070000:
51  print("Python 2.7 or newer is required.", file=sys.stderr)
52  sys.exit(1)
53
54import datetime
55import hashlib
56import os
57import shlex
58import shutil
59import subprocess
60import tempfile
61import uuid
62import zipfile
63
64import build_image
65import common
66import rangelib
67import sparse_img
68
69OPTIONS = common.OPTIONS
70
71OPTIONS.add_missing = False
72OPTIONS.rebuild_recovery = False
73OPTIONS.replace_updated_files_list = []
74OPTIONS.replace_verity_public_key = False
75OPTIONS.replace_verity_private_key = False
76OPTIONS.is_signing = False
77
78
79class OutputFile(object):
80  def __init__(self, output_zip, input_dir, prefix, name):
81    self._output_zip = output_zip
82    self.input_name = os.path.join(input_dir, prefix, name)
83
84    if self._output_zip:
85      self._zip_name = os.path.join(prefix, name)
86
87      root, suffix = os.path.splitext(name)
88      self.name = common.MakeTempFile(prefix=root + '-', suffix=suffix)
89    else:
90      self.name = self.input_name
91
92  def Write(self):
93    if self._output_zip:
94      common.ZipWrite(self._output_zip, self.name, self._zip_name)
95
96
97def GetCareMap(which, imgname):
98  """Generate care_map of system (or vendor) partition"""
99
100  assert which in ("system", "vendor")
101
102  simg = sparse_img.SparseImage(imgname)
103  care_map_list = [which]
104
105  care_map_ranges = simg.care_map
106  key = which + "_adjusted_partition_size"
107  adjusted_blocks = OPTIONS.info_dict.get(key)
108  if adjusted_blocks:
109    assert adjusted_blocks > 0, "blocks should be positive for " + which
110    care_map_ranges = care_map_ranges.intersect(rangelib.RangeSet(
111        "0-%d" % (adjusted_blocks,)))
112
113  care_map_list.append(care_map_ranges.to_string_raw())
114  return care_map_list
115
116
117def AddSystem(output_zip, prefix="IMAGES/", recovery_img=None, boot_img=None):
118  """Turn the contents of SYSTEM into a system image and store it in
119  output_zip. Returns the name of the system image file."""
120
121  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "system.img")
122  if os.path.exists(img.input_name):
123    print("system.img already exists in %s, no need to rebuild..." % (prefix,))
124    return img.input_name
125
126  def output_sink(fn, data):
127    ofile = open(os.path.join(OPTIONS.input_tmp, "SYSTEM", fn), "w")
128    ofile.write(data)
129    ofile.close()
130
131    arc_name = "SYSTEM/" + fn
132    if arc_name in output_zip.namelist():
133      OPTIONS.replace_updated_files_list.append(arc_name)
134    else:
135      common.ZipWrite(output_zip, ofile.name, arc_name)
136
137  if OPTIONS.rebuild_recovery:
138    print("Building new recovery patch")
139    common.MakeRecoveryPatch(OPTIONS.input_tmp, output_sink, recovery_img,
140                             boot_img, info_dict=OPTIONS.info_dict)
141
142  block_list = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "system.map")
143  CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "system", img,
144              block_list=block_list)
145
146  return img.name
147
148
149def AddSystemOther(output_zip, prefix="IMAGES/"):
150  """Turn the contents of SYSTEM_OTHER into a system_other image
151  and store it in output_zip."""
152
153  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "system_other.img")
154  if os.path.exists(img.input_name):
155    print("system_other.img already exists in %s, no need to rebuild..." % (
156        prefix,))
157    return
158
159  CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "system_other", img)
160
161
162def AddVendor(output_zip, prefix="IMAGES/"):
163  """Turn the contents of VENDOR into a vendor image and store in it
164  output_zip."""
165
166  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "vendor.img")
167  if os.path.exists(img.input_name):
168    print("vendor.img already exists in %s, no need to rebuild..." % (prefix,))
169    return img.input_name
170
171  block_list = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "vendor.map")
172  CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "vendor", img,
173              block_list=block_list)
174  return img.name
175
176
177def AddDtbo(output_zip, prefix="IMAGES/"):
178  """Adds the DTBO image.
179
180  Uses the image under prefix if it already exists. Otherwise looks for the
181  image under PREBUILT_IMAGES/, signs it as needed, and returns the image name.
182  """
183
184  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "dtbo.img")
185  if os.path.exists(img.input_name):
186    print("dtbo.img already exists in %s, no need to rebuild..." % (prefix,))
187    return img.input_name
188
189  dtbo_prebuilt_path = os.path.join(
190      OPTIONS.input_tmp, "PREBUILT_IMAGES", "dtbo.img")
191  assert os.path.exists(dtbo_prebuilt_path)
192  shutil.copy(dtbo_prebuilt_path, img.name)
193
194  # AVB-sign the image as needed.
195  if OPTIONS.info_dict.get("avb_enable") == "true":
196    avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
197    part_size = OPTIONS.info_dict["dtbo_size"]
198    # The AVB hash footer will be replaced if already present.
199    cmd = [avbtool, "add_hash_footer", "--image", img.name,
200           "--partition_size", str(part_size), "--partition_name", "dtbo"]
201    common.AppendAVBSigningArgs(cmd, "dtbo")
202    args = OPTIONS.info_dict.get("avb_dtbo_add_hash_footer_args")
203    if args and args.strip():
204      cmd.extend(shlex.split(args))
205    p = common.Run(cmd, stdout=subprocess.PIPE)
206    p.communicate()
207    assert p.returncode == 0, \
208        "avbtool add_hash_footer of %s failed" % (img.name,)
209
210  img.Write()
211  return img.name
212
213
214def CreateImage(input_dir, info_dict, what, output_file, block_list=None):
215  print("creating " + what + ".img...")
216
217  image_props = build_image.ImagePropFromGlobalDict(info_dict, what)
218  fstab = info_dict["fstab"]
219  mount_point = "/" + what
220  if fstab and mount_point in fstab:
221    image_props["fs_type"] = fstab[mount_point].fs_type
222
223  # Use a fixed timestamp (01/01/2009) when packaging the image.
224  # Bug: 24377993
225  epoch = datetime.datetime.fromtimestamp(0)
226  timestamp = (datetime.datetime(2009, 1, 1) - epoch).total_seconds()
227  image_props["timestamp"] = int(timestamp)
228
229  if what == "system":
230    fs_config_prefix = ""
231  else:
232    fs_config_prefix = what + "_"
233
234  fs_config = os.path.join(
235      input_dir, "META/" + fs_config_prefix + "filesystem_config.txt")
236  if not os.path.exists(fs_config):
237    fs_config = None
238
239  # Override values loaded from info_dict.
240  if fs_config:
241    image_props["fs_config"] = fs_config
242  if block_list:
243    image_props["block_list"] = block_list.name
244
245  # Use repeatable ext4 FS UUID and hash_seed UUID (based on partition name and
246  # build fingerprint).
247  uuid_seed = what + "-"
248  if "build.prop" in info_dict:
249    build_prop = info_dict["build.prop"]
250    if "ro.build.fingerprint" in build_prop:
251      uuid_seed += build_prop["ro.build.fingerprint"]
252    elif "ro.build.thumbprint" in build_prop:
253      uuid_seed += build_prop["ro.build.thumbprint"]
254  image_props["uuid"] = str(uuid.uuid5(uuid.NAMESPACE_URL, uuid_seed))
255  hash_seed = "hash_seed-" + uuid_seed
256  image_props["hash_seed"] = str(uuid.uuid5(uuid.NAMESPACE_URL, hash_seed))
257
258  succ = build_image.BuildImage(os.path.join(input_dir, what.upper()),
259                                image_props, output_file.name)
260  assert succ, "build " + what + ".img image failed"
261
262  output_file.Write()
263  if block_list:
264    block_list.Write()
265
266  # Set the 'adjusted_partition_size' that excludes the verity blocks of the
267  # given image. When avb is enabled, this size is the max image size returned
268  # by the avb tool.
269  is_verity_partition = "verity_block_device" in image_props
270  verity_supported = (image_props.get("verity") == "true" or
271                      image_props.get("avb_enable") == "true")
272  is_avb_enable = image_props.get("avb_hashtree_enable") == "true"
273  if verity_supported and (is_verity_partition or is_avb_enable):
274    adjusted_blocks_value = image_props.get("partition_size")
275    if adjusted_blocks_value:
276      adjusted_blocks_key = what + "_adjusted_partition_size"
277      info_dict[adjusted_blocks_key] = int(adjusted_blocks_value)/4096 - 1
278
279
280def AddUserdata(output_zip, prefix="IMAGES/"):
281  """Create a userdata image and store it in output_zip.
282
283  In most case we just create and store an empty userdata.img;
284  But the invoker can also request to create userdata.img with real
285  data from the target files, by setting "userdata_img_with_data=true"
286  in OPTIONS.info_dict.
287  """
288
289  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "userdata.img")
290  if os.path.exists(img.input_name):
291    print("userdata.img already exists in %s, no need to rebuild..." % (
292        prefix,))
293    return
294
295  # Skip userdata.img if no size.
296  image_props = build_image.ImagePropFromGlobalDict(OPTIONS.info_dict, "data")
297  if not image_props.get("partition_size"):
298    return
299
300  print("creating userdata.img...")
301
302  # Use a fixed timestamp (01/01/2009) when packaging the image.
303  # Bug: 24377993
304  epoch = datetime.datetime.fromtimestamp(0)
305  timestamp = (datetime.datetime(2009, 1, 1) - epoch).total_seconds()
306  image_props["timestamp"] = int(timestamp)
307
308  if OPTIONS.info_dict.get("userdata_img_with_data") == "true":
309    user_dir = os.path.join(OPTIONS.input_tmp, "DATA")
310  else:
311    user_dir = tempfile.mkdtemp()
312    OPTIONS.tempfiles.append(user_dir)
313
314  fstab = OPTIONS.info_dict["fstab"]
315  if fstab:
316    image_props["fs_type"] = fstab["/data"].fs_type
317  succ = build_image.BuildImage(user_dir, image_props, img.name)
318  assert succ, "build userdata.img image failed"
319
320  common.CheckSize(img.name, "userdata.img", OPTIONS.info_dict)
321  img.Write()
322
323
324def AppendVBMetaArgsForPartition(cmd, partition, img_path, public_key_dir):
325  if not img_path:
326    return
327
328  # Check if chain partition is used.
329  key_path = OPTIONS.info_dict.get("avb_" + partition + "_key_path")
330  if key_path:
331    # extract public key in AVB format to be included in vbmeta.img
332    avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
333    public_key_path = os.path.join(public_key_dir, "%s.avbpubkey" % partition)
334    p = common.Run([avbtool, "extract_public_key", "--key", key_path,
335                    "--output", public_key_path],
336                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
337    p.communicate()
338    assert p.returncode == 0, \
339        "avbtool extract_public_key fail for partition: %r" % partition
340
341    rollback_index_location = OPTIONS.info_dict[
342        "avb_" + partition + "_rollback_index_location"]
343    cmd.extend(["--chain_partition", "%s:%s:%s" % (
344        partition, rollback_index_location, public_key_path)])
345  else:
346    cmd.extend(["--include_descriptors_from_image", img_path])
347
348
349def AddVBMeta(output_zip, boot_img_path, system_img_path, vendor_img_path,
350              dtbo_img_path, prefix="IMAGES/"):
351  """Create a VBMeta image and store it in output_zip."""
352  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "vbmeta.img")
353  if os.path.exists(img.input_name):
354    print("vbmeta.img already exists in %s; not rebuilding..." % (prefix,))
355    return img.input_name
356
357  avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
358  cmd = [avbtool, "make_vbmeta_image", "--output", img.name]
359  common.AppendAVBSigningArgs(cmd, "vbmeta")
360
361  public_key_dir = tempfile.mkdtemp(prefix="avbpubkey-")
362  OPTIONS.tempfiles.append(public_key_dir)
363
364  AppendVBMetaArgsForPartition(cmd, "boot", boot_img_path, public_key_dir)
365  AppendVBMetaArgsForPartition(cmd, "system", system_img_path, public_key_dir)
366  AppendVBMetaArgsForPartition(cmd, "vendor", vendor_img_path, public_key_dir)
367  AppendVBMetaArgsForPartition(cmd, "dtbo", dtbo_img_path, public_key_dir)
368
369  args = OPTIONS.info_dict.get("avb_vbmeta_args")
370  if args and args.strip():
371    split_args = shlex.split(args)
372    for index, arg in enumerate(split_args[:-1]):
373      # Sanity check that the image file exists. Some images might be defined
374      # as a path relative to source tree, which may not be available at the
375      # same location when running this script (we have the input target_files
376      # zip only). For such cases, we additionally scan other locations (e.g.
377      # IMAGES/, RADIO/, etc) before bailing out.
378      if arg == '--include_descriptors_from_image':
379        image_path = split_args[index + 1]
380        if os.path.exists(image_path):
381          continue
382        found = False
383        for dir in ['IMAGES', 'RADIO', 'VENDOR_IMAGES', 'PREBUILT_IMAGES']:
384          alt_path = os.path.join(
385              OPTIONS.input_tmp, dir, os.path.basename(image_path))
386          if os.path.exists(alt_path):
387            split_args[index + 1] = alt_path
388            found = True
389            break
390        assert found, 'failed to find %s' % (image_path,)
391    cmd.extend(split_args)
392
393  p = common.Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
394  p.communicate()
395  assert p.returncode == 0, "avbtool make_vbmeta_image failed"
396  img.Write()
397
398
399def AddPartitionTable(output_zip, prefix="IMAGES/"):
400  """Create a partition table image and store it in output_zip."""
401
402  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "partition-table.img")
403  bpt = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "partition-table.bpt")
404
405  # use BPTTOOL from environ, or "bpttool" if empty or not set.
406  bpttool = os.getenv("BPTTOOL") or "bpttool"
407  cmd = [bpttool, "make_table", "--output_json", bpt.name,
408         "--output_gpt", img.name]
409  input_files_str = OPTIONS.info_dict["board_bpt_input_files"]
410  input_files = input_files_str.split(" ")
411  for i in input_files:
412    cmd.extend(["--input", i])
413  disk_size = OPTIONS.info_dict.get("board_bpt_disk_size")
414  if disk_size:
415    cmd.extend(["--disk_size", disk_size])
416  args = OPTIONS.info_dict.get("board_bpt_make_table_args")
417  if args:
418    cmd.extend(shlex.split(args))
419
420  p = common.Run(cmd, stdout=subprocess.PIPE)
421  p.communicate()
422  assert p.returncode == 0, "bpttool make_table failed"
423
424  img.Write()
425  bpt.Write()
426
427
428def AddCache(output_zip, prefix="IMAGES/"):
429  """Create an empty cache image and store it in output_zip."""
430
431  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "cache.img")
432  if os.path.exists(img.input_name):
433    print("cache.img already exists in %s, no need to rebuild..." % (prefix,))
434    return
435
436  image_props = build_image.ImagePropFromGlobalDict(OPTIONS.info_dict, "cache")
437  # The build system has to explicitly request for cache.img.
438  if "fs_type" not in image_props:
439    return
440
441  print("creating cache.img...")
442
443  # Use a fixed timestamp (01/01/2009) when packaging the image.
444  # Bug: 24377993
445  epoch = datetime.datetime.fromtimestamp(0)
446  timestamp = (datetime.datetime(2009, 1, 1) - epoch).total_seconds()
447  image_props["timestamp"] = int(timestamp)
448
449  user_dir = tempfile.mkdtemp()
450  OPTIONS.tempfiles.append(user_dir)
451
452  fstab = OPTIONS.info_dict["fstab"]
453  if fstab:
454    image_props["fs_type"] = fstab["/cache"].fs_type
455  succ = build_image.BuildImage(user_dir, image_props, img.name)
456  assert succ, "build cache.img image failed"
457
458  common.CheckSize(img.name, "cache.img", OPTIONS.info_dict)
459  img.Write()
460
461
462def ReplaceUpdatedFiles(zip_filename, files_list):
463  """Update all the zip entries listed in the files_list.
464
465  For now the list includes META/care_map.txt, and the related files under
466  SYSTEM/ after rebuilding recovery.
467  """
468
469  cmd = ["zip", "-d", zip_filename] + files_list
470  p = common.Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
471  p.communicate()
472
473  output_zip = zipfile.ZipFile(zip_filename, "a",
474                               compression=zipfile.ZIP_DEFLATED,
475                               allowZip64=True)
476  for item in files_list:
477    file_path = os.path.join(OPTIONS.input_tmp, item)
478    assert os.path.exists(file_path)
479    common.ZipWrite(output_zip, file_path, arcname=item)
480  common.ZipClose(output_zip)
481
482
483def AddImagesToTargetFiles(filename):
484  """Creates and adds images (boot/recovery/system/...) to a target_files.zip.
485
486  It works with either a zip file (zip mode), or a directory that contains the
487  files to be packed into a target_files.zip (dir mode). The latter is used when
488  being called from build/make/core/Makefile.
489
490  The images will be created under IMAGES/ in the input target_files.zip.
491
492  Args:
493      filename: the target_files.zip, or the zip root directory.
494  """
495  if os.path.isdir(filename):
496    OPTIONS.input_tmp = os.path.abspath(filename)
497    input_zip = None
498  else:
499    OPTIONS.input_tmp, input_zip = common.UnzipTemp(filename)
500
501  if not OPTIONS.add_missing:
502    if os.path.isdir(os.path.join(OPTIONS.input_tmp, "IMAGES")):
503      print("target_files appears to already contain images.")
504      sys.exit(1)
505
506  # vendor.img is unlike system.img or system_other.img. Because it could be
507  # built from source, or dropped into target_files.zip as a prebuilt blob. We
508  # consider either of them as vendor.img being available, which could be used
509  # when generating vbmeta.img for AVB.
510  has_vendor = (os.path.isdir(os.path.join(OPTIONS.input_tmp, "VENDOR")) or
511                os.path.exists(os.path.join(OPTIONS.input_tmp, "IMAGES",
512                                            "vendor.img")))
513  has_system_other = os.path.isdir(os.path.join(OPTIONS.input_tmp,
514                                                "SYSTEM_OTHER"))
515
516  if input_zip:
517    OPTIONS.info_dict = common.LoadInfoDict(input_zip, OPTIONS.input_tmp)
518
519    common.ZipClose(input_zip)
520    output_zip = zipfile.ZipFile(filename, "a",
521                                 compression=zipfile.ZIP_DEFLATED,
522                                 allowZip64=True)
523  else:
524    OPTIONS.info_dict = common.LoadInfoDict(filename, filename)
525    output_zip = None
526
527  # Always make input_tmp/IMAGES available, since we may stage boot / recovery
528  # images there even under zip mode. The directory will be cleaned up as part
529  # of OPTIONS.input_tmp.
530  images_dir = os.path.join(OPTIONS.input_tmp, "IMAGES")
531  if not os.path.isdir(images_dir):
532    os.makedirs(images_dir)
533
534  has_recovery = (OPTIONS.info_dict.get("no_recovery") != "true")
535
536  if OPTIONS.info_dict.get("avb_enable") == "true":
537    fp = None
538    if "build.prop" in OPTIONS.info_dict:
539      build_prop = OPTIONS.info_dict["build.prop"]
540      if "ro.build.fingerprint" in build_prop:
541        fp = build_prop["ro.build.fingerprint"]
542      elif "ro.build.thumbprint" in build_prop:
543        fp = build_prop["ro.build.thumbprint"]
544    if fp:
545      OPTIONS.info_dict["avb_salt"] = hashlib.sha256(fp).hexdigest()
546
547  def banner(s):
548    print("\n\n++++ " + s + " ++++\n\n")
549
550  banner("boot")
551  # common.GetBootableImage() returns the image directly if present.
552  boot_image = common.GetBootableImage(
553      "IMAGES/boot.img", "boot.img", OPTIONS.input_tmp, "BOOT")
554  # boot.img may be unavailable in some targets (e.g. aosp_arm64).
555  if boot_image:
556    boot_img_path = os.path.join(OPTIONS.input_tmp, "IMAGES", "boot.img")
557    if not os.path.exists(boot_img_path):
558      boot_image.WriteToDir(OPTIONS.input_tmp)
559      if output_zip:
560        boot_image.AddToZip(output_zip)
561
562  recovery_image = None
563  if has_recovery:
564    banner("recovery")
565    recovery_image = common.GetBootableImage(
566        "IMAGES/recovery.img", "recovery.img", OPTIONS.input_tmp, "RECOVERY")
567    assert recovery_image, "Failed to create recovery.img."
568    recovery_img_path = os.path.join(
569        OPTIONS.input_tmp, "IMAGES", "recovery.img")
570    if not os.path.exists(recovery_img_path):
571      recovery_image.WriteToDir(OPTIONS.input_tmp)
572      if output_zip:
573        recovery_image.AddToZip(output_zip)
574
575      banner("recovery (two-step image)")
576      # The special recovery.img for two-step package use.
577      recovery_two_step_image = common.GetBootableImage(
578          "IMAGES/recovery-two-step.img", "recovery-two-step.img",
579          OPTIONS.input_tmp, "RECOVERY", two_step_image=True)
580      assert recovery_two_step_image, "Failed to create recovery-two-step.img."
581      recovery_two_step_image_path = os.path.join(
582          OPTIONS.input_tmp, "IMAGES", "recovery-two-step.img")
583      if not os.path.exists(recovery_two_step_image_path):
584        recovery_two_step_image.WriteToDir(OPTIONS.input_tmp)
585        if output_zip:
586          recovery_two_step_image.AddToZip(output_zip)
587
588  banner("system")
589  system_img_path = AddSystem(
590      output_zip, recovery_img=recovery_image, boot_img=boot_image)
591  vendor_img_path = None
592  if has_vendor:
593    banner("vendor")
594    vendor_img_path = AddVendor(output_zip)
595  if has_system_other:
596    banner("system_other")
597    AddSystemOther(output_zip)
598  if not OPTIONS.is_signing:
599    banner("userdata")
600    AddUserdata(output_zip)
601    banner("cache")
602    AddCache(output_zip)
603
604  if OPTIONS.info_dict.get("board_bpt_enable") == "true":
605    banner("partition-table")
606    AddPartitionTable(output_zip)
607
608  dtbo_img_path = None
609  if OPTIONS.info_dict.get("has_dtbo") == "true":
610    banner("dtbo")
611    dtbo_img_path = AddDtbo(output_zip)
612
613  if OPTIONS.info_dict.get("avb_enable") == "true":
614    banner("vbmeta")
615    AddVBMeta(output_zip, boot_img_path, system_img_path, vendor_img_path,
616              dtbo_img_path)
617
618  # For devices using A/B update, copy over images from RADIO/ and/or
619  # VENDOR_IMAGES/ to IMAGES/ and make sure we have all the needed
620  # images ready under IMAGES/. All images should have '.img' as extension.
621  banner("radio")
622  ab_partitions = os.path.join(OPTIONS.input_tmp, "META", "ab_partitions.txt")
623  if os.path.exists(ab_partitions):
624    with open(ab_partitions, 'r') as f:
625      lines = f.readlines()
626    # For devices using A/B update, generate care_map for system and vendor
627    # partitions (if present), then write this file to target_files package.
628    care_map_list = []
629    for line in lines:
630      if line.strip() == "system" and (
631          "system_verity_block_device" in OPTIONS.info_dict or
632          OPTIONS.info_dict.get("avb_system_hashtree_enable") == "true"):
633        assert os.path.exists(system_img_path)
634        care_map_list += GetCareMap("system", system_img_path)
635      if line.strip() == "vendor" and (
636          "vendor_verity_block_device" in OPTIONS.info_dict or
637          OPTIONS.info_dict.get("avb_vendor_hashtree_enable") == "true"):
638        assert os.path.exists(vendor_img_path)
639        care_map_list += GetCareMap("vendor", vendor_img_path)
640
641      img_name = line.strip() + ".img"
642      prebuilt_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
643      if os.path.exists(prebuilt_path):
644        print("%s already exists, no need to overwrite..." % (img_name,))
645        continue
646
647      img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
648      img_vendor_dir = os.path.join(
649        OPTIONS.input_tmp, "VENDOR_IMAGES")
650      if os.path.exists(img_radio_path):
651        if output_zip:
652          common.ZipWrite(output_zip, img_radio_path,
653                          os.path.join("IMAGES", img_name))
654        else:
655          shutil.copy(img_radio_path, prebuilt_path)
656      else:
657        for root, _, files in os.walk(img_vendor_dir):
658          if img_name in files:
659            if output_zip:
660              common.ZipWrite(output_zip, os.path.join(root, img_name),
661                os.path.join("IMAGES", img_name))
662            else:
663              shutil.copy(os.path.join(root, img_name), prebuilt_path)
664            break
665
666      if output_zip:
667        # Zip spec says: All slashes MUST be forward slashes.
668        img_path = 'IMAGES/' + img_name
669        assert img_path in output_zip.namelist(), "cannot find " + img_name
670      else:
671        img_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
672        assert os.path.exists(img_path), "cannot find " + img_name
673
674    if care_map_list:
675      care_map_path = "META/care_map.txt"
676      if output_zip and care_map_path not in output_zip.namelist():
677        common.ZipWriteStr(output_zip, care_map_path, '\n'.join(care_map_list))
678      else:
679        with open(os.path.join(OPTIONS.input_tmp, care_map_path), 'w') as fp:
680          fp.write('\n'.join(care_map_list))
681        if output_zip:
682          OPTIONS.replace_updated_files_list.append(care_map_path)
683
684  # Radio images that need to be packed into IMAGES/, and product-img.zip.
685  pack_radioimages = os.path.join(
686      OPTIONS.input_tmp, "META", "pack_radioimages.txt")
687  if os.path.exists(pack_radioimages):
688    with open(pack_radioimages, 'r') as f:
689      lines = f.readlines()
690    for line in lines:
691      img_name = line.strip()
692      _, ext = os.path.splitext(img_name)
693      if not ext:
694        img_name += ".img"
695      prebuilt_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
696      if os.path.exists(prebuilt_path):
697        print("%s already exists, no need to overwrite..." % (img_name,))
698        continue
699
700      img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
701      assert os.path.exists(img_radio_path), \
702          "Failed to find %s at %s" % (img_name, img_radio_path)
703      if output_zip:
704        common.ZipWrite(output_zip, img_radio_path,
705                        os.path.join("IMAGES", img_name))
706      else:
707        shutil.copy(img_radio_path, prebuilt_path)
708
709  if output_zip:
710    common.ZipClose(output_zip)
711    if OPTIONS.replace_updated_files_list:
712      ReplaceUpdatedFiles(output_zip.filename,
713                          OPTIONS.replace_updated_files_list)
714
715
716def main(argv):
717  def option_handler(o, a):
718    if o in ("-a", "--add_missing"):
719      OPTIONS.add_missing = True
720    elif o in ("-r", "--rebuild_recovery",):
721      OPTIONS.rebuild_recovery = True
722    elif o == "--replace_verity_private_key":
723      OPTIONS.replace_verity_private_key = (True, a)
724    elif o == "--replace_verity_public_key":
725      OPTIONS.replace_verity_public_key = (True, a)
726    elif o == "--is_signing":
727      OPTIONS.is_signing = True
728    else:
729      return False
730    return True
731
732  args = common.ParseOptions(
733      argv, __doc__, extra_opts="ar",
734      extra_long_opts=["add_missing", "rebuild_recovery",
735                       "replace_verity_public_key=",
736                       "replace_verity_private_key=",
737                       "is_signing"],
738      extra_option_handler=option_handler)
739
740
741  if len(args) != 1:
742    common.Usage(__doc__)
743    sys.exit(1)
744
745  AddImagesToTargetFiles(args[0])
746  print("done.")
747
748if __name__ == '__main__':
749  try:
750    common.CloseInheritedPipes()
751    main(sys.argv[1:])
752  except common.ExternalError as e:
753    print("\n   ERROR: %s\n" % (e,))
754    sys.exit(1)
755  finally:
756    common.Cleanup()
757