Merge "releasetools: Add common.ZipDelete()."

tools/releasetools/add_img_to_target_files.py

@@ -467,16 +467,12 @@ def AddCache(output_zip, prefix="IMAGES/"):
 
 
 def ReplaceUpdatedFiles(zip_filename, files_list):
-  """Update all the zip entries listed in the files_list.
+  """Updates all the ZIP entries listed in files_list.
 
   For now the list includes META/care_map.txt, and the related files under
   SYSTEM/ after rebuilding recovery.
   """
-
-  cmd = ["zip", "-d", zip_filename] + files_list
-  p = common.Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-  p.communicate()
-
+  common.ZipDelete(zip_filename, files_list)
   output_zip = zipfile.ZipFile(zip_filename, "a",
                                compression=zipfile.ZIP_DEFLATED,
                                allowZip64=True)

tools/releasetools/common.py

@@ -1147,6 +1147,28 @@ def ZipWriteStr(zip_file, zinfo_or_arcname, data, perms=None,
   zipfile.ZIP64_LIMIT = saved_zip64_limit
 
 
+def ZipDelete(zip_filename, entries):
+  """Deletes entries from a ZIP file.
+
+  Since deleting entries from a ZIP file is not supported, it shells out to
+  'zip -d'.
+
+  Args:
+    zip_filename: The name of the ZIP file.
+    entries: The name of the entry, or the list of names to be deleted.
+
+  Raises:
+    AssertionError: In case of non-zero return from 'zip'.
+  """
+  if isinstance(entries, basestring):
+    entries = [entries]
+  cmd = ["zip", "-d", zip_filename] + entries
+  proc = Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+  stdoutdata, _ = proc.communicate()
+  assert proc.returncode == 0, "Failed to delete %s:\n%s" % (entries,
+                                                             stdoutdata)
+
+
 def ZipClose(zip_file):
   # http://b/18015246
   # zipfile also refers to ZIP64_LIMIT during close() when it writes out the

tools/releasetools/ota_from_target_files.py

@@ -137,7 +137,6 @@ if sys.hexversion < 0x02070000:
   print("Python 2.7 or newer is required.", file=sys.stderr)
   sys.exit(1)
 
-import copy
 import multiprocessing
 import os.path
 import subprocess
@@ -1239,44 +1238,33 @@ def WriteABOTAPackageWithBrilloScript(target_file, output_file,
   common.ZipClose(output_zip)
 
   # SignOutput(), which in turn calls signapk.jar, will possibly reorder the
-  # zip entries, as well as padding the entry headers. We do a preliminary
+  # ZIP entries, as well as padding the entry headers. We do a preliminary
   # signing (with an incomplete metadata entry) to allow that to happen. Then
-  # compute the zip entry offsets, write back the final metadata and do the
+  # compute the ZIP entry offsets, write back the final metadata and do the
   # final signing.
-  prelim_signing = tempfile.NamedTemporaryFile()
-  SignOutput(temp_zip_file.name, prelim_signing.name)
+  prelim_signing = common.MakeTempFile(suffix='.zip')
+  SignOutput(temp_zip_file.name, prelim_signing)
   common.ZipClose(temp_zip_file)
 
   # Open the signed zip. Compute the final metadata that's needed for streaming.
-  prelim_zip = zipfile.ZipFile(prelim_signing, "r",
-                               compression=zipfile.ZIP_DEFLATED)
+  prelim_signing_zip = zipfile.ZipFile(prelim_signing, 'r')
   expected_length = len(metadata['ota-streaming-property-files'])
   metadata['ota-streaming-property-files'] = ComputeStreamingMetadata(
-      prelim_zip, reserve_space=False, expected_length=expected_length)
+      prelim_signing_zip, reserve_space=False, expected_length=expected_length)
+  common.ZipClose(prelim_signing_zip)
 
-  # Copy the zip entries, as we cannot update / delete entries with zipfile.
-  final_signing = tempfile.NamedTemporaryFile()
-  output_zip = zipfile.ZipFile(final_signing, "w",
+  # Replace the METADATA entry.
+  common.ZipDelete(prelim_signing, METADATA_NAME)
+  output_zip = zipfile.ZipFile(prelim_signing, 'a',
                                compression=zipfile.ZIP_DEFLATED)
-  for item in prelim_zip.infolist():
-    if item.filename == METADATA_NAME:
-      continue
-
-    data = prelim_zip.read(item.filename)
-    out_info = copy.copy(item)
-    common.ZipWriteStr(output_zip, out_info, data)
-
-  # Now write the final metadata entry.
   WriteMetadata(metadata, output_zip)
-  common.ZipClose(prelim_zip)
   common.ZipClose(output_zip)
 
   # Re-sign the package after updating the metadata entry.
-  SignOutput(final_signing.name, output_file)
-  final_signing.close()
+  SignOutput(prelim_signing, output_file)
 
   # Reopen the final signed zip to double check the streaming metadata.
-  output_zip = zipfile.ZipFile(output_file, "r")
+  output_zip = zipfile.ZipFile(output_file, 'r')
   actual = metadata['ota-streaming-property-files'].strip()
   expected = ComputeStreamingMetadata(output_zip)
   assert actual == expected, \

tools/releasetools/test_common.py

@@ -309,6 +309,51 @@ class CommonZipTest(unittest.TestCase):
     finally:
       os.remove(zip_file_name)
 
+  def test_ZipDelete(self):
+    zip_file = tempfile.NamedTemporaryFile(delete=False, suffix='.zip')
+    output_zip = zipfile.ZipFile(zip_file.name, 'w',
+                                 compression=zipfile.ZIP_DEFLATED)
+    with tempfile.NamedTemporaryFile() as entry_file:
+      entry_file.write(os.urandom(1024))
+      common.ZipWrite(output_zip, entry_file.name, arcname='Test1')
+      common.ZipWrite(output_zip, entry_file.name, arcname='Test2')
+      common.ZipWrite(output_zip, entry_file.name, arcname='Test3')
+      common.ZipClose(output_zip)
+    zip_file.close()
+
+    try:
+      common.ZipDelete(zip_file.name, 'Test2')
+      with zipfile.ZipFile(zip_file.name, 'r') as check_zip:
+        entries = check_zip.namelist()
+        self.assertTrue('Test1' in entries)
+        self.assertFalse('Test2' in entries)
+        self.assertTrue('Test3' in entries)
+
+      self.assertRaises(AssertionError, common.ZipDelete, zip_file.name,
+                        'Test2')
+      with zipfile.ZipFile(zip_file.name, 'r') as check_zip:
+        entries = check_zip.namelist()
+        self.assertTrue('Test1' in entries)
+        self.assertFalse('Test2' in entries)
+        self.assertTrue('Test3' in entries)
+
+      common.ZipDelete(zip_file.name, ['Test3'])
+      with zipfile.ZipFile(zip_file.name, 'r') as check_zip:
+        entries = check_zip.namelist()
+        self.assertTrue('Test1' in entries)
+        self.assertFalse('Test2' in entries)
+        self.assertFalse('Test3' in entries)
+
+      common.ZipDelete(zip_file.name, ['Test1', 'Test2'])
+      with zipfile.ZipFile(zip_file.name, 'r') as check_zip:
+        entries = check_zip.namelist()
+        self.assertFalse('Test1' in entries)
+        self.assertFalse('Test2' in entries)
+        self.assertFalse('Test3' in entries)
+    finally:
+      os.remove(zip_file.name)
+
+
 class InstallRecoveryScriptFormatTest(unittest.TestCase):
   """Check the format of install-recovery.sh