releasetools: Add support for compressed APKs.

Compressed APKs can be identified by a "compressed=<ext>" entry in
the apkcerts.txt file. When we encounter such an entry, we need to
decompress the file to a temporary location before we process its
certs. When we're signing, we should also recompress the package
after it's signed.

Bug: 64531948
Test: ./build/tools/releasetools/check_target_files_signatures.py
Test: ./build/tools/releasetools/sign_target_files_apks.py
Test: compared signed output before / after this change, verify that
      it's bitwise identical when no compressed APKs are present.

Change-Id: Id32e52f9c11023955330c113117daaf6b73bd8c2

tools/releasetools/check_target_files_signatures.py

@@ -235,12 +235,40 @@ class TargetFiles(object):
     self.certmap = None
 
   def LoadZipFile(self, filename):
-    d, z = common.UnzipTemp(filename, ['*.apk'])
+    # First read the APK certs file to figure out whether there are compressed
+    # APKs in the archive. If we do have compressed APKs in the archive, then we
+    # must decompress them individually before we perform any analysis.
+
+    # This is the list of wildcards of files we extract from |filename|.
+    apk_extensions = ['*.apk']
+
+    self.certmap, compressed_extension = common.ReadApkCerts(zipfile.ZipFile(filename, "r"))
+    if compressed_extension:
+      apk_extensions.append("*.apk" + compressed_extension)
+
+    d, z = common.UnzipTemp(filename, apk_extensions)
     try:
       self.apks = {}
       self.apks_by_basename = {}
       for dirpath, _, filenames in os.walk(d):
         for fn in filenames:
+          # Decompress compressed APKs before we begin processing them.
+          if compressed_extension and fn.endswith(compressed_extension):
+            # First strip the compressed extension from the file.
+            uncompressed_fn = fn[:-len(compressed_extension)]
+
+            # Decompress the compressed file to the output file.
+            common.Gunzip(os.path.join(dirpath, fn),
+                          os.path.join(dirpath, uncompressed_fn))
+
+            # Finally, delete the compressed file and use the uncompressed file
+            # for further processing. Note that the deletion is not strictly required,
+            # but is done here to ensure that we're not using too much space in
+            # the temporary directory.
+            os.remove(os.path.join(dirpath, fn))
+            fn = uncompressed_fn
+
+
           if fn.endswith(".apk"):
             fullname = os.path.join(dirpath, fn)
             displayname = fullname[len(d)+1:]
@@ -253,7 +281,6 @@ class TargetFiles(object):
     finally:
       shutil.rmtree(d)
 
-    self.certmap = common.ReadApkCerts(z)
     z.close()
 
   def CheckSharedUids(self):

tools/releasetools/common.py

@@ -18,6 +18,7 @@ import copy
 import errno
 import getopt
 import getpass
+import gzip
 import imp
 import os
 import platform
@@ -552,6 +553,13 @@ def GetBootableImage(name, prebuilt_name, unpack_dir, tree_subdir,
   return None
 
 
+def Gunzip(in_filename, out_filename):
+  """Gunzip the given gzip compressed file to a given output file.
+  """
+  with gzip.open(in_filename, "rb") as in_file, open(out_filename, "wb") as out_file:
+    shutil.copyfileobj(in_file, out_file)
+
+
 def UnzipTemp(filename, pattern=None):
   """Unzip the given archive into a temporary directory and return the name.
 
@@ -757,16 +765,26 @@ def CheckSize(data, target, info_dict):
 
 def ReadApkCerts(tf_zip):
   """Given a target_files ZipFile, parse the META/apkcerts.txt file
-  and return a {package: cert} dict."""
+  and return a tuple with the following elements: (1) a dictionary that maps
+  packages to certs (based on the "certificate" and "private_key" attributes
+  in the file. (2) A string representing the extension of compressed APKs in
+  the target files (e.g ".gz" ".bro")."""
   certmap = {}
+  compressed_extension = None
+
   for line in tf_zip.read("META/apkcerts.txt").split("\n"):
     line = line.strip()
     if not line:
       continue
-    m = re.match(r'^name="(.*)"\s+certificate="(.*)"\s+'
+    m = re.match(r'^name="(?P<NAME>.*)"\s+certificate="(?P<CERT>.*)"\s+'
-                 r'private_key="(.*)"$', line)
+                 r'private_key="(?P<PRIVKEY>.*?)"(\s+compressed="(?P<COMPRESSED>.*)")?$',
+                 line)
     if m:
-      name, cert, privkey = m.groups()
+      matches = m.groupdict()
+      cert = matches["CERT"]
+      privkey = matches["PRIVKEY"]
+      name = matches["NAME"]
+      this_compressed_extension = matches["COMPRESSED"]
       public_key_suffix_len = len(OPTIONS.public_key_suffix)
       private_key_suffix_len = len(OPTIONS.private_key_suffix)
       if cert in SPECIAL_CERT_STRINGS and not privkey:
@@ -777,7 +795,18 @@ def ReadApkCerts(tf_zip):
         certmap[name] = cert[:-public_key_suffix_len]
       else:
         raise ValueError("failed to parse line from apkcerts.txt:\n" + line)
-  return certmap
+      if this_compressed_extension:
+        # Make sure that all the values in the compression map have the same
+        # extension. We don't support multiple compression methods in the same
+        # system image.
+        if compressed_extension:
+          if this_compressed_extension != compressed_extension:
+            raise ValueError("multiple compressed extensions : %s vs %s",
+                             (compressed_extension, this_compressed_extension))
+        else:
+          compressed_extension = this_compressed_extension
+
+  return (certmap, ("." + compressed_extension) if compressed_extension else None)
 
 
 COMMON_DOCSTRING = """

tools/releasetools/sign_target_files_apks.py

@@ -100,8 +100,10 @@ import base64
 import cStringIO
 import copy
 import errno
+import gzip
 import os
 import re
+import shutil
 import stat
 import subprocess
 import tempfile
@@ -124,9 +126,7 @@ OPTIONS.avb_keys = {}
 OPTIONS.avb_algorithms = {}
 OPTIONS.avb_extra_args = {}
 
-def GetApkCerts(tf_zip):
+def GetApkCerts(certmap):
-  certmap = common.ReadApkCerts(tf_zip)
-
   # apply the key remapping to the contents of the file
   for apk, cert in certmap.iteritems():
     certmap[apk] = OPTIONS.key_map.get(cert, cert)
@@ -140,13 +140,19 @@ def GetApkCerts(tf_zip):
   return certmap
 
 
-def CheckAllApksSigned(input_tf_zip, apk_key_map):
+def CheckAllApksSigned(input_tf_zip, apk_key_map, compressed_extension):
   """Check that all the APKs we want to sign have keys specified, and
   error out if they don't."""
   unknown_apks = []
+  compressed_apk_extension = None
+  if compressed_extension:
+    compressed_apk_extension = ".apk" + compressed_extension
   for info in input_tf_zip.infolist():
-    if info.filename.endswith(".apk"):
+    if (info.filename.endswith(".apk") or
+        (compressed_apk_extension and info.filename.endswith(compressed_apk_extension))):
       name = os.path.basename(info.filename)
+      if compressed_apk_extension and name.endswith(compressed_apk_extension):
+        name = name[:-len(compressed_extension)]
       if name not in apk_key_map:
         unknown_apks.append(name)
   if unknown_apks:
@@ -157,11 +163,25 @@ def CheckAllApksSigned(input_tf_zip, apk_key_map):
     sys.exit(1)
 
 
-def SignApk(data, keyname, pw, platform_api_level, codename_to_api_level_map):
+def SignApk(data, keyname, pw, platform_api_level, codename_to_api_level_map,
+            is_compressed):
   unsigned = tempfile.NamedTemporaryFile()
   unsigned.write(data)
   unsigned.flush()
 
+  if is_compressed:
+    uncompressed = tempfile.NamedTemporaryFile()
+    with gzip.open(unsigned.name, "rb") as in_file, open(uncompressed.name, "wb") as out_file:
+      shutil.copyfileobj(in_file, out_file)
+
+    # Finally, close the "unsigned" file (which is gzip compressed), and then
+    # replace it with the uncompressed version.
+    #
+    # TODO(narayan): All this nastiness can be avoided if python 3.2 is in use,
+    # we could just gzip / gunzip in-memory buffers instead.
+    unsigned.close()
+    unsigned = uncompressed
+
   signed = tempfile.NamedTemporaryFile()
 
   # For pre-N builds, don't upgrade to SHA-256 JAR signatures based on the APK's
@@ -186,7 +206,18 @@ def SignApk(data, keyname, pw, platform_api_level, codename_to_api_level_map):
       min_api_level=min_api_level,
       codename_to_api_level_map=codename_to_api_level_map)
 
-  data = signed.read()
+  data = None;
+  if is_compressed:
+    # Recompress the file after it has been signed.
+    compressed = tempfile.NamedTemporaryFile()
+    with open(signed.name, "rb") as in_file, gzip.open(compressed.name, "wb") as out_file:
+      shutil.copyfileobj(in_file, out_file)
+
+    data = compressed.read()
+    compressed.close()
+  else:
+    data = signed.read()
+
   unsigned.close()
   signed.close()
 
@@ -195,11 +226,17 @@ def SignApk(data, keyname, pw, platform_api_level, codename_to_api_level_map):
 
 def ProcessTargetFiles(input_tf_zip, output_tf_zip, misc_info,
                        apk_key_map, key_passwords, platform_api_level,
-                       codename_to_api_level_map):
+                       codename_to_api_level_map,
+                       compressed_extension):
+
+  compressed_apk_extension = None
+  if compressed_extension:
+    compressed_apk_extension = ".apk" + compressed_extension
 
   maxsize = max([len(os.path.basename(i.filename))
                  for i in input_tf_zip.infolist()
-                 if i.filename.endswith('.apk')])
+                 if i.filename.endswith('.apk') or
+                 (compressed_apk_extension and i.filename.endswith(compressed_apk_extension))])
   system_root_image = misc_info.get("system_root_image") == "true"
 
   for info in input_tf_zip.infolist():
@@ -210,13 +247,18 @@ def ProcessTargetFiles(input_tf_zip, output_tf_zip, misc_info,
     out_info = copy.copy(info)
 
     # Sign APKs.
-    if info.filename.endswith(".apk"):
+    if (info.filename.endswith(".apk") or
+        (compressed_apk_extension and info.filename.endswith(compressed_apk_extension))):
+      is_compressed = compressed_extension and info.filename.endswith(compressed_apk_extension)
       name = os.path.basename(info.filename)
+      if is_compressed:
+        name = name[:-len(compressed_extension)]
+
       key = apk_key_map[name]
       if key not in common.SPECIAL_CERT_STRINGS:
         print "    signing: %-*s (%s)" % (maxsize, name, key)
         signed_data = SignApk(data, key, key_passwords[key], platform_api_level,
-            codename_to_api_level_map)
+            codename_to_api_level_map, is_compressed)
         common.ZipWriteStr(output_tf_zip, out_info, signed_data)
       else:
         # an APK we're not supposed to sign.
@@ -748,8 +790,9 @@ def main(argv):
 
   BuildKeyMap(misc_info, key_mapping_options)
 
-  apk_key_map = GetApkCerts(input_zip)
+  certmap, compressed_extension = common.ReadApkCerts(input_zip)
-  CheckAllApksSigned(input_zip, apk_key_map)
+  apk_key_map = GetApkCerts(certmap)
+  CheckAllApksSigned(input_zip, apk_key_map, compressed_extension)
 
   key_passwords = common.GetKeyPasswords(set(apk_key_map.values()))
   platform_api_level, _ = GetApiLevelAndCodename(input_zip)
@@ -758,7 +801,8 @@ def main(argv):
   ProcessTargetFiles(input_zip, output_zip, misc_info,
                      apk_key_map, key_passwords,
                      platform_api_level,
-                     codename_to_api_level_map)
+                     codename_to_api_level_map,
+                     compressed_extension)
 
   common.ZipClose(input_zip)
   common.ZipClose(output_zip)