Merge "releasetools: Prefer the avbtool specified in target_files."

tools/releasetools/apex_utils.py

@@ -18,7 +18,6 @@ import logging
 import os.path
 import re
 import shlex
-import sys
 import zipfile
 
 import common
@@ -42,11 +41,11 @@ class ApexSigningError(Exception):
     Exception.__init__(self, message)
 
 
-def SignApexPayload(payload_file, payload_key_path, payload_key_name, algorithm,
-                    salt, signing_args=None):
+def SignApexPayload(avbtool, payload_file, payload_key_path, payload_key_name,
+                    algorithm, salt, signing_args=None):
   """Signs a given payload_file with the payload key."""
   # Add the new footer. Old footer, if any, will be replaced by avbtool.
-  cmd = ['avbtool', 'add_hashtree_footer',
+  cmd = [avbtool, 'add_hashtree_footer',
          '--do_not_generate_fec',
          '--algorithm', algorithm,
          '--key', payload_key_path,
@@ -65,12 +64,12 @@ def SignApexPayload(payload_file, payload_key_path, payload_key_name, algorithm,
 
   # Verify the signed payload image with specified public key.
   logger.info('Verifying %s', payload_file)
-  VerifyApexPayload(payload_file, payload_key_path)
+  VerifyApexPayload(avbtool, payload_file, payload_key_path)
 
 
-def VerifyApexPayload(payload_file, payload_key):
+def VerifyApexPayload(avbtool, payload_file, payload_key):
   """Verifies the APEX payload signature with the given key."""
-  cmd = ['avbtool', 'verify_image', '--image', payload_file,
+  cmd = [avbtool, 'verify_image', '--image', payload_file,
          '--key', payload_key]
   try:
     common.RunAndCheckOutput(cmd)
@@ -80,10 +79,11 @@ def VerifyApexPayload(payload_file, payload_key):
             payload_file, payload_key, e))
 
 
-def ParseApexPayloadInfo(payload_path):
+def ParseApexPayloadInfo(avbtool, payload_path):
   """Parses the APEX payload info.
 
   Args:
+    avbtool: The AVB tool to use.
     payload_path: The path to the payload image.
 
   Raises:
@@ -96,7 +96,7 @@ def ParseApexPayloadInfo(payload_path):
   if not os.path.exists(payload_path):
     raise ApexInfoError('Failed to find image: {}'.format(payload_path))
 
-  cmd = ['avbtool', 'info_image', '--image', payload_path]
+  cmd = [avbtool, 'info_image', '--image', payload_path]
   try:
     output = common.RunAndCheckOutput(cmd)
   except common.ExternalError as e:
@@ -150,7 +150,7 @@ def ParseApexPayloadInfo(payload_path):
   return payload_info
 
 
-def SignApex(apex_data, payload_key, container_key, container_pw,
+def SignApex(avbtool, apex_data, payload_key, container_key, container_pw,
              codename_to_api_level_map, signing_args=None):
   """Signs the current APEX with the given payload/container keys.
 
@@ -178,8 +178,9 @@ def SignApex(apex_data, payload_key, container_key, container_pw,
   with zipfile.ZipFile(apex_file) as apex_fd:
     payload_file = apex_fd.extract(APEX_PAYLOAD_IMAGE, payload_dir)
 
-  payload_info = ParseApexPayloadInfo(payload_file)
+  payload_info = ParseApexPayloadInfo(avbtool, payload_file)
   SignApexPayload(
+      avbtool,
       payload_file,
       payload_key,
       payload_info['apex.key'],
@@ -188,7 +189,7 @@ def SignApex(apex_data, payload_key, container_key, container_pw,
       signing_args)
 
   # 1b. Update the embedded payload public key.
-  payload_public_key = common.ExtractAvbPublicKey(payload_key)
+  payload_public_key = common.ExtractAvbPublicKey(avbtool, payload_key)
 
   common.ZipDelete(apex_file, APEX_PAYLOAD_IMAGE)
   common.ZipDelete(apex_file, APEX_PUBKEY)

tools/releasetools/common.py

@@ -578,7 +578,7 @@ def GetAvbChainedPartitionArg(partition, info_dict, key=None):
   """
   if key is None:
     key = info_dict["avb_" + partition + "_key_path"]
-  pubkey_path = ExtractAvbPublicKey(key)
+  pubkey_path = ExtractAvbPublicKey(info_dict["avb_avbtool"], key)
   rollback_index_location = info_dict[
       "avb_" + partition + "_rollback_index_location"]
   return "{}:{}:{}".format(partition, rollback_index_location, pubkey_path)
@@ -2239,10 +2239,11 @@ def ExtractPublicKey(cert):
   return pubkey
 
 
-def ExtractAvbPublicKey(key):
+def ExtractAvbPublicKey(avbtool, key):
   """Extracts the AVB public key from the given public or private key.
 
   Args:
+    avbtool: The AVB tool to use.
     key: The input key file, which should be PEM-encoded public or private key.
 
   Returns:
@@ -2250,7 +2251,7 @@ def ExtractAvbPublicKey(key):
   """
   output = MakeTempFile(prefix='avb-', suffix='.avbpubkey')
   RunAndCheckOutput(
-      ['avbtool', 'extract_public_key', "--key", key, "--output", output])
+      [avbtool, 'extract_public_key', "--key", key, "--output", output])
   return output
 
 

tools/releasetools/sign_apex.py

@@ -19,6 +19,9 @@ Signs a standalone APEX file.
 
 Usage:  sign_apex [flags] input_apex_file output_apex_file
 
+  --avbtool <avbtool>
+      Optional flag that specifies the AVB tool to use. Defaults to `avbtool`.
+
   --container_key <key>
       Mandatory flag that specifies the container signing key.
 
@@ -40,12 +43,14 @@ import common
 logger = logging.getLogger(__name__)
 
 
-def SignApexFile(apex_file, payload_key, container_key, signing_args=None):
+def SignApexFile(avbtool, apex_file, payload_key, container_key,
+                 signing_args=None):
   """Signs the given apex file."""
   with open(apex_file, 'rb') as input_fp:
     apex_data = input_fp.read()
 
   return apex_utils.SignApex(
+      avbtool,
       apex_data,
       payload_key=payload_key,
       container_key=container_key,
@@ -59,7 +64,9 @@ def main(argv):
   options = {}
 
   def option_handler(o, a):
-    if o == '--container_key':
+    if o == '--avbtool':
+      options['avbtool'] = a
+    elif o == '--container_key':
       # Strip the suffix if any, as common.SignFile expects no suffix.
       DEFAULT_CONTAINER_KEY_SUFFIX = '.x509.pem'
       if a.endswith(DEFAULT_CONTAINER_KEY_SUFFIX):
@@ -77,6 +84,7 @@ def main(argv):
       argv, __doc__,
       extra_opts='',
       extra_long_opts=[
+          'avbtool=',
           'container_key=',
           'payload_extra_args=',
           'payload_key=',
@@ -91,6 +99,7 @@ def main(argv):
   common.InitLogging()
 
   signed_apex = SignApexFile(
+      options.get('avbtool', 'avbtool'),
       args[0],
       options['payload_key'],
       options['container_key'],

tools/releasetools/sign_target_files_apks.py

@@ -468,6 +468,7 @@ def ProcessTargetFiles(input_tf_zip, output_tf_zip, misc_info,
             maxsize, name, payload_key))
 
         signed_apex = apex_utils.SignApex(
+            misc_info['avb_avbtool'],
             data,
             payload_key,
             container_key,
@@ -572,7 +573,8 @@ def ProcessTargetFiles(input_tf_zip, output_tf_zip, misc_info,
       # key is specified via --avb_system_other_key.
       signing_key = OPTIONS.avb_keys.get("system_other")
       if signing_key:
-        public_key = common.ExtractAvbPublicKey(signing_key)
+        public_key = common.ExtractAvbPublicKey(
+            misc_info['avb_avbtool'], signing_key)
         print("    Rewriting AVB public key of system_other in /product")
         common.ZipWrite(output_tf_zip, public_key, filename)
 

tools/releasetools/test_apex_utils.py

@@ -43,8 +43,9 @@ class ApexUtilsTest(test_utils.ReleaseToolsTestCase):
   def test_ParseApexPayloadInfo(self):
     payload_file = self._GetTestPayload()
     apex_utils.SignApexPayload(
-        payload_file, self.payload_key, 'testkey', 'SHA256_RSA2048', self.SALT)
-    payload_info = apex_utils.ParseApexPayloadInfo(payload_file)
+        'avbtool', payload_file, self.payload_key, 'testkey', 'SHA256_RSA2048',
+        self.SALT)
+    payload_info = apex_utils.ParseApexPayloadInfo('avbtool', payload_file)
     self.assertEqual('SHA256_RSA2048', payload_info['Algorithm'])
     self.assertEqual(self.SALT, payload_info['Salt'])
     self.assertEqual('testkey', payload_info['apex.key'])
@@ -53,8 +54,9 @@ class ApexUtilsTest(test_utils.ReleaseToolsTestCase):
   def test_SignApexPayload(self):
     payload_file = self._GetTestPayload()
     apex_utils.SignApexPayload(
-        payload_file, self.payload_key, 'testkey', 'SHA256_RSA2048', self.SALT)
-    apex_utils.VerifyApexPayload(payload_file, self.payload_key)
+        'avbtool', payload_file, self.payload_key, 'testkey', 'SHA256_RSA2048',
+        self.SALT)
+    apex_utils.VerifyApexPayload('avbtool', payload_file, self.payload_key)
 
   @test_utils.SkipIfExternalToolsUnavailable()
   def test_SignApexPayload_withSignerHelper(self):
@@ -64,17 +66,19 @@ class ApexUtilsTest(test_utils.ReleaseToolsTestCase):
     payload_signer_args = '--signing_helper_with_files {}'.format(
         signing_helper)
     apex_utils.SignApexPayload(
+        'avbtool',
         payload_file,
         self.payload_key,
         'testkey', 'SHA256_RSA2048', self.SALT,
         payload_signer_args)
-    apex_utils.VerifyApexPayload(payload_file, self.payload_key)
+    apex_utils.VerifyApexPayload('avbtool', payload_file, self.payload_key)
 
   @test_utils.SkipIfExternalToolsUnavailable()
   def test_SignApexPayload_invalidKey(self):
     self.assertRaises(
         apex_utils.ApexSigningError,
         apex_utils.SignApexPayload,
+        'avbtool',
         self._GetTestPayload(),
         os.path.join(self.testdata_dir, 'testkey.x509.pem'),
         'testkey',
@@ -85,10 +89,12 @@ class ApexUtilsTest(test_utils.ReleaseToolsTestCase):
   def test_VerifyApexPayload_wrongKey(self):
     payload_file = self._GetTestPayload()
     apex_utils.SignApexPayload(
-        payload_file, self.payload_key, 'testkey', 'SHA256_RSA2048', self.SALT)
-    apex_utils.VerifyApexPayload(payload_file, self.payload_key)
+        'avbtool', payload_file, self.payload_key, 'testkey', 'SHA256_RSA2048',
+        self.SALT)
+    apex_utils.VerifyApexPayload('avbtool', payload_file, self.payload_key)
     self.assertRaises(
         apex_utils.ApexSigningError,
         apex_utils.VerifyApexPayload,
+        'avbtool',
         payload_file,
         os.path.join(self.testdata_dir, 'testkey_with_passwd.key'))

tools/releasetools/test_common.py

@@ -590,8 +590,10 @@ class CommonApkUtilsTest(test_utils.ReleaseToolsTestCase):
   def test_ExtractAvbPublicKey(self):
     privkey = os.path.join(self.testdata_dir, 'testkey.key')
     pubkey = os.path.join(self.testdata_dir, 'testkey.pubkey.pem')
-    with open(common.ExtractAvbPublicKey(privkey), 'rb') as privkey_fp, \
-        open(common.ExtractAvbPublicKey(pubkey), 'rb') as pubkey_fp:
+    extracted_from_privkey = common.ExtractAvbPublicKey('avbtool', privkey)
+    extracted_from_pubkey = common.ExtractAvbPublicKey('avbtool', pubkey)
+    with open(extracted_from_privkey, 'rb') as privkey_fp, \
+        open(extracted_from_pubkey, 'rb') as pubkey_fp:
       self.assertEqual(privkey_fp.read(), pubkey_fp.read())
 
   def test_ParseCertificate(self):

tools/releasetools/test_sign_apex.py

@@ -35,6 +35,7 @@ class SignApexTest(test_utils.ReleaseToolsTestCase):
     payload_key = os.path.join(self.testdata_dir, 'testkey_RSA4096.key')
     container_key = os.path.join(self.testdata_dir, 'testkey')
     signed_foo_apex = sign_apex.SignApexFile(
+        'avbtool',
         foo_apex,
         payload_key,
         container_key)

tools/releasetools/validate_target_files.py

@@ -324,7 +324,8 @@ def ValidateVerifiedBootImages(input_tmp, info_dict, options):
 
     # avbtool verifies all the images that have descriptors listed in vbmeta.
     image = os.path.join(input_tmp, 'IMAGES', 'vbmeta.img')
-    cmd = ['avbtool', 'verify_image', '--image', image, '--key', key]
+    cmd = [info_dict['avb_avbtool'], 'verify_image', '--image', image,
+           '--key', key]
 
     # Append the args for chained partitions if any.
     for partition in common.AVB_PARTITIONS + common.AVB_VBMETA_PARTITIONS: