1#!/usr/bin/env python
2#
3# Copyright (C) 2008 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"""
18Signs all the APK files in a target-files zipfile, producing a new
19target-files zip.
20
21Usage:  sign_target_files_apks [flags] input_target_files output_target_files
22
23  -e  (--extra_apks)  <name,name,...=key>
24      Add extra APK/APEX name/key pairs as though they appeared in apkcerts.txt
25      or apexkeys.txt (so mappings specified by -k and -d are applied). Keys
26      specified in -e override any value for that app contained in the
27      apkcerts.txt file, or the container key for an APEX. Option may be
28      repeated to give multiple extra packages.
29
30  --extra_apex_payload_key <name=key>
31      Add a mapping for APEX package name to payload signing key, which will
32      override the default payload signing key in apexkeys.txt. Note that the
33      container key should be overridden via the `--extra_apks` flag above.
34      Option may be repeated for multiple APEXes.
35
36  --skip_apks_with_path_prefix  <prefix>
37      Skip signing an APK if it has the matching prefix in its path. The prefix
38      should be matching the entry name, which has partition names in upper
39      case, e.g. "VENDOR/app/", or "SYSTEM_OTHER/preloads/". Option may be
40      repeated to give multiple prefixes.
41
42  -k  (--key_mapping)  <src_key=dest_key>
43      Add a mapping from the key name as specified in apkcerts.txt (the
44      src_key) to the real key you wish to sign the package with
45      (dest_key).  Option may be repeated to give multiple key
46      mappings.
47
48  -d  (--default_key_mappings)  <dir>
49      Set up the following key mappings:
50
51        $devkey/devkey    ==>  $dir/releasekey
52        $devkey/testkey   ==>  $dir/releasekey
53        $devkey/media     ==>  $dir/media
54        $devkey/shared    ==>  $dir/shared
55        $devkey/platform  ==>  $dir/platform
56
57      where $devkey is the directory part of the value of
58      default_system_dev_certificate from the input target-files's
59      META/misc_info.txt.  (Defaulting to "build/make/target/product/security"
60      if the value is not present in misc_info.
61
62      -d and -k options are added to the set of mappings in the order
63      in which they appear on the command line.
64
65  -o  (--replace_ota_keys)
66      Replace the certificate (public key) used by OTA package verification
67      with the ones specified in the input target_files zip (in the
68      META/otakeys.txt file). Key remapping (-k and -d) is performed on the
69      keys. For A/B devices, the payload verification key will be replaced
70      as well. If there're multiple OTA keys, only the first one will be used
71      for payload verification.
72
73  -t  (--tag_changes)  <+tag>,<-tag>,...
74      Comma-separated list of changes to make to the set of tags (in
75      the last component of the build fingerprint).  Prefix each with
76      '+' or '-' to indicate whether that tag should be added or
77      removed.  Changes are processed in the order they appear.
78      Default value is "-test-keys,-dev-keys,+release-keys".
79
80  --replace_verity_private_key <key>
81      Replace the private key used for verity signing. It expects a filename
82      WITHOUT the extension (e.g. verity_key).
83
84  --replace_verity_public_key <key>
85      Replace the certificate (public key) used for verity verification. The
86      key file replaces the one at BOOT/RAMDISK/verity_key (or ROOT/verity_key
87      for devices using system_root_image). It expects the key filename WITH
88      the extension (e.g. verity_key.pub).
89
90  --replace_verity_keyid <path_to_X509_PEM_cert_file>
91      Replace the veritykeyid in BOOT/cmdline of input_target_file_zip
92      with keyid of the cert pointed by <path_to_X509_PEM_cert_file>.
93
94  --remove_avb_public_keys <key1>,<key2>,...
95      Remove AVB public keys from the first-stage ramdisk. The key file to
96      remove is located at either of the following dirs:
97        - BOOT/RAMDISK/avb/ or
98        - BOOT/RAMDISK/first_stage_ramdisk/avb/
99      The second dir will be used for lookup if BOARD_USES_RECOVERY_AS_BOOT is
100      set to true.
101
102  --avb_{boot,system,system_other,vendor,dtbo,vbmeta,vbmeta_system,
103         vbmeta_vendor}_algorithm <algorithm>
104  --avb_{boot,system,system_other,vendor,dtbo,vbmeta,vbmeta_system,
105         vbmeta_vendor}_key <key>
106      Use the specified algorithm (e.g. SHA256_RSA4096) and the key to AVB-sign
107      the specified image. Otherwise it uses the existing values in info dict.
108
109  --avb_{apex,boot,system,system_other,vendor,dtbo,vbmeta,vbmeta_system,
110         vbmeta_vendor}_extra_args <args>
111      Specify any additional args that are needed to AVB-sign the image
112      (e.g. "--signing_helper /path/to/helper"). The args will be appended to
113      the existing ones in info dict.
114
115  --avb_extra_custom_image_key <partition=key>
116  --avb_extra_custom_image_algorithm <partition=algorithm>
117      Use the specified algorithm (e.g. SHA256_RSA4096) and the key to AVB-sign
118      the specified custom images mounted on the partition. Otherwise it uses
119      the existing values in info dict.
120
121  --avb_extra_custom_image_extra_args <partition=extra_args>
122      Specify any additional args that are needed to AVB-sign the custom images
123      mounted on the partition (e.g. "--signing_helper /path/to/helper"). The
124      args will be appended to the existing ones in info dict.
125
126  --android_jar_path <path>
127      Path to the android.jar to repack the apex file.
128"""
129
130from __future__ import print_function
131
132import base64
133import copy
134import errno
135import gzip
136import io
137import itertools
138import logging
139import os
140import re
141import shutil
142import stat
143import subprocess
144import sys
145import tempfile
146import zipfile
147from xml.etree import ElementTree
148
149import add_img_to_target_files
150import apex_utils
151import common
152
153
154if sys.hexversion < 0x02070000:
155  print("Python 2.7 or newer is required.", file=sys.stderr)
156  sys.exit(1)
157
158
159logger = logging.getLogger(__name__)
160
161OPTIONS = common.OPTIONS
162
163OPTIONS.extra_apks = {}
164OPTIONS.extra_apex_payload_keys = {}
165OPTIONS.skip_apks_with_path_prefix = set()
166OPTIONS.key_map = {}
167OPTIONS.rebuild_recovery = False
168OPTIONS.replace_ota_keys = False
169OPTIONS.replace_verity_public_key = False
170OPTIONS.replace_verity_private_key = False
171OPTIONS.replace_verity_keyid = False
172OPTIONS.remove_avb_public_keys = None
173OPTIONS.tag_changes = ("-test-keys", "-dev-keys", "+release-keys")
174OPTIONS.avb_keys = {}
175OPTIONS.avb_algorithms = {}
176OPTIONS.avb_extra_args = {}
177OPTIONS.android_jar_path = None
178
179
180AVB_FOOTER_ARGS_BY_PARTITION = {
181    'boot' : 'avb_boot_add_hash_footer_args',
182    'dtbo' : 'avb_dtbo_add_hash_footer_args',
183    'recovery' : 'avb_recovery_add_hash_footer_args',
184    'system' : 'avb_system_add_hashtree_footer_args',
185    'system_other' : 'avb_system_other_add_hashtree_footer_args',
186    'vendor' : 'avb_vendor_add_hashtree_footer_args',
187    'vendor_boot' : 'avb_vendor_boot_add_hash_footer_args',
188    'vbmeta' : 'avb_vbmeta_args',
189    'vbmeta_system' : 'avb_vbmeta_system_args',
190    'vbmeta_vendor' : 'avb_vbmeta_vendor_args',
191}
192
193
194def GetApkCerts(certmap):
195  # apply the key remapping to the contents of the file
196  for apk, cert in certmap.items():
197    certmap[apk] = OPTIONS.key_map.get(cert, cert)
198
199  # apply all the -e options, overriding anything in the file
200  for apk, cert in OPTIONS.extra_apks.items():
201    if not cert:
202      cert = "PRESIGNED"
203    certmap[apk] = OPTIONS.key_map.get(cert, cert)
204
205  return certmap
206
207
208def GetApexKeys(keys_info, key_map):
209  """Gets APEX payload and container signing keys by applying the mapping rules.
210
211  Presigned payload / container keys will be set accordingly.
212
213  Args:
214    keys_info: A dict that maps from APEX filenames to a tuple of (payload_key,
215        container_key).
216    key_map: A dict that overrides the keys, specified via command-line input.
217
218  Returns:
219    A dict that contains the updated APEX key mapping, which should be used for
220    the current signing.
221
222  Raises:
223    AssertionError: On invalid container / payload key overrides.
224  """
225  # Apply all the --extra_apex_payload_key options to override the payload
226  # signing keys in the given keys_info.
227  for apex, key in OPTIONS.extra_apex_payload_keys.items():
228    if not key:
229      key = 'PRESIGNED'
230    if apex not in keys_info:
231      logger.warning('Failed to find %s in target_files; Ignored', apex)
232      continue
233    keys_info[apex] = (key, keys_info[apex][1])
234
235  # Apply the key remapping to container keys.
236  for apex, (payload_key, container_key) in keys_info.items():
237    keys_info[apex] = (payload_key, key_map.get(container_key, container_key))
238
239  # Apply all the --extra_apks options to override the container keys.
240  for apex, key in OPTIONS.extra_apks.items():
241    # Skip non-APEX containers.
242    if apex not in keys_info:
243      continue
244    if not key:
245      key = 'PRESIGNED'
246    keys_info[apex] = (keys_info[apex][0], key_map.get(key, key))
247
248  # A PRESIGNED container entails a PRESIGNED payload. Apply this to all the
249  # APEX key pairs. However, a PRESIGNED container with non-PRESIGNED payload
250  # (overridden via commandline) indicates a config error, which should not be
251  # allowed.
252  for apex, (payload_key, container_key) in keys_info.items():
253    if container_key != 'PRESIGNED':
254      continue
255    if apex in OPTIONS.extra_apex_payload_keys:
256      payload_override = OPTIONS.extra_apex_payload_keys[apex]
257      assert payload_override == '', \
258          ("Invalid APEX key overrides: {} has PRESIGNED container but "
259           "non-PRESIGNED payload key {}").format(apex, payload_override)
260    if payload_key != 'PRESIGNED':
261      print(
262          "Setting {} payload as PRESIGNED due to PRESIGNED container".format(
263              apex))
264    keys_info[apex] = ('PRESIGNED', 'PRESIGNED')
265
266  return keys_info
267
268
269def GetApkFileInfo(filename, compressed_extension, skipped_prefixes):
270  """Returns the APK info based on the given filename.
271
272  Checks if the given filename (with path) looks like an APK file, by taking the
273  compressed extension into consideration. If it appears to be an APK file,
274  further checks if the APK file should be skipped when signing, based on the
275  given path prefixes.
276
277  Args:
278    filename: Path to the file.
279    compressed_extension: The extension string of compressed APKs (e.g. ".gz"),
280        or None if there's no compressed APKs.
281    skipped_prefixes: A set/list/tuple of the path prefixes to be skipped.
282
283  Returns:
284    (is_apk, is_compressed, should_be_skipped): is_apk indicates whether the
285    given filename is an APK file. is_compressed indicates whether the APK file
286    is compressed (only meaningful when is_apk is True). should_be_skipped
287    indicates whether the filename matches any of the given prefixes to be
288    skipped.
289
290  Raises:
291    AssertionError: On invalid compressed_extension or skipped_prefixes inputs.
292  """
293  assert compressed_extension is None or compressed_extension.startswith('.'), \
294      "Invalid compressed_extension arg: '{}'".format(compressed_extension)
295
296  # skipped_prefixes should be one of set/list/tuple types. Other types such as
297  # str shouldn't be accepted.
298  assert isinstance(skipped_prefixes, (set, list, tuple)), \
299      "Invalid skipped_prefixes input type: {}".format(type(skipped_prefixes))
300
301  compressed_apk_extension = (
302      ".apk" + compressed_extension if compressed_extension else None)
303  is_apk = (filename.endswith(".apk") or
304            (compressed_apk_extension and
305             filename.endswith(compressed_apk_extension)))
306  if not is_apk:
307    return (False, False, False)
308
309  is_compressed = (compressed_apk_extension and
310                   filename.endswith(compressed_apk_extension))
311  should_be_skipped = filename.startswith(tuple(skipped_prefixes))
312  return (True, is_compressed, should_be_skipped)
313
314
315def CheckApkAndApexKeysAvailable(input_tf_zip, known_keys,
316                                 compressed_extension, apex_keys):
317  """Checks that all the APKs and APEXes have keys specified.
318
319  Args:
320    input_tf_zip: An open target_files zip file.
321    known_keys: A set of APKs and APEXes that have known signing keys.
322    compressed_extension: The extension string of compressed APKs, such as
323        '.gz', or None if there's no compressed APKs.
324    apex_keys: A dict that contains the key mapping from APEX name to
325        (payload_key, container_key).
326
327  Raises:
328    AssertionError: On finding unknown APKs and APEXes.
329  """
330  unknown_files = []
331  for info in input_tf_zip.infolist():
332    # Handle APEXes first, e.g. SYSTEM/apex/com.android.tzdata.apex.
333    if (info.filename.startswith('SYSTEM/apex') and
334        info.filename.endswith('.apex')):
335      name = os.path.basename(info.filename)
336      if name not in known_keys:
337        unknown_files.append(name)
338      continue
339
340    # And APKs.
341    (is_apk, is_compressed, should_be_skipped) = GetApkFileInfo(
342        info.filename, compressed_extension, OPTIONS.skip_apks_with_path_prefix)
343    if not is_apk or should_be_skipped:
344      continue
345
346    name = os.path.basename(info.filename)
347    if is_compressed:
348      name = name[:-len(compressed_extension)]
349    if name not in known_keys:
350      unknown_files.append(name)
351
352  assert not unknown_files, \
353      ("No key specified for:\n  {}\n"
354       "Use '-e <apkname>=' to specify a key (which may be an empty string to "
355       "not sign this apk).".format("\n  ".join(unknown_files)))
356
357  # For all the APEXes, double check that we won't have an APEX that has only
358  # one of the payload / container keys set. Note that non-PRESIGNED container
359  # with PRESIGNED payload could be allowed but currently unsupported. It would
360  # require changing SignApex implementation.
361  if not apex_keys:
362    return
363
364  invalid_apexes = []
365  for info in input_tf_zip.infolist():
366    if (not info.filename.startswith('SYSTEM/apex') or
367        not info.filename.endswith('.apex')):
368      continue
369
370    name = os.path.basename(info.filename)
371    (payload_key, container_key) = apex_keys[name]
372    if ((payload_key in common.SPECIAL_CERT_STRINGS and
373         container_key not in common.SPECIAL_CERT_STRINGS) or
374        (payload_key not in common.SPECIAL_CERT_STRINGS and
375         container_key in common.SPECIAL_CERT_STRINGS)):
376      invalid_apexes.append(
377          "{}: payload_key {}, container_key {}".format(
378              name, payload_key, container_key))
379
380  assert not invalid_apexes, \
381      "Invalid APEX keys specified:\n  {}\n".format(
382          "\n  ".join(invalid_apexes))
383
384
385def SignApk(data, keyname, pw, platform_api_level, codename_to_api_level_map,
386            is_compressed, apk_name):
387  unsigned = tempfile.NamedTemporaryFile(suffix='_' + apk_name)
388  unsigned.write(data)
389  unsigned.flush()
390
391  if is_compressed:
392    uncompressed = tempfile.NamedTemporaryFile()
393    with gzip.open(unsigned.name, "rb") as in_file, \
394         open(uncompressed.name, "wb") as out_file:
395      shutil.copyfileobj(in_file, out_file)
396
397    # Finally, close the "unsigned" file (which is gzip compressed), and then
398    # replace it with the uncompressed version.
399    #
400    # TODO(narayan): All this nastiness can be avoided if python 3.2 is in use,
401    # we could just gzip / gunzip in-memory buffers instead.
402    unsigned.close()
403    unsigned = uncompressed
404
405  signed = tempfile.NamedTemporaryFile(suffix='_' + apk_name)
406
407  # For pre-N builds, don't upgrade to SHA-256 JAR signatures based on the APK's
408  # minSdkVersion to avoid increasing incremental OTA update sizes. If an APK
409  # didn't change, we don't want its signature to change due to the switch
410  # from SHA-1 to SHA-256.
411  # By default, APK signer chooses SHA-256 signatures if the APK's minSdkVersion
412  # is 18 or higher. For pre-N builds we disable this mechanism by pretending
413  # that the APK's minSdkVersion is 1.
414  # For N+ builds, we let APK signer rely on the APK's minSdkVersion to
415  # determine whether to use SHA-256.
416  min_api_level = None
417  if platform_api_level > 23:
418    # Let APK signer choose whether to use SHA-1 or SHA-256, based on the APK's
419    # minSdkVersion attribute
420    min_api_level = None
421  else:
422    # Force APK signer to use SHA-1
423    min_api_level = 1
424
425  common.SignFile(unsigned.name, signed.name, keyname, pw,
426                  min_api_level=min_api_level,
427                  codename_to_api_level_map=codename_to_api_level_map)
428
429  data = None
430  if is_compressed:
431    # Recompress the file after it has been signed.
432    compressed = tempfile.NamedTemporaryFile()
433    with open(signed.name, "rb") as in_file, \
434         gzip.open(compressed.name, "wb") as out_file:
435      shutil.copyfileobj(in_file, out_file)
436
437    data = compressed.read()
438    compressed.close()
439  else:
440    data = signed.read()
441
442  unsigned.close()
443  signed.close()
444
445  return data
446
447
448def ProcessTargetFiles(input_tf_zip, output_tf_zip, misc_info,
449                       apk_keys, apex_keys, key_passwords,
450                       platform_api_level, codename_to_api_level_map,
451                       compressed_extension):
452  # maxsize measures the maximum filename length, including the ones to be
453  # skipped.
454  maxsize = max(
455      [len(os.path.basename(i.filename)) for i in input_tf_zip.infolist()
456       if GetApkFileInfo(i.filename, compressed_extension, [])[0]])
457  system_root_image = misc_info.get("system_root_image") == "true"
458
459  for info in input_tf_zip.infolist():
460    filename = info.filename
461    if filename.startswith("IMAGES/"):
462      continue
463
464    # Skip OTA-specific images (e.g. split super images), which will be
465    # re-generated during signing.
466    if filename.startswith("OTA/") and filename.endswith(".img"):
467      continue
468
469    data = input_tf_zip.read(filename)
470    out_info = copy.copy(info)
471    (is_apk, is_compressed, should_be_skipped) = GetApkFileInfo(
472        filename, compressed_extension, OPTIONS.skip_apks_with_path_prefix)
473
474    if is_apk and should_be_skipped:
475      # Copy skipped APKs verbatim.
476      print(
477          "NOT signing: %s\n"
478          "        (skipped due to matching prefix)" % (filename,))
479      common.ZipWriteStr(output_tf_zip, out_info, data)
480
481    # Sign APKs.
482    elif is_apk:
483      name = os.path.basename(filename)
484      if is_compressed:
485        name = name[:-len(compressed_extension)]
486
487      key = apk_keys[name]
488      if key not in common.SPECIAL_CERT_STRINGS:
489        print("    signing: %-*s (%s)" % (maxsize, name, key))
490        signed_data = SignApk(data, key, key_passwords[key], platform_api_level,
491                              codename_to_api_level_map, is_compressed, name)
492        common.ZipWriteStr(output_tf_zip, out_info, signed_data)
493      else:
494        # an APK we're not supposed to sign.
495        print(
496            "NOT signing: %s\n"
497            "        (skipped due to special cert string)" % (name,))
498        common.ZipWriteStr(output_tf_zip, out_info, data)
499
500    # Sign bundled APEX files.
501    elif filename.startswith("SYSTEM/apex") and filename.endswith(".apex"):
502      name = os.path.basename(filename)
503      payload_key, container_key = apex_keys[name]
504
505      # We've asserted not having a case with only one of them PRESIGNED.
506      if (payload_key not in common.SPECIAL_CERT_STRINGS and
507          container_key not in common.SPECIAL_CERT_STRINGS):
508        print("    signing: %-*s container (%s)" % (
509            maxsize, name, container_key))
510        print("           : %-*s payload   (%s)" % (
511            maxsize, name, payload_key))
512
513        signed_apex = apex_utils.SignApex(
514            misc_info['avb_avbtool'],
515            data,
516            payload_key,
517            container_key,
518            key_passwords[container_key],
519            apk_keys,
520            codename_to_api_level_map,
521            no_hashtree=True,
522            signing_args=OPTIONS.avb_extra_args.get('apex'))
523        common.ZipWrite(output_tf_zip, signed_apex, filename)
524
525      else:
526        print(
527            "NOT signing: %s\n"
528            "        (skipped due to special cert string)" % (name,))
529        common.ZipWriteStr(output_tf_zip, out_info, data)
530
531    # AVB public keys for the installed APEXes, which will be updated later.
532    elif (os.path.dirname(filename) == 'SYSTEM/etc/security/apex' and
533          filename != 'SYSTEM/etc/security/apex/'):
534      continue
535
536    # System properties.
537    elif filename in (
538        "SYSTEM/build.prop",
539
540        "VENDOR/build.prop",
541        "SYSTEM/vendor/build.prop",
542
543        "ODM/etc/build.prop",
544        "VENDOR/odm/etc/build.prop",
545
546        "PRODUCT/build.prop",
547        "SYSTEM/product/build.prop",
548
549        "SYSTEM_EXT/build.prop",
550        "SYSTEM/system_ext/build.prop",
551
552        "SYSTEM/etc/prop.default",
553        "BOOT/RAMDISK/prop.default",
554        "RECOVERY/RAMDISK/prop.default",
555
556        # ROOT/default.prop is a legacy path, but may still exist for upgrading
557        # devices that don't support `property_overrides_split_enabled`.
558        "ROOT/default.prop",
559
560        # RECOVERY/RAMDISK/default.prop is a legacy path, but will always exist
561        # as a symlink in the current code. So it's a no-op here. Keeping the
562        # path here for clarity.
563        "RECOVERY/RAMDISK/default.prop"):
564      print("Rewriting %s:" % (filename,))
565      if stat.S_ISLNK(info.external_attr >> 16):
566        new_data = data
567      else:
568        new_data = RewriteProps(data.decode())
569      common.ZipWriteStr(output_tf_zip, out_info, new_data)
570
571    # Replace the certs in *mac_permissions.xml (there could be multiple, such
572    # as {system,vendor}/etc/selinux/{plat,nonplat}_mac_permissions.xml).
573    elif filename.endswith("mac_permissions.xml"):
574      print("Rewriting %s with new keys." % (filename,))
575      new_data = ReplaceCerts(data.decode())
576      common.ZipWriteStr(output_tf_zip, out_info, new_data)
577
578    # Ask add_img_to_target_files to rebuild the recovery patch if needed.
579    elif filename in ("SYSTEM/recovery-from-boot.p",
580                      "VENDOR/recovery-from-boot.p",
581
582                      "SYSTEM/etc/recovery.img",
583                      "VENDOR/etc/recovery.img",
584
585                      "SYSTEM/bin/install-recovery.sh",
586                      "VENDOR/bin/install-recovery.sh"):
587      OPTIONS.rebuild_recovery = True
588
589    # Don't copy OTA certs if we're replacing them.
590    # Replacement of update-payload-key.pub.pem was removed in b/116660991.
591    elif (
592        OPTIONS.replace_ota_keys and
593        filename in (
594            "BOOT/RAMDISK/system/etc/security/otacerts.zip",
595            "RECOVERY/RAMDISK/system/etc/security/otacerts.zip",
596            "SYSTEM/etc/security/otacerts.zip")):
597      pass
598
599    # Skip META/misc_info.txt since we will write back the new values later.
600    elif filename == "META/misc_info.txt":
601      pass
602
603    # Skip verity public key if we will replace it.
604    elif (OPTIONS.replace_verity_public_key and
605          filename in ("BOOT/RAMDISK/verity_key",
606                       "ROOT/verity_key")):
607      pass
608    elif (OPTIONS.remove_avb_public_keys and
609          (filename.startswith("BOOT/RAMDISK/avb/") or
610           filename.startswith("BOOT/RAMDISK/first_stage_ramdisk/avb/"))):
611      matched_removal = False
612      for key_to_remove in OPTIONS.remove_avb_public_keys:
613        if filename.endswith(key_to_remove):
614          matched_removal = True
615          print("Removing AVB public key from ramdisk: %s" % filename)
616          break
617      if not matched_removal:
618        # Copy it verbatim if we don't want to remove it.
619        common.ZipWriteStr(output_tf_zip, out_info, data)
620
621    # Skip verity keyid (for system_root_image use) if we will replace it.
622    elif OPTIONS.replace_verity_keyid and filename == "BOOT/cmdline":
623      pass
624
625    # Skip the care_map as we will regenerate the system/vendor images.
626    elif filename in ["META/care_map.pb", "META/care_map.txt"]:
627      pass
628
629    # Updates system_other.avbpubkey in /product/etc/.
630    elif filename in (
631        "PRODUCT/etc/security/avb/system_other.avbpubkey",
632        "SYSTEM/product/etc/security/avb/system_other.avbpubkey"):
633      # Only update system_other's public key, if the corresponding signing
634      # key is specified via --avb_system_other_key.
635      signing_key = OPTIONS.avb_keys.get("system_other")
636      if signing_key:
637        public_key = common.ExtractAvbPublicKey(
638            misc_info['avb_avbtool'], signing_key)
639        print("    Rewriting AVB public key of system_other in /product")
640        common.ZipWrite(output_tf_zip, public_key, filename)
641
642    # Should NOT sign boot-debug.img.
643    elif filename in (
644        "BOOT/RAMDISK/force_debuggable",
645        "BOOT/RAMDISK/first_stage_ramdisk/force_debuggable"):
646      raise common.ExternalError("debuggable boot.img cannot be signed")
647
648    # A non-APK file; copy it verbatim.
649    else:
650      common.ZipWriteStr(output_tf_zip, out_info, data)
651
652  if OPTIONS.replace_ota_keys:
653    ReplaceOtaKeys(input_tf_zip, output_tf_zip, misc_info)
654
655  # Replace the keyid string in misc_info dict.
656  if OPTIONS.replace_verity_private_key:
657    ReplaceVerityPrivateKey(misc_info, OPTIONS.replace_verity_private_key[1])
658
659  if OPTIONS.replace_verity_public_key:
660    # Replace the one in root dir in system.img.
661    ReplaceVerityPublicKey(
662        output_tf_zip, 'ROOT/verity_key', OPTIONS.replace_verity_public_key[1])
663
664    if not system_root_image:
665      # Additionally replace the copy in ramdisk if not using system-as-root.
666      ReplaceVerityPublicKey(
667          output_tf_zip,
668          'BOOT/RAMDISK/verity_key',
669          OPTIONS.replace_verity_public_key[1])
670
671  # Replace the keyid string in BOOT/cmdline.
672  if OPTIONS.replace_verity_keyid:
673    ReplaceVerityKeyId(input_tf_zip, output_tf_zip,
674                       OPTIONS.replace_verity_keyid[1])
675
676  # Replace the AVB signing keys, if any.
677  ReplaceAvbSigningKeys(misc_info)
678
679  # Rewrite the props in AVB signing args.
680  if misc_info.get('avb_enable') == 'true':
681    RewriteAvbProps(misc_info)
682
683  # Write back misc_info with the latest values.
684  ReplaceMiscInfoTxt(input_tf_zip, output_tf_zip, misc_info)
685
686
687def ReplaceCerts(data):
688  """Replaces all the occurences of X.509 certs with the new ones.
689
690  The mapping info is read from OPTIONS.key_map. Non-existent certificate will
691  be skipped. After the replacement, it additionally checks for duplicate
692  entries, which would otherwise fail the policy loading code in
693  frameworks/base/services/core/java/com/android/server/pm/SELinuxMMAC.java.
694
695  Args:
696    data: Input string that contains a set of X.509 certs.
697
698  Returns:
699    A string after the replacement.
700
701  Raises:
702    AssertionError: On finding duplicate entries.
703  """
704  for old, new in OPTIONS.key_map.items():
705    if OPTIONS.verbose:
706      print("    Replacing %s.x509.pem with %s.x509.pem" % (old, new))
707
708    try:
709      with open(old + ".x509.pem") as old_fp:
710        old_cert16 = base64.b16encode(
711            common.ParseCertificate(old_fp.read())).decode().lower()
712      with open(new + ".x509.pem") as new_fp:
713        new_cert16 = base64.b16encode(
714            common.ParseCertificate(new_fp.read())).decode().lower()
715    except IOError as e:
716      if OPTIONS.verbose or e.errno != errno.ENOENT:
717        print("    Error accessing %s: %s.\nSkip replacing %s.x509.pem with "
718              "%s.x509.pem." % (e.filename, e.strerror, old, new))
719      continue
720
721    # Only match entire certs.
722    pattern = "\\b" + old_cert16 + "\\b"
723    (data, num) = re.subn(pattern, new_cert16, data, flags=re.IGNORECASE)
724
725    if OPTIONS.verbose:
726      print("    Replaced %d occurence(s) of %s.x509.pem with %s.x509.pem" % (
727          num, old, new))
728
729  # Verify that there're no duplicate entries after the replacement. Note that
730  # it's only checking entries with global seinfo at the moment (i.e. ignoring
731  # the ones with inner packages). (Bug: 69479366)
732  root = ElementTree.fromstring(data)
733  signatures = [signer.attrib['signature'] for signer in root.findall('signer')]
734  assert len(signatures) == len(set(signatures)), \
735      "Found duplicate entries after cert replacement: {}".format(data)
736
737  return data
738
739
740def EditTags(tags):
741  """Applies the edits to the tag string as specified in OPTIONS.tag_changes.
742
743  Args:
744    tags: The input string that contains comma-separated tags.
745
746  Returns:
747    The updated tags (comma-separated and sorted).
748  """
749  tags = set(tags.split(","))
750  for ch in OPTIONS.tag_changes:
751    if ch[0] == "-":
752      tags.discard(ch[1:])
753    elif ch[0] == "+":
754      tags.add(ch[1:])
755  return ",".join(sorted(tags))
756
757
758def RewriteProps(data):
759  """Rewrites the system properties in the given string.
760
761  Each property is expected in 'key=value' format. The properties that contain
762  build tags (i.e. test-keys, dev-keys) will be updated accordingly by calling
763  EditTags().
764
765  Args:
766    data: Input string, separated by newlines.
767
768  Returns:
769    The string with modified properties.
770  """
771  output = []
772  for line in data.split("\n"):
773    line = line.strip()
774    original_line = line
775    if line and line[0] != '#' and "=" in line:
776      key, value = line.split("=", 1)
777      if (key.startswith("ro.") and
778          key.endswith((".build.fingerprint", ".build.thumbprint"))):
779        pieces = value.split("/")
780        pieces[-1] = EditTags(pieces[-1])
781        value = "/".join(pieces)
782      elif key == "ro.bootimage.build.fingerprint":
783        pieces = value.split("/")
784        pieces[-1] = EditTags(pieces[-1])
785        value = "/".join(pieces)
786      elif key == "ro.build.description":
787        pieces = value.split(" ")
788        assert len(pieces) == 5
789        pieces[-1] = EditTags(pieces[-1])
790        value = " ".join(pieces)
791      elif key.startswith("ro.") and key.endswith(".build.tags"):
792        value = EditTags(value)
793      elif key == "ro.build.display.id":
794        # change, eg, "JWR66N dev-keys" to "JWR66N"
795        value = value.split()
796        if len(value) > 1 and value[-1].endswith("-keys"):
797          value.pop()
798        value = " ".join(value)
799      line = key + "=" + value
800    if line != original_line:
801      print("  replace: ", original_line)
802      print("     with: ", line)
803    output.append(line)
804  return "\n".join(output) + "\n"
805
806
807def WriteOtacerts(output_zip, filename, keys):
808  """Constructs a zipfile from given keys; and writes it to output_zip.
809
810  Args:
811    output_zip: The output target_files zip.
812    filename: The archive name in the output zip.
813    keys: A list of public keys to use during OTA package verification.
814  """
815  temp_file = io.BytesIO()
816  certs_zip = zipfile.ZipFile(temp_file, "w")
817  for k in keys:
818    common.ZipWrite(certs_zip, k)
819  common.ZipClose(certs_zip)
820  common.ZipWriteStr(output_zip, filename, temp_file.getvalue())
821
822
823def ReplaceOtaKeys(input_tf_zip, output_tf_zip, misc_info):
824  try:
825    keylist = input_tf_zip.read("META/otakeys.txt").split()
826  except KeyError:
827    raise common.ExternalError("can't read META/otakeys.txt from input")
828
829  extra_recovery_keys = misc_info.get("extra_recovery_keys")
830  if extra_recovery_keys:
831    extra_recovery_keys = [OPTIONS.key_map.get(k, k) + ".x509.pem"
832                           for k in extra_recovery_keys.split()]
833    if extra_recovery_keys:
834      print("extra recovery-only key(s): " + ", ".join(extra_recovery_keys))
835  else:
836    extra_recovery_keys = []
837
838  mapped_keys = []
839  for k in keylist:
840    m = re.match(r"^(.*)\.x509\.pem$", k)
841    if not m:
842      raise common.ExternalError(
843          "can't parse \"%s\" from META/otakeys.txt" % (k,))
844    k = m.group(1)
845    mapped_keys.append(OPTIONS.key_map.get(k, k) + ".x509.pem")
846
847  if mapped_keys:
848    print("using:\n   ", "\n   ".join(mapped_keys))
849    print("for OTA package verification")
850  else:
851    devkey = misc_info.get("default_system_dev_certificate",
852                           "build/make/target/product/security/testkey")
853    mapped_devkey = OPTIONS.key_map.get(devkey, devkey)
854    if mapped_devkey != devkey:
855      misc_info["default_system_dev_certificate"] = mapped_devkey
856    mapped_keys.append(mapped_devkey + ".x509.pem")
857    print("META/otakeys.txt has no keys; using %s for OTA package"
858          " verification." % (mapped_keys[0],))
859
860  # recovery now uses the same x509.pem version of the keys.
861  # extra_recovery_keys are used only in recovery.
862  if misc_info.get("recovery_as_boot") == "true":
863    recovery_keys_location = "BOOT/RAMDISK/system/etc/security/otacerts.zip"
864  else:
865    recovery_keys_location = "RECOVERY/RAMDISK/system/etc/security/otacerts.zip"
866
867  WriteOtacerts(output_tf_zip, recovery_keys_location,
868                mapped_keys + extra_recovery_keys)
869
870  # SystemUpdateActivity uses the x509.pem version of the keys, but
871  # put into a zipfile system/etc/security/otacerts.zip.
872  # We DO NOT include the extra_recovery_keys (if any) here.
873  WriteOtacerts(output_tf_zip, "SYSTEM/etc/security/otacerts.zip", mapped_keys)
874
875
876
877def ReplaceVerityPublicKey(output_zip, filename, key_path):
878  """Replaces the verity public key at the given path in the given zip.
879
880  Args:
881    output_zip: The output target_files zip.
882    filename: The archive name in the output zip.
883    key_path: The path to the public key.
884  """
885  print("Replacing verity public key with %s" % (key_path,))
886  common.ZipWrite(output_zip, key_path, arcname=filename)
887
888
889def ReplaceVerityPrivateKey(misc_info, key_path):
890  """Replaces the verity private key in misc_info dict.
891
892  Args:
893    misc_info: The info dict.
894    key_path: The path to the private key in PKCS#8 format.
895  """
896  print("Replacing verity private key with %s" % (key_path,))
897  misc_info["verity_key"] = key_path
898
899
900def ReplaceVerityKeyId(input_zip, output_zip, key_path):
901  """Replaces the veritykeyid parameter in BOOT/cmdline.
902
903  Args:
904    input_zip: The input target_files zip, which should be already open.
905    output_zip: The output target_files zip, which should be already open and
906        writable.
907    key_path: The path to the PEM encoded X.509 certificate.
908  """
909  in_cmdline = input_zip.read("BOOT/cmdline").decode()
910  # Copy in_cmdline to output_zip if veritykeyid is not present.
911  if "veritykeyid" not in in_cmdline:
912    common.ZipWriteStr(output_zip, "BOOT/cmdline", in_cmdline)
913    return
914
915  out_buffer = []
916  for param in in_cmdline.split():
917    if "veritykeyid" not in param:
918      out_buffer.append(param)
919      continue
920
921    # Extract keyid using openssl command.
922    p = common.Run(["openssl", "x509", "-in", key_path, "-text"],
923                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
924    keyid, stderr = p.communicate()
925    assert p.returncode == 0, "Failed to dump certificate: {}".format(stderr)
926    keyid = re.search(
927        r'keyid:([0-9a-fA-F:]*)', keyid).group(1).replace(':', '').lower()
928    print("Replacing verity keyid with {}".format(keyid))
929    out_buffer.append("veritykeyid=id:%s" % (keyid,))
930
931  out_cmdline = ' '.join(out_buffer).strip() + '\n'
932  common.ZipWriteStr(output_zip, "BOOT/cmdline", out_cmdline)
933
934
935def ReplaceMiscInfoTxt(input_zip, output_zip, misc_info):
936  """Replaces META/misc_info.txt.
937
938  Only writes back the ones in the original META/misc_info.txt. Because the
939  current in-memory dict contains additional items computed at runtime.
940  """
941  misc_info_old = common.LoadDictionaryFromLines(
942      input_zip.read('META/misc_info.txt').decode().split('\n'))
943  items = []
944  for key in sorted(misc_info):
945    if key in misc_info_old:
946      items.append('%s=%s' % (key, misc_info[key]))
947  common.ZipWriteStr(output_zip, "META/misc_info.txt", '\n'.join(items))
948
949
950def ReplaceAvbSigningKeys(misc_info):
951  """Replaces the AVB signing keys."""
952
953  def ReplaceAvbPartitionSigningKey(partition):
954    key = OPTIONS.avb_keys.get(partition)
955    if not key:
956      return
957
958    algorithm = OPTIONS.avb_algorithms.get(partition)
959    assert algorithm, 'Missing AVB signing algorithm for %s' % (partition,)
960
961    print('Replacing AVB signing key for %s with "%s" (%s)' % (
962        partition, key, algorithm))
963    misc_info['avb_' + partition + '_algorithm'] = algorithm
964    misc_info['avb_' + partition + '_key_path'] = key
965
966    extra_args = OPTIONS.avb_extra_args.get(partition)
967    if extra_args:
968      print('Setting extra AVB signing args for %s to "%s"' % (
969          partition, extra_args))
970      args_key = AVB_FOOTER_ARGS_BY_PARTITION.get(
971          partition,
972          # custom partition
973          "avb_{}_add_hashtree_footer_args".format(partition))
974      misc_info[args_key] = (misc_info.get(args_key, '') + ' ' + extra_args)
975
976  for partition in AVB_FOOTER_ARGS_BY_PARTITION:
977    ReplaceAvbPartitionSigningKey(partition)
978
979  for custom_partition in misc_info.get(
980      "avb_custom_images_partition_list", "").strip().split():
981    ReplaceAvbPartitionSigningKey(custom_partition)
982
983
984def RewriteAvbProps(misc_info):
985  """Rewrites the props in AVB signing args."""
986  for partition, args_key in AVB_FOOTER_ARGS_BY_PARTITION.items():
987    args = misc_info.get(args_key)
988    if not args:
989      continue
990
991    tokens = []
992    changed = False
993    for token in args.split(' '):
994      fingerprint_key = 'com.android.build.{}.fingerprint'.format(partition)
995      if not token.startswith(fingerprint_key):
996        tokens.append(token)
997        continue
998      prefix, tag = token.rsplit('/', 1)
999      tokens.append('{}/{}'.format(prefix, EditTags(tag)))
1000      changed = True
1001
1002    if changed:
1003      result = ' '.join(tokens)
1004      print('Rewriting AVB prop for {}:\n'.format(partition))
1005      print('  replace: {}'.format(args))
1006      print('     with: {}'.format(result))
1007      misc_info[args_key] = result
1008
1009
1010def BuildKeyMap(misc_info, key_mapping_options):
1011  for s, d in key_mapping_options:
1012    if s is None:   # -d option
1013      devkey = misc_info.get("default_system_dev_certificate",
1014                             "build/make/target/product/security/testkey")
1015      devkeydir = os.path.dirname(devkey)
1016
1017      OPTIONS.key_map.update({
1018          devkeydir + "/testkey":  d + "/releasekey",
1019          devkeydir + "/devkey":   d + "/releasekey",
1020          devkeydir + "/media":    d + "/media",
1021          devkeydir + "/shared":   d + "/shared",
1022          devkeydir + "/platform": d + "/platform",
1023          devkeydir + "/networkstack": d + "/networkstack",
1024          })
1025    else:
1026      OPTIONS.key_map[s] = d
1027
1028
1029def GetApiLevelAndCodename(input_tf_zip):
1030  data = input_tf_zip.read("SYSTEM/build.prop").decode()
1031  api_level = None
1032  codename = None
1033  for line in data.split("\n"):
1034    line = line.strip()
1035    if line and line[0] != '#' and "=" in line:
1036      key, value = line.split("=", 1)
1037      key = key.strip()
1038      if key == "ro.build.version.sdk":
1039        api_level = int(value.strip())
1040      elif key == "ro.build.version.codename":
1041        codename = value.strip()
1042
1043  if api_level is None:
1044    raise ValueError("No ro.build.version.sdk in SYSTEM/build.prop")
1045  if codename is None:
1046    raise ValueError("No ro.build.version.codename in SYSTEM/build.prop")
1047
1048  return (api_level, codename)
1049
1050
1051def GetCodenameToApiLevelMap(input_tf_zip):
1052  data = input_tf_zip.read("SYSTEM/build.prop").decode()
1053  api_level = None
1054  codenames = None
1055  for line in data.split("\n"):
1056    line = line.strip()
1057    if line and line[0] != '#' and "=" in line:
1058      key, value = line.split("=", 1)
1059      key = key.strip()
1060      if key == "ro.build.version.sdk":
1061        api_level = int(value.strip())
1062      elif key == "ro.build.version.all_codenames":
1063        codenames = value.strip().split(",")
1064
1065  if api_level is None:
1066    raise ValueError("No ro.build.version.sdk in SYSTEM/build.prop")
1067  if codenames is None:
1068    raise ValueError("No ro.build.version.all_codenames in SYSTEM/build.prop")
1069
1070  result = {}
1071  for codename in codenames:
1072    codename = codename.strip()
1073    if codename:
1074      result[codename] = api_level
1075  return result
1076
1077
1078def ReadApexKeysInfo(tf_zip):
1079  """Parses the APEX keys info from a given target-files zip.
1080
1081  Given a target-files ZipFile, parses the META/apexkeys.txt entry and returns a
1082  dict that contains the mapping from APEX names (e.g. com.android.tzdata) to a
1083  tuple of (payload_key, container_key).
1084
1085  Args:
1086    tf_zip: The input target_files ZipFile (already open).
1087
1088  Returns:
1089    (payload_key, container_key): payload_key contains the path to the payload
1090        signing key; container_key contains the path to the container signing
1091        key.
1092  """
1093  keys = {}
1094  for line in tf_zip.read('META/apexkeys.txt').decode().split('\n'):
1095    line = line.strip()
1096    if not line:
1097      continue
1098    matches = re.match(
1099        r'^name="(?P<NAME>.*)"\s+'
1100        r'public_key="(?P<PAYLOAD_PUBLIC_KEY>.*)"\s+'
1101        r'private_key="(?P<PAYLOAD_PRIVATE_KEY>.*)"\s+'
1102        r'container_certificate="(?P<CONTAINER_CERT>.*)"\s+'
1103        r'container_private_key="(?P<CONTAINER_PRIVATE_KEY>.*?)"'
1104        r'(\s+partition="(?P<PARTITION>.*?)")?$',
1105        line)
1106    if not matches:
1107      continue
1108
1109    name = matches.group('NAME')
1110    payload_private_key = matches.group("PAYLOAD_PRIVATE_KEY")
1111
1112    def CompareKeys(pubkey, pubkey_suffix, privkey, privkey_suffix):
1113      pubkey_suffix_len = len(pubkey_suffix)
1114      privkey_suffix_len = len(privkey_suffix)
1115      return (pubkey.endswith(pubkey_suffix) and
1116              privkey.endswith(privkey_suffix) and
1117              pubkey[:-pubkey_suffix_len] == privkey[:-privkey_suffix_len])
1118
1119    # Check the container key names, as we'll carry them without the
1120    # extensions. This doesn't apply to payload keys though, which we will use
1121    # full names only.
1122    container_cert = matches.group("CONTAINER_CERT")
1123    container_private_key = matches.group("CONTAINER_PRIVATE_KEY")
1124    if container_cert == 'PRESIGNED' and container_private_key == 'PRESIGNED':
1125      container_key = 'PRESIGNED'
1126    elif CompareKeys(
1127        container_cert, OPTIONS.public_key_suffix,
1128        container_private_key, OPTIONS.private_key_suffix):
1129      container_key = container_cert[:-len(OPTIONS.public_key_suffix)]
1130    else:
1131      raise ValueError("Failed to parse container keys: \n{}".format(line))
1132
1133    keys[name] = (payload_private_key, container_key)
1134
1135  return keys
1136
1137
1138def main(argv):
1139
1140  key_mapping_options = []
1141
1142  def option_handler(o, a):
1143    if o in ("-e", "--extra_apks"):
1144      names, key = a.split("=")
1145      names = names.split(",")
1146      for n in names:
1147        OPTIONS.extra_apks[n] = key
1148    elif o == "--extra_apex_payload_key":
1149      apex_name, key = a.split("=")
1150      OPTIONS.extra_apex_payload_keys[apex_name] = key
1151    elif o == "--skip_apks_with_path_prefix":
1152      # Check the prefix, which must be in all upper case.
1153      prefix = a.split('/')[0]
1154      if not prefix or prefix != prefix.upper():
1155        raise ValueError("Invalid path prefix '%s'" % (a,))
1156      OPTIONS.skip_apks_with_path_prefix.add(a)
1157    elif o in ("-d", "--default_key_mappings"):
1158      key_mapping_options.append((None, a))
1159    elif o in ("-k", "--key_mapping"):
1160      key_mapping_options.append(a.split("=", 1))
1161    elif o in ("-o", "--replace_ota_keys"):
1162      OPTIONS.replace_ota_keys = True
1163    elif o in ("-t", "--tag_changes"):
1164      new = []
1165      for i in a.split(","):
1166        i = i.strip()
1167        if not i or i[0] not in "-+":
1168          raise ValueError("Bad tag change '%s'" % (i,))
1169        new.append(i[0] + i[1:].strip())
1170      OPTIONS.tag_changes = tuple(new)
1171    elif o == "--replace_verity_public_key":
1172      OPTIONS.replace_verity_public_key = (True, a)
1173    elif o == "--replace_verity_private_key":
1174      OPTIONS.replace_verity_private_key = (True, a)
1175    elif o == "--replace_verity_keyid":
1176      OPTIONS.replace_verity_keyid = (True, a)
1177    elif o == "--remove_avb_public_keys":
1178      OPTIONS.remove_avb_public_keys = a.split(",")
1179    elif o == "--avb_vbmeta_key":
1180      OPTIONS.avb_keys['vbmeta'] = a
1181    elif o == "--avb_vbmeta_algorithm":
1182      OPTIONS.avb_algorithms['vbmeta'] = a
1183    elif o == "--avb_vbmeta_extra_args":
1184      OPTIONS.avb_extra_args['vbmeta'] = a
1185    elif o == "--avb_boot_key":
1186      OPTIONS.avb_keys['boot'] = a
1187    elif o == "--avb_boot_algorithm":
1188      OPTIONS.avb_algorithms['boot'] = a
1189    elif o == "--avb_boot_extra_args":
1190      OPTIONS.avb_extra_args['boot'] = a
1191    elif o == "--avb_dtbo_key":
1192      OPTIONS.avb_keys['dtbo'] = a
1193    elif o == "--avb_dtbo_algorithm":
1194      OPTIONS.avb_algorithms['dtbo'] = a
1195    elif o == "--avb_dtbo_extra_args":
1196      OPTIONS.avb_extra_args['dtbo'] = a
1197    elif o == "--avb_system_key":
1198      OPTIONS.avb_keys['system'] = a
1199    elif o == "--avb_system_algorithm":
1200      OPTIONS.avb_algorithms['system'] = a
1201    elif o == "--avb_system_extra_args":
1202      OPTIONS.avb_extra_args['system'] = a
1203    elif o == "--avb_system_other_key":
1204      OPTIONS.avb_keys['system_other'] = a
1205    elif o == "--avb_system_other_algorithm":
1206      OPTIONS.avb_algorithms['system_other'] = a
1207    elif o == "--avb_system_other_extra_args":
1208      OPTIONS.avb_extra_args['system_other'] = a
1209    elif o == "--avb_vendor_key":
1210      OPTIONS.avb_keys['vendor'] = a
1211    elif o == "--avb_vendor_algorithm":
1212      OPTIONS.avb_algorithms['vendor'] = a
1213    elif o == "--avb_vendor_extra_args":
1214      OPTIONS.avb_extra_args['vendor'] = a
1215    elif o == "--avb_vbmeta_system_key":
1216      OPTIONS.avb_keys['vbmeta_system'] = a
1217    elif o == "--avb_vbmeta_system_algorithm":
1218      OPTIONS.avb_algorithms['vbmeta_system'] = a
1219    elif o == "--avb_vbmeta_system_extra_args":
1220      OPTIONS.avb_extra_args['vbmeta_system'] = a
1221    elif o == "--avb_vbmeta_vendor_key":
1222      OPTIONS.avb_keys['vbmeta_vendor'] = a
1223    elif o == "--avb_vbmeta_vendor_algorithm":
1224      OPTIONS.avb_algorithms['vbmeta_vendor'] = a
1225    elif o == "--avb_vbmeta_vendor_extra_args":
1226      OPTIONS.avb_extra_args['vbmeta_vendor'] = a
1227    elif o == "--avb_apex_extra_args":
1228      OPTIONS.avb_extra_args['apex'] = a
1229    elif o == "--avb_extra_custom_image_key":
1230      partition, key = a.split("=")
1231      OPTIONS.avb_keys[partition] = key
1232    elif o == "--avb_extra_custom_image_algorithm":
1233      partition, algorithm = a.split("=")
1234      OPTIONS.avb_algorithms[partition] = algorithm
1235    elif o == "--avb_extra_custom_image_extra_args":
1236      # Setting the maxsplit parameter to one, which will return a list with
1237      # two elements. e.g., the second '=' should not be splitted for
1238      # 'oem=--signing_helper_with_files=/tmp/avbsigner.sh'.
1239      partition, extra_args = a.split("=", 1)
1240      OPTIONS.avb_extra_args[partition] = extra_args
1241    else:
1242      return False
1243    return True
1244
1245  args = common.ParseOptions(
1246      argv, __doc__,
1247      extra_opts="e:d:k:ot:",
1248      extra_long_opts=[
1249          "extra_apks=",
1250          "extra_apex_payload_key=",
1251          "skip_apks_with_path_prefix=",
1252          "default_key_mappings=",
1253          "key_mapping=",
1254          "replace_ota_keys",
1255          "tag_changes=",
1256          "replace_verity_public_key=",
1257          "replace_verity_private_key=",
1258          "replace_verity_keyid=",
1259          "remove_avb_public_keys=",
1260          "avb_apex_extra_args=",
1261          "avb_vbmeta_algorithm=",
1262          "avb_vbmeta_key=",
1263          "avb_vbmeta_extra_args=",
1264          "avb_boot_algorithm=",
1265          "avb_boot_key=",
1266          "avb_boot_extra_args=",
1267          "avb_dtbo_algorithm=",
1268          "avb_dtbo_key=",
1269          "avb_dtbo_extra_args=",
1270          "avb_system_algorithm=",
1271          "avb_system_key=",
1272          "avb_system_extra_args=",
1273          "avb_system_other_algorithm=",
1274          "avb_system_other_key=",
1275          "avb_system_other_extra_args=",
1276          "avb_vendor_algorithm=",
1277          "avb_vendor_key=",
1278          "avb_vendor_extra_args=",
1279          "avb_vbmeta_system_algorithm=",
1280          "avb_vbmeta_system_key=",
1281          "avb_vbmeta_system_extra_args=",
1282          "avb_vbmeta_vendor_algorithm=",
1283          "avb_vbmeta_vendor_key=",
1284          "avb_vbmeta_vendor_extra_args=",
1285          "avb_extra_custom_image_key=",
1286          "avb_extra_custom_image_algorithm=",
1287          "avb_extra_custom_image_extra_args=",
1288      ],
1289      extra_option_handler=option_handler)
1290
1291  if len(args) != 2:
1292    common.Usage(__doc__)
1293    sys.exit(1)
1294
1295  common.InitLogging()
1296
1297  input_zip = zipfile.ZipFile(args[0], "r")
1298  output_zip = zipfile.ZipFile(args[1], "w",
1299                               compression=zipfile.ZIP_DEFLATED,
1300                               allowZip64=True)
1301
1302  misc_info = common.LoadInfoDict(input_zip)
1303
1304  BuildKeyMap(misc_info, key_mapping_options)
1305
1306  apk_keys_info, compressed_extension = common.ReadApkCerts(input_zip)
1307  apk_keys = GetApkCerts(apk_keys_info)
1308
1309  apex_keys_info = ReadApexKeysInfo(input_zip)
1310  apex_keys = GetApexKeys(apex_keys_info, apk_keys)
1311
1312  # TODO(xunchang) check for the apks inside the apex files, and abort early if
1313  # the keys are not available.
1314  CheckApkAndApexKeysAvailable(
1315      input_zip,
1316      set(apk_keys.keys()) | set(apex_keys.keys()),
1317      compressed_extension,
1318      apex_keys)
1319
1320  key_passwords = common.GetKeyPasswords(
1321      set(apk_keys.values()) | set(itertools.chain(*apex_keys.values())))
1322  platform_api_level, _ = GetApiLevelAndCodename(input_zip)
1323  codename_to_api_level_map = GetCodenameToApiLevelMap(input_zip)
1324
1325  ProcessTargetFiles(input_zip, output_zip, misc_info,
1326                     apk_keys, apex_keys, key_passwords,
1327                     platform_api_level, codename_to_api_level_map,
1328                     compressed_extension)
1329
1330  common.ZipClose(input_zip)
1331  common.ZipClose(output_zip)
1332
1333  # Skip building userdata.img and cache.img when signing the target files.
1334  new_args = ["--is_signing"]
1335  # add_img_to_target_files builds the system image from scratch, so the
1336  # recovery patch is guaranteed to be regenerated there.
1337  if OPTIONS.rebuild_recovery:
1338    new_args.append("--rebuild_recovery")
1339  new_args.append(args[1])
1340  add_img_to_target_files.main(new_args)
1341
1342  print("done.")
1343
1344
1345if __name__ == '__main__':
1346  try:
1347    main(sys.argv[1:])
1348  except common.ExternalError as e:
1349    print("\n   ERROR: %s\n" % (e,))
1350    sys.exit(1)
1351  finally:
1352    common.Cleanup()
1353