Sign OTA packages inside target_files during signing

Test: th
Bug: 293313353
Change-Id: Ifd5dd08153c5970dac8166808173f7dfbbb3411d
diff --git a/tools/releasetools/sign_target_files_apks.py b/tools/releasetools/sign_target_files_apks.py
index 2b45825..c595bb8 100755
--- a/tools/releasetools/sign_target_files_apks.py
+++ b/tools/releasetools/sign_target_files_apks.py
@@ -147,6 +147,34 @@
 
   --override_apex_keys <path>
       Replace all APEX keys with this private key
+
+  -k  (--package_key) <key>
+      Key to use to sign the package (default is the value of
+      default_system_dev_certificate from the input target-files's
+      META/misc_info.txt, or "build/make/target/product/security/testkey" if
+      that value is not specified).
+
+      For incremental OTAs, the default value is based on the source
+      target-file, not the target build.
+
+  --payload_signer <signer>
+      Specify the signer when signing the payload and metadata for A/B OTAs.
+      By default (i.e. without this flag), it calls 'openssl pkeyutl' to sign
+      with the package private key. If the private key cannot be accessed
+      directly, a payload signer that knows how to do that should be specified.
+      The signer will be supplied with "-inkey <path_to_key>",
+      "-in <input_file>" and "-out <output_file>" parameters.
+
+  --payload_signer_args <args>
+      Specify the arguments needed for payload signer.
+
+  --payload_signer_maximum_signature_size <signature_size>
+      The maximum signature size (in bytes) that would be generated by the given
+      payload signer. Only meaningful when custom payload signer is specified
+      via '--payload_signer'.
+      If the signer uses a RSA key, this should be the number of bytes to
+      represent the modulus. If it uses an EC key, this is the size of a
+      DER-encoded ECDSA signature.
 """
 
 from __future__ import print_function
@@ -162,7 +190,6 @@
 import re
 import shutil
 import stat
-import subprocess
 import sys
 import tempfile
 import zipfile
@@ -171,6 +198,8 @@
 import add_img_to_target_files
 import apex_utils
 import common
+import payload_signer
+from payload_signer import SignOtaPackage, PAYLOAD_BIN
 
 
 if sys.hexversion < 0x02070000:
@@ -241,6 +270,20 @@
   return filename.endswith(".apex") or filename.endswith(".capex")
 
 
+def IsOtaPackage(fp):
+  with zipfile.ZipFile(fp) as zfp:
+    if not PAYLOAD_BIN in zfp.namelist():
+      return False
+    with zfp.open(PAYLOAD_BIN, "r") as payload:
+      magic = payload.read(4)
+      return magic == b"CrAU"
+
+
+def IsEntryOtaPackage(input_zip, filename):
+  with input_zip.open(filename, "r") as fp:
+    return IsOtaPackage(fp)
+
+
 def GetApexFilename(filename):
   name = os.path.basename(filename)
   # Replace the suffix for compressed apex
@@ -515,6 +558,7 @@
   return data
 
 
+
 def IsBuildPropFile(filename):
   return filename in (
       "SYSTEM/etc/prop.default",
@@ -541,7 +585,7 @@
         filename.endswith("/prop.default")
 
 
-def ProcessTargetFiles(input_tf_zip, output_tf_zip, misc_info,
+def ProcessTargetFiles(input_tf_zip: zipfile.ZipFile, output_tf_zip, misc_info,
                        apk_keys, apex_keys, key_passwords,
                        platform_api_level, codename_to_api_level_map,
                        compressed_extension):
@@ -631,6 +675,15 @@
             "        (skipped due to special cert string)" % (name,))
         common.ZipWriteStr(output_tf_zip, out_info, data)
 
+    elif filename.endswith(".zip") and IsEntryOtaPackage(input_tf_zip, filename):
+      logger.info("Re-signing OTA package {}".format(filename))
+      with tempfile.NamedTemporaryFile() as input_ota, tempfile.NamedTemporaryFile() as output_ota:
+        with input_tf_zip.open(filename, "r") as in_fp:
+          shutil.copyfileobj(in_fp, input_ota)
+          input_ota.flush()
+        SignOtaPackage(input_ota.name, output_ota.name)
+        common.ZipWrite(output_tf_zip, output_ota.name, filename,
+                        compress_type=zipfile.ZIP_STORED)
     # System properties.
     elif IsBuildPropFile(filename):
       print("Rewriting %s:" % (filename,))
@@ -1504,7 +1557,7 @@
           "override_apk_keys=",
           "override_apex_keys=",
       ],
-      extra_option_handler=option_handler)
+      extra_option_handler=[option_handler, payload_signer.signer_options])
 
   if len(args) != 2:
     common.Usage(__doc__)
@@ -1518,6 +1571,10 @@
                                allowZip64=True)
 
   misc_info = common.LoadInfoDict(input_zip)
+  if OPTIONS.package_key is None:
+      OPTIONS.package_key = misc_info.get(
+          "default_system_dev_certificate",
+          "build/make/target/product/security/testkey")
 
   BuildKeyMap(misc_info, key_mapping_options)