1 files changed, 41 insertions, 6 deletions
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index b132963..51d64b8 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -47,7 +47,8 @@ OPTIONS.replace_verity_public_key = False
 OPTIONS.replace_verity_private_key = False
 OPTIONS.verity_signer_path = None
 
-def AddSystem(output_zip, prefix="IMAGES/", recovery_img=None, boot_img=None):
+def AddSystem(output_zip, prefix="IMAGES/", recovery_img=None, boot_img=None,
+    rebuild_recovery=OPTIONS.rebuild_recovery):
   """Turn the contents of SYSTEM into a system image and store it in
   output_zip."""
 
@@ -61,7 +62,7 @@ def AddSystem(output_zip, prefix="IMAGES/", recovery_img=None, boot_img=None):
     ofile.write(data)
     ofile.close()
 
-  if OPTIONS.rebuild_recovery:
+  if rebuild_recovery:
     print("Building new recovery patch")
     common.MakeRecoveryPatch(OPTIONS.input_tmp, output_sink, recovery_img,
                              boot_img, info_dict=OPTIONS.info_dict)
@@ -100,6 +101,29 @@ def BuildVendor(input_dir, info_dict, block_list=None):
   file containing it."""
   return CreateImage(input_dir, info_dict, "vendor", block_list=block_list)
 
+def AddOem(output_zip, prefix="IMAGES/"):
+  """Turn the contents of OEM into a oem image and store in it
+  output_zip."""
+
+  prebuilt_path = os.path.join(OPTIONS.input_tmp, prefix, "oem.img")
+  if os.path.exists(prebuilt_path):
+    print("oem.img already exists in %s, no need to rebuild..." % prefix)
+    return
+
+  block_list = common.MakeTempFile(prefix="oem-blocklist-", suffix=".map")
+  imgname = BuildOem(OPTIONS.input_tmp, OPTIONS.info_dict,
+                     block_list=block_list)
+  with open(imgname, "rb") as f:
+    common.ZipWriteStr(output_zip, prefix + "oem.img", f.read())
+  with open(block_list, "rb") as f:
+    common.ZipWriteStr(output_zip, prefix + "oem.map", f.read())
+
+
+def BuildOem(input_dir, info_dict, block_list=None):
+  """Build the (sparse) oem image and return the name of a temp
+  file containing it."""
+  return CreateImage(input_dir, info_dict, "oem", block_list=block_list)
+
 
 def CreateImage(input_dir, info_dict, what, block_list=None):
   print("creating " + what + ".img...")
@@ -299,7 +323,7 @@ def AddCache(output_zip, prefix="IMAGES/"):
   os.rmdir(temp_dir)
 
 
-def AddImagesToTargetFiles(filename):
+def AddImagesToTargetFiles(filename, rebuild_recovery=OPTIONS.rebuild_recovery):
   OPTIONS.input_tmp, input_zip = common.UnzipTemp(filename)
 
   if not OPTIONS.add_missing:
@@ -314,6 +338,12 @@ def AddImagesToTargetFiles(filename):
   except KeyError:
     has_vendor = False
 
+  try:
+    input_zip.getinfo("OEM/")
+    has_oem = True
+  except KeyError:
+    has_oem = False
+
   OPTIONS.info_dict = common.LoadInfoDict(input_zip)
   if "selinux_fc" in OPTIONS.info_dict:
     OPTIONS.info_dict["selinux_fc"] = os.path.join(
@@ -331,7 +361,7 @@ def AddImagesToTargetFiles(filename):
   boot_image = None
   if os.path.exists(prebuilt_path):
     print("boot.img already exists in IMAGES/, no need to rebuild...")
-    if OPTIONS.rebuild_recovery:
+    if rebuild_recovery:
       boot_image = common.GetBootableImage(
           "IMAGES/boot.img", "boot.img", OPTIONS.input_tmp, "BOOT")
   else:
@@ -345,7 +375,7 @@ def AddImagesToTargetFiles(filename):
   prebuilt_path = os.path.join(OPTIONS.input_tmp, "IMAGES", "recovery.img")
   if os.path.exists(prebuilt_path):
     print("recovery.img already exists in IMAGES/, no need to rebuild...")
-    if OPTIONS.rebuild_recovery:
+    if rebuild_recovery:
       recovery_image = common.GetBootableImage(
           "IMAGES/recovery.img", "recovery.img", OPTIONS.input_tmp, "RECOVERY")
   else:
@@ -355,7 +385,8 @@ def AddImagesToTargetFiles(filename):
       recovery_image.AddToZip(output_zip)
 
   banner("system")
-  AddSystem(output_zip, recovery_img=recovery_image, boot_img=boot_image)
+  AddSystem(output_zip, recovery_img=recovery_image, boot_img=boot_image,
+      rebuild_recovery = rebuild_recovery)
   if has_vendor:
     banner("vendor")
     AddVendor(output_zip)
@@ -365,6 +396,10 @@ def AddImagesToTargetFiles(filename):
   AddUserdataExtra(output_zip)
   banner("cache")
   AddCache(output_zip)
+  if has_oem:
+    banner("oem")
+    AddOem(output_zip)
+
 
   common.ZipClose(output_zip)