releasetools.py: Support flashing modem as a patch

The modem image is rather large, so support flashing
this image as a binary patchfile.

Also:

   - Refactor motoboot parsing code into a class. This
     helps clean up extraction of individual packed images.
   - Replace our Image data wrapper with common.File.

Change-Id: Ia6b2510b6b0ad1ff583a426327908e84bab50c02

1 files changed, 121 insertions, 57 deletions

diff --git a/releasetools.py b/releasetools.py
index 8b0a3079..635edcd0 100644
--- a/releasetools.py
+++ b/releasetools.py
@@ -1,11 +1,6 @@
 import common
 import struct
 
-class Image:
-  def __init__(self, name, data):
-    self.name = name
-    self.data = data
-
 class BadMagicError(Exception):
     __str__ = "bad magic value"
 
@@ -28,34 +23,48 @@ HEADER_SIZE = 1024
 SECTOR_SIZE = 512
 NUM_MAX_IMAGES = 20
 MAGIC = "MBOOTV1\0"
-def UnpackMotobootImage(data):
-  """ Unpack the data blob as a motoboot image and return the list
-  of contained image objects"""
-  num_imgs_fmt = "<L"
-  num_imgs_size = struct.calcsize(num_imgs_fmt)
-  num_imgs, = struct.unpack(num_imgs_fmt, data[:num_imgs_size])
+class MotobootImage(object):
+
+  def __init__(self, data, name = None):
+
+    self.name = name
+    self._unpack(data)
+
+  def _unpack(self, data):
+    """ Unpack the data blob as a motoboot image and return the list
+    of contained image objects"""
+    num_imgs_fmt = "<L"
+    num_imgs_size = struct.calcsize(num_imgs_fmt)
+    num_imgs, = struct.unpack(num_imgs_fmt, data[:num_imgs_size])
 
-  img_info_format = "<24sLL"
-  img_info_size = struct.calcsize(img_info_format)
+    img_info_format = "<24sLL"
+    img_info_size = struct.calcsize(img_info_format)
 
-  imgs = [ struct.unpack(img_info_format, data[num_imgs_size + i*img_info_size:num_imgs_size + (i+1)*img_info_size]) for i in range(num_imgs) ]
+    imgs = [ struct.unpack(img_info_format, data[num_imgs_size + i*img_info_size:num_imgs_size + (i+1)*img_info_size]) for i in range(num_imgs) ]
 
-  magic_format = "<8s"
-  magic_size = struct.calcsize(magic_format)
-  magic, = struct.unpack(magic_format, data[num_imgs_size + NUM_MAX_IMAGES*img_info_size:num_imgs_size + NUM_MAX_IMAGES*img_info_size + magic_size])
-  if magic != MAGIC:
-    raise BadMagicError
+    magic_format = "<8s"
+    magic_size = struct.calcsize(magic_format)
+    magic, = struct.unpack(magic_format, data[num_imgs_size + NUM_MAX_IMAGES*img_info_size:num_imgs_size + NUM_MAX_IMAGES*img_info_size + magic_size])
+    if magic != MAGIC:
+      raise BadMagicError
 
-  img_objs = []
-  for name, start, end in imgs:
-    start_offset = HEADER_SIZE + start * SECTOR_SIZE
-    end_offset = HEADER_SIZE + (end + 1) * SECTOR_SIZE - 1
-    img = Image(trunc_to_null(name), data[start_offset:end_offset+1])
-    img_objs.append(img)
+    img_objs = []
+    for name, start, end in imgs:
+      start_offset = HEADER_SIZE + start * SECTOR_SIZE
+      end_offset = HEADER_SIZE + (end + 1) * SECTOR_SIZE - 1
+      img = common.File(trunc_to_null(name), data[start_offset:end_offset+1])
+      img_objs.append(img)
 
-  imgs = img_objs
+    self.unpacked_images = img_objs
 
-  return imgs
+  def GetUnpackedImage(self, name):
+
+    found_image = None
+    for image in self.unpacked_images:
+      if image.name == name:
+          found_image = image
+          break
+    return found_image
 
 
 def FindRadio(zipfile):
@@ -64,7 +73,6 @@ def FindRadio(zipfile):
   except KeyError:
     return None
 
-
 def FullOTA_InstallEnd(info):
   try:
     bootloader_img = info.input_zip.read("RADIO/bootloader.img")
@@ -80,19 +88,20 @@ def FullOTA_InstallEnd(info):
     print "no radio.img in target_files; skipping install"
 
 def IncrementalOTA_VerifyEnd(info):
-  # HACK: We don't yet have radio image patching working,
-  #       so just return
-  return
   target_radio_img = FindRadio(info.target_zip)
   source_radio_img = FindRadio(info.source_zip)
   if not target_radio_img or not source_radio_img: return
-  if source_radio_img != target_radio_img:
-    info.script.CacheFreeSpaceCheck(len(source_radio_img))
-    radio_type, radio_device = common.GetTypeAndDevice("/radio", info.info_dict)
+  target_modem_img = MotobootImage(target_radio_img).GetUnpackedImage("modem")
+  if not target_modem_img: return
+  source_modem_img = MotobootImage(source_radio_img).GetUnpackedImage("modem")
+  if not source_modem_img: return
+  if target_modem_img.sha1 != source_modem_img.sha1:
+    info.script.CacheFreeSpaceCheck(len(source_modem_img.data))
+    radio_type, radio_device = common.GetTypeAndDevice("/modem", info.info_dict)
     info.script.PatchCheck("%s:%s:%d:%s:%d:%s" % (
         radio_type, radio_device,
-        len(source_radio_img), common.sha1(source_radio_img).hexdigest(),
-        len(target_radio_img), common.sha1(target_radio_img).hexdigest()))
+        len(source_modem_img.data), source_modem_img.sha1,
+        len(target_modem_img.data), target_modem_img.sha1))
 
 def IncrementalOTA_InstallEnd(info):
   try:
@@ -120,6 +129,7 @@ def IncrementalOTA_InstallEnd(info):
     if not sf:
       # failed to read SOURCE radio image: include the whole target
       # radio image.
+      print "no radio image in source target_files; installing complete image"
       WriteRadio(info, tf.data)
     else:
       sf = common.File("radio.img", sf)
@@ -127,45 +137,101 @@ def IncrementalOTA_InstallEnd(info):
       if tf.sha1 == sf.sha1:
         print "radio image unchanged; skipping"
       else:
-        WriteRadioPatchset(info, tf, sf)
+        WriteIncrementalRadio(info, tf, sf)
+
+def WriteIncrementalRadio(info, target_imagefile, source_imagefile):
+  try:
+    target_radio_img = MotobootImage(target_imagefile.data, "radio")
+  except BadMagicError:
+    assert False, "radio.img bad magic value"
 
-def WriteRadioPatchset(info, target_imagefile, source_imagefile):
-    # TODO: Support patching modem partition images.
-    #       For now copy the whole radio.img to the OTA package
-    WriteRadio(info, target_imagefile.data)
+  try:
+    source_radio_img = MotobootImage(source_imagefile.data, "radio")
+  except BadMagicError:
+    source_radio_img = None
+
+  write_full_modem = True
+  if source_radio_img:
+    target_modem_img = target_radio_img.GetUnpackedImage("modem")
+    if target_modem_img:
+      source_modem_img = source_radio_img.GetUnpackedImage("modem")
+      if source_modem_img:
+        WriteIncrementalModemPartition(info, target_modem_img, source_modem_img)
+        write_full_modem = False
+
+  # Write the full images, skipping modem if so directed.
+  blacklist = []
+  if not write_full_modem:
+    blacklist.append('modem')
+  WriteMotobootPartitionImages(info, target_radio_img, blacklist)
+
+def WriteIncrementalModemPartition(info, target_modem_image, source_modem_image):
+  tf = target_modem_image
+  sf = source_modem_image
+
+  patchfile_name = "radio.modem.img.p"
+
+  diff = common.Difference(tf, sf, diff_program="bsdiff")
+  common.ComputeDifferences([diff])
+  _, _, d = diff.GetPatch()
+  if d is None or len(d) > tf.size * common.OPTIONS.patch_threshold:
+    # computing difference failed, or difference is nearly as
+    # big as the target:  simply send the target.
+    WritePartitionImage(info, tf, "radio", "modem")
+  else:
+    common.ZipWriteStr(info.output_zip, patchfile_name, d)
+    info.script.Print("Patching modem image...")
+    radio_type, radio_device = common.GetTypeAndDevice(
+      "/modem", info.info_dict)
+    info.script.ApplyPatch("%s:%s:%d:%s:%d:%s" % (radio_type, radio_device,
+                           sf.size, sf.sha1, tf.size, tf.sha1),
+                           "-", tf.size, tf.sha1, sf.sha1, patchfile_name)
 
 def WriteRadio(info, radio_img):
   info.script.Print("Writing radio...")
 
   try:
-    imgs = UnpackMotobootImage(radio_img)
+    motoboot_image = MotobootImage(radio_img, "radio")
   except BadMagicError:
       assert False, "radio.img bad magic value"
 
-  WriteGroupedImages(info, "radio", imgs)
+  WriteMotobootPartitionImages(info, motoboot_image)
+
+def WriteMotobootPartitionImages(info, motoboot_image, blacklist = []):
+  WriteGroupedImages(info, motoboot_image.name, motoboot_image.unpacked_images, blacklist)
 
-def WriteGroupedImages(info, group_name, images):
+def WriteGroupedImages(info, group_name, images, blacklist = []):
   """ Write a group of partition images to the OTA package,
   and add the corresponding flash instructions to the recovery
   script.  Skip any images that do not have a corresponding
   entry in recovery.fstab."""
 
   for i in images:
-    try:
-      _, device = common.GetTypeAndDevice("/"+i.name, info.info_dict)
-    except KeyError:
-      print "skipping flash of %s; not in recovery.fstab" % (i.name,)
-      continue
-    common.ZipWriteStr(info.output_zip, "%s.%s.img" % (group_name,i.name),i.data)
+    if i.name not in blacklist:
+      WritePartitionImage(info, i, group_name)
+
+def WritePartitionImage(info, image, group_name = None):
 
-    info.script.AppendExtra('package_extract_file("%s.%s.img", "%s");' %
-                            (group_name, i.name, device))
+  filename = "%s.img" % image.name
+  if group_name:
+    filename = "%s.%s" % (group_name,filename)
+
+  try:
+    _, device = common.GetTypeAndDevice("/"+image.name, info.info_dict)
+  except KeyError:
+    print "skipping flash of %s; not in recovery.fstab" % (image.name,)
+    return
+
+  common.ZipWriteStr(info.output_zip, filename, image.data)
+
+  info.script.AppendExtra('package_extract_file("%s", "%s");' %
+                          (filename, device))
 
 def WriteBootloader(info, bootloader):
   info.script.Print("Writing bootloader...")
 
   try:
-    imgs = UnpackMotobootImage(bootloader)
+    motoboot_image = MotobootImage(bootloader,"bootloader")
   except BadMagicError:
       assert False, "bootloader.img bad magic value"
 
@@ -181,10 +247,8 @@ def WriteBootloader(info, bootloader):
 
   # OTA does not support partition changes, so
   # do not bundle the partition image in the OTA package.
-  black_list = [ 'partition' ]
-  imgs_to_write = [ k for k in imgs if k.name not in black_list ]
-
-  WriteGroupedImages(info, "bootloader", imgs_to_write)
+  blacklist = [ 'partition' ]
+  WriteMotobootPartitionImages(info, motoboot_image, blacklist)
 
   info.script.AppendExtra(
       'package_extract_file("bootloader-flag-clear.txt", "%s");' %