support use of prebuilt bootable images

img_from_target_files now, with the -z flag, will produce an output
zip with only the bootable partitions (boot and recovery).

img_ and ota_from_target_files can take, instead of a simple
"target_files.zip", a name of the form
"target_files.zip+bootable_images.zip", where the second zip contains
bootable images that should be used instead of building them from the
target_files.zip.  (This should be the zip produced by the above -z
flag, perhaps with the images messed with in some way, such as by an
unnamed OEM's extra signature wrapper for their "secure boot"
process.)

Bug: 3391371
Change-Id: Iaf96dfc8f30e806ae342dcf3241566e76ae372d4

3 files changed, 96 insertions, 69 deletions

diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index f3b9ed5..3cc86bf 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -19,7 +19,6 @@ import getpass
 import imp
 import os
 import re
-import sha
 import shutil
 import subprocess
 import sys
@@ -28,6 +27,13 @@ import threading
 import time
 import zipfile
 
+try:
+  import hashlib
+  sha1 = hashlib.sha1
+except ImportError:
+  import sha
+  sha1 = sha.sha
+
 # missing in Python 2.4 and before
 if not hasattr(os, "SEEK_SET"):
   os.SEEK_SET = 0
@@ -163,23 +169,6 @@ def DumpInfoDict(d):
   for k, v in sorted(d.items()):
     print "%-25s = (%s) %s" % (k, type(v).__name__, v)
 
-def BuildAndAddBootableImage(sourcedir, targetname, output_zip, info_dict):
-  """Take a kernel, cmdline, and ramdisk directory from the input (in
-  'sourcedir'), and turn them into a boot image.  Put the boot image
-  into the output zip file under the name 'targetname'.  Returns
-  targetname on success or None on failure (if sourcedir does not
-  appear to contain files for the requested image)."""
-
-  print "creating %s..." % (targetname,)
-
-  img = BuildBootableImage(sourcedir)
-  if img is None:
-    return None
-
-  CheckSize(img, targetname, info_dict)
-  ZipWriteStr(output_zip, targetname, img)
-  return targetname
-
 def BuildBootableImage(sourcedir):
   """Take a kernel, cmdline, and ramdisk directory from the input (in
   'sourcedir'), and turn them into a boot image.  Return the image
@@ -237,28 +226,53 @@ def BuildBootableImage(sourcedir):
   return data
 
 
-def AddRecovery(output_zip, info_dict):
-  BuildAndAddBootableImage(os.path.join(OPTIONS.input_tmp, "RECOVERY"),
-                           "recovery.img", output_zip, info_dict)
+def GetBootableImage(name, prebuilt_name, unpack_dir, tree_subdir):
+  """Return a File object (with name 'name') with the desired bootable
+  image.  Look for it in 'unpack_dir'/BOOTABLE_IMAGES under the name
+  'prebuilt_name', otherwise construct it from the source files in
+  'unpack_dir'/'tree_subdir'."""
+
+  prebuilt_path = os.path.join(unpack_dir, "BOOTABLE_IMAGES", prebuilt_name)
+  if os.path.exists(prebuilt_path):
+    print "using prebuilt %s..." % (prebuilt_name,)
+    return File.FromLocalFile(name, prebuilt_path)
+  else:
+    print "building image from target_files %s..." % (tree_subdir,)
+    return File(name, BuildBootableImage(os.path.join(unpack_dir, tree_subdir)))
 
-def AddBoot(output_zip, info_dict):
-  BuildAndAddBootableImage(os.path.join(OPTIONS.input_tmp, "BOOT"),
-                           "boot.img", output_zip, info_dict)
 
 def UnzipTemp(filename, pattern=None):
-  """Unzip the given archive into a temporary directory and return the name."""
+  """Unzip the given archive into a temporary directory and return the name.
+
+  If filename is of the form "foo.zip+bar.zip", unzip foo.zip into a
+  temp dir, then unzip bar.zip into that_dir/BOOTABLE_IMAGES.
+
+  Returns (tempdir, zipobj) where zipobj is a zipfile.ZipFile (of the
+  main file), open for reading.
+  """
 
   tmp = tempfile.mkdtemp(prefix="targetfiles-")
   OPTIONS.tempfiles.append(tmp)
-  cmd = ["unzip", "-o", "-q", filename, "-d", tmp]
-  if pattern is not None:
-    cmd.append(pattern)
-  p = Run(cmd, stdout=subprocess.PIPE)
-  p.communicate()
-  if p.returncode != 0:
-    raise ExternalError("failed to unzip input target-files \"%s\"" %
-                        (filename,))
-  return tmp
+
+  def unzip_to_dir(filename, dirname):
+    cmd = ["unzip", "-o", "-q", filename, "-d", dirname]
+    if pattern is not None:
+      cmd.append(pattern)
+    p = Run(cmd, stdout=subprocess.PIPE)
+    p.communicate()
+    if p.returncode != 0:
+      raise ExternalError("failed to unzip input target-files \"%s\"" %
+                          (filename,))
+
+  m = re.match(r"^(.*[.]zip)\+(.*[.]zip)$", filename, re.IGNORECASE)
+  if m:
+    unzip_to_dir(m.group(1), tmp)
+    unzip_to_dir(m.group(2), os.path.join(tmp, "BOOTABLE_IMAGES"))
+    filename = m.group(1)
+  else:
+    unzip_to_dir(filename, tmp)
+
+  return tmp, zipfile.ZipFile(filename, "r")
 
 
 def GetKeyPasswords(keylist):
@@ -650,7 +664,14 @@ class File(object):
     self.name = name
     self.data = data
     self.size = len(data)
-    self.sha1 = sha.sha(data).hexdigest()
+    self.sha1 = sha1(data).hexdigest()
+
+  @classmethod
+  def FromLocalFile(cls, name, diskname):
+    f = open(diskname, "rb")
+    data = f.read()
+    f.close()
+    return File(name, data)
 
   def WriteToTemp(self):
     t = tempfile.NamedTemporaryFile()
diff --git a/tools/releasetools/img_from_target_files b/tools/releasetools/img_from_target_files
index 793f7ac..08817ca 100755
--- a/tools/releasetools/img_from_target_files
+++ b/tools/releasetools/img_from_target_files
@@ -23,6 +23,10 @@ Usage:  img_from_target_files [flags] input_target_files output_image_zip
   -b  (--board_config)  <file>
       Deprecated.
 
+  -z  (--bootable_zip)
+      Include only the bootable images (eg 'boot' and 'recovery') in
+      the output.
+
 """
 
 import sys
@@ -149,35 +153,44 @@ def CopyInfo(output_zip):
 
 
 def main(argv):
+  bootable_only = [False]
 
   def option_handler(o, a):
     if o in ("-b", "--board_config"):
       pass       # deprecated
+    if o in ("-z", "--bootable_zip"):
+      bootable_only[0] = True
     else:
       return False
     return True
 
   args = common.ParseOptions(argv, __doc__,
-                             extra_opts="b:",
-                             extra_long_opts=["board_config="],
+                             extra_opts="b:z",
+                             extra_long_opts=["board_config=",
+                                              "bootable_zip"],
                              extra_option_handler=option_handler)
 
+  bootable_only = bootable_only[0]
+
   if len(args) != 2:
     common.Usage(__doc__)
     sys.exit(1)
 
-  OPTIONS.input_tmp = common.UnzipTemp(args[0])
-
-  input_zip = zipfile.ZipFile(args[0], "r")
+  OPTIONS.input_tmp, input_zip = common.UnzipTemp(args[0])
   OPTIONS.info_dict = common.LoadInfoDict(input_zip)
 
   output_zip = zipfile.ZipFile(args[1], "w", compression=zipfile.ZIP_DEFLATED)
 
-  common.AddBoot(output_zip, OPTIONS.info_dict)
-  common.AddRecovery(output_zip, OPTIONS.info_dict)
-  AddSystem(output_zip)
-  AddUserdata(output_zip)
-  CopyInfo(output_zip)
+  common.GetBootableImage(
+      "boot.img", "boot.img", OPTIONS.input_tmp, "BOOT").AddToZip(output_zip)
+  common.GetBootableImage(
+      "recovery.img", "recovery.img", OPTIONS.input_tmp,
+      "RECOVERY").AddToZip(output_zip)
+
+  if not bootable_only:
+    AddSystem(output_zip)
+    AddUserdata(output_zip)
+    CopyInfo(output_zip)
 
   print "cleaning up..."
   output_zip.close()
diff --git a/tools/releasetools/ota_from_target_files b/tools/releasetools/ota_from_target_files
index 9aeee27..79c5465 100755
--- a/tools/releasetools/ota_from_target_files
+++ b/tools/releasetools/ota_from_target_files
@@ -58,7 +58,6 @@ import copy
 import errno
 import os
 import re
-import sha
 import subprocess
 import tempfile
 import time
@@ -279,7 +278,7 @@ def CopySystemFiles(input_zip, output_zip=None,
             data = input_zip.read(info.filename)
           if info.filename.startswith("SYSTEM/lib/") and IsRegular(info):
             retouch_files.append(("/system/" + basefilename,
-                                  sha.sha(data).hexdigest()))
+                                  common.sha1(data).hexdigest()))
           output_zip.writestr(info2, data)
         if fn.endswith("/"):
           Item.Get(fn[:-1], dir=True)
@@ -331,7 +330,7 @@ def MakeRecoveryPatch(output_zip, recovery_img, boot_img):
   # we check to see if this recovery has already been installed by
   # testing just the first 2k.
   HEADER_SIZE = 2048
-  header_sha1 = sha.sha(recovery_img.data[:HEADER_SIZE]).hexdigest()
+  header_sha1 = common.sha1(recovery_img.data[:HEADER_SIZE]).hexdigest()
   sh = """#!/system/bin/sh
 if ! applypatch -c %(recovery_type)s:%(recovery_device)s:%(header_size)d:%(header_sha1)s; then
   log -t recovery "Installing new recovery image"
@@ -398,10 +397,10 @@ def WriteFullOTAPackage(input_zip, output_zip):
   else:
     script.UndoRetouchBinaries(retouch_files)
 
-  boot_img = common.File("boot.img", common.BuildBootableImage(
-      os.path.join(OPTIONS.input_tmp, "BOOT")))
-  recovery_img = common.File("recovery.img", common.BuildBootableImage(
-      os.path.join(OPTIONS.input_tmp, "RECOVERY")))
+  boot_img = common.GetBootableImage("boot.img", "boot.img",
+                                     OPTIONS.input_tmp, "BOOT")
+  recovery_img = common.GetBootableImage("recovery.img", "recovery.img",
+                                         OPTIONS.input_tmp, "RECOVERY")
   MakeRecoveryPatch(output_zip, recovery_img, boot_img)
 
   Item.GetMetadata(input_zip)
@@ -523,7 +522,7 @@ def WriteIncrementalOTAPackage(target_zip, source_zip, output_zip):
       verbatim_targets.append((tf.name, tf.size))
     else:
       common.ZipWriteStr(output_zip, "patch/" + tf.name + ".p", d)
-      patch_list.append((tf.name, tf, sf, tf.size, sha.sha(d).hexdigest()))
+      patch_list.append((tf.name, tf, sf, tf.size, common.sha1(d).hexdigest()))
       largest_source_size = max(largest_source_size, sf.size)
 
   source_fp = GetBuildProp("ro.build.fingerprint", source_zip)
@@ -534,20 +533,16 @@ def WriteIncrementalOTAPackage(target_zip, source_zip, output_zip):
   script.Mount("/system")
   script.AssertSomeFingerprint(source_fp, target_fp)
 
-  source_boot = common.File("/tmp/boot.img",
-                            common.BuildBootableImage(
-                                os.path.join(OPTIONS.source_tmp, "BOOT")))
-  target_boot = common.File("/tmp/boot.img",
-                            common.BuildBootableImage(
-                                os.path.join(OPTIONS.target_tmp, "BOOT")))
+  source_boot = common.GetBootableImage(
+      "/tmp/boot.img", "boot.img", OPTIONS.source_tmp, "BOOT")
+  target_boot = common.GetBootableImage(
+      "/tmp/boot.img", "boot.img", OPTIONS.target_tmp, "BOOT")
   updating_boot = (source_boot.data != target_boot.data)
 
-  source_recovery = common.File("system/recovery.img",
-                                common.BuildBootableImage(
-                                    os.path.join(OPTIONS.source_tmp, "RECOVERY")))
-  target_recovery = common.File("system/recovery.img",
-                                common.BuildBootableImage(
-                                    os.path.join(OPTIONS.target_tmp, "RECOVERY")))
+  source_recovery = common.GetBootableImage(
+      "/tmp/recovery.img", "recovery.img", OPTIONS.source_tmp, "RECOVERY")
+  target_recovery = common.GetBootableImage(
+      "/tmp/recovery.img", "recovery.img", OPTIONS.target_tmp, "RECOVERY")
   updating_recovery = (source_recovery.data != target_recovery.data)
 
   # Here's how we divide up the progress bar:
@@ -766,10 +761,9 @@ def main(argv):
     OPTIONS.extra_script = open(OPTIONS.extra_script).read()
 
   print "unzipping target target-files..."
-  OPTIONS.input_tmp = common.UnzipTemp(args[0])
+  OPTIONS.input_tmp, input_zip = common.UnzipTemp(args[0])
 
   OPTIONS.target_tmp = OPTIONS.input_tmp
-  input_zip = zipfile.ZipFile(args[0], "r")
   OPTIONS.info_dict = common.LoadInfoDict(input_zip)
   if OPTIONS.verbose:
     print "--- target info ---"
@@ -793,8 +787,7 @@ def main(argv):
     WriteFullOTAPackage(input_zip, output_zip)
   else:
     print "unzipping source target-files..."
-    OPTIONS.source_tmp = common.UnzipTemp(OPTIONS.incremental_source)
-    source_zip = zipfile.ZipFile(OPTIONS.incremental_source, "r")
+    OPTIONS.source_tmp, source_zip = common.UnzipTemp(OPTIONS.incremental_source)
     OPTIONS.target_info_dict = OPTIONS.info_dict
     OPTIONS.source_info_dict = common.LoadInfoDict(source_zip)
     if OPTIONS.verbose: