Sign OTA packages inside target_files during signing

Test: th
Bug: 293313353
Change-Id: Ifd5dd08153c5970dac8166808173f7dfbbb3411d

tools/releasetools/payload_signer.py

@@ -17,7 +17,12 @@
 import common
 import logging
 import shlex
+import argparse
+import tempfile
+import zipfile
+import shutil
 from common import OPTIONS, OptionHandler
+from ota_signing_utils import AddSigningArgumentParse
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +31,8 @@ OPTIONS.payload_signer_args = []
 OPTIONS.payload_signer_maximum_signature_size = None
 OPTIONS.package_key = None
 
+PAYLOAD_BIN = 'payload.bin'
+PAYLOAD_PROPERTIES_TXT = 'payload_properties.txt'
 
 class SignerOptions(OptionHandler):
 
@@ -165,3 +172,52 @@ class PayloadSigner(object):
     cmd = [self.signer] + self.signer_args + ['-in', in_file, '-out', out_file]
     common.RunAndCheckOutput(cmd)
     return out_file
+
+def GeneratePayloadProperties(payload_file):
+    properties_file = common.MakeTempFile(prefix="payload-properties-",
+                                          suffix=".txt")
+    cmd = ["delta_generator",
+           "--in_file=" + payload_file,
+           "--properties_file=" + properties_file]
+    common.RunAndCheckOutput(cmd)
+    return properties_file
+
+def SignOtaPackage(input_path, output_path):
+  payload_signer = PayloadSigner(
+      OPTIONS.package_key, OPTIONS.private_key_suffix,
+      None, OPTIONS.payload_signer, OPTIONS.payload_signer_args)
+  common.ZipExclude(input_path, output_path, [PAYLOAD_BIN, PAYLOAD_PROPERTIES_TXT])
+  with tempfile.NamedTemporaryFile() as unsigned_payload, zipfile.ZipFile(input_path, "r", allowZip64=True) as zfp:
+    with zfp.open("payload.bin") as payload_fp:
+      shutil.copyfileobj(payload_fp, unsigned_payload)
+    signed_payload = payload_signer.SignPayload(unsigned_payload.name)
+    properties_file = GeneratePayloadProperties(signed_payload)
+    with zipfile.ZipFile(output_path, "a", compression=zipfile.ZIP_STORED, allowZip64=True) as output_zfp:
+      common.ZipWrite(output_zfp, signed_payload, PAYLOAD_BIN)
+      common.ZipWrite(output_zfp, properties_file, PAYLOAD_PROPERTIES_TXT)
+
+
+def main(argv):
+  parser = argparse.ArgumentParser(
+      prog=argv[0], description="Given a series of .img files, produces a full OTA package that installs thoese images")
+  parser.add_argument("input_ota", type=str,
+                      help="Input OTA for signing")
+  parser.add_argument('output_ota', type=str,
+                      help='Output OTA for the signed package')
+  parser.add_argument("-v", action="store_true",
+                      help="Enable verbose logging", dest="verbose")
+  AddSigningArgumentParse(parser)
+  args = parser.parse_args(argv[1:])
+  input_ota = args.input_ota
+  output_ota = args.output_ota
+  if args.verbose:
+    OPTIONS.verbose = True
+  common.InitLogging()
+  if args.package_key:
+    OPTIONS.package_key = args.package_key
+  logger.info("Re-signing OTA package {}".format(input_ota))
+  SignOtaPackage(input_ota, output_ota)
+
+if __name__ == "__main__":
+  import sys
+  main(sys.argv)