Resign apks contained in apex

Some apex payload images contain apk files. And these apks need to be
signed during the signing processed when sign_target_files_apks is
called. To support the signing, we can extract the payload and repack
the apex file with the (de)apexer tool. Add the signing support in the
apex_util.

Bug: 146508800
Test: unit tests pass, run sign_apex, sign_target_files_apks
Change-Id: If6d58975248709a144b07dbabf47c27916e5695e
diff --git a/tools/releasetools/apex_utils.py b/tools/releasetools/apex_utils.py
index ee3c463..3d16d30 100644
--- a/tools/releasetools/apex_utils.py
+++ b/tools/releasetools/apex_utils.py
@@ -18,6 +18,7 @@
 import os.path
 import re
 import shlex
+import shutil
 import zipfile
 
 import common
@@ -41,6 +42,130 @@
     Exception.__init__(self, message)
 
 
+class ApexApkSigner(object):
+  """Class to sign the apk files in a apex payload image and repack the apex"""
+
+  def __init__(self, apex_path, key_passwords, codename_to_api_level_map):
+    self.apex_path = apex_path
+    self.key_passwords = key_passwords
+    self.codename_to_api_level_map = codename_to_api_level_map
+
+  def ProcessApexFile(self, apk_keys, payload_key, payload_public_key):
+    """Scans and signs the apk files and repack the apex
+
+    Args:
+      apk_keys: A dict that holds the signing keys for apk files.
+      payload_key: The path to the apex payload signing key.
+      payload_public_key: The path to the public key corresponding to the
+       payload signing key.
+
+    Returns:
+      The repacked apex file containing the signed apk files.
+    """
+    list_cmd = ['deapexer', 'list', self.apex_path]
+    entries_names = common.RunAndCheckOutput(list_cmd).split()
+    apk_entries = [name for name in entries_names if name.endswith('.apk')]
+
+    # No need to sign and repack, return the original apex path.
+    if not apk_entries:
+      logger.info('No apk file to sign in %s', self.apex_path)
+      return self.apex_path
+
+    for entry in apk_entries:
+      apk_name = os.path.basename(entry)
+      if apk_name not in apk_keys:
+        raise ApexSigningError('Failed to find signing keys for apk file {} in'
+                               ' apex {}.  Use "-e <apkname>=" to specify a key'
+                               .format(entry, self.apex_path))
+      if not any(dirname in entry for dirname in ['app/', 'priv-app/',
+                                                  'overlay/']):
+        logger.warning('Apk path does not contain the intended directory name:'
+                       ' %s', entry)
+
+    payload_dir, has_signed_apk = self.ExtractApexPayloadAndSignApks(
+        apk_entries, apk_keys)
+    if not has_signed_apk:
+      logger.info('No apk file has been signed in %s', self.apex_path)
+      return self.apex_path
+
+    return self.RepackApexPayload(payload_dir, payload_key, payload_public_key)
+
+  def ExtractApexPayloadAndSignApks(self, apk_entries, apk_keys):
+    """Extracts the payload image and signs the containing apk files."""
+    payload_dir = common.MakeTempDir()
+    extract_cmd = ['deapexer', 'extract', self.apex_path, payload_dir]
+    common.RunAndCheckOutput(extract_cmd)
+
+    has_signed_apk = False
+    for entry in apk_entries:
+      apk_path = os.path.join(payload_dir, entry)
+      assert os.path.exists(self.apex_path)
+
+      key_name = apk_keys.get(os.path.basename(entry))
+      if key_name in common.SPECIAL_CERT_STRINGS:
+        logger.info('Not signing: %s due to special cert string', apk_path)
+        continue
+
+      logger.info('Signing apk file %s in apex %s', apk_path, self.apex_path)
+      # Rename the unsigned apk and overwrite the original apk path with the
+      # signed apk file.
+      unsigned_apk = common.MakeTempFile()
+      os.rename(apk_path, unsigned_apk)
+      common.SignFile(unsigned_apk, apk_path, key_name, self.key_passwords,
+                      codename_to_api_level_map=self.codename_to_api_level_map)
+      has_signed_apk = True
+    return payload_dir, has_signed_apk
+
+  def RepackApexPayload(self, payload_dir, payload_key, payload_public_key):
+    """Rebuilds the apex file with the updated payload directory."""
+    apex_dir = common.MakeTempDir()
+    # Extract the apex file and reuse its meta files as repack parameters.
+    common.UnzipToDir(self.apex_path, apex_dir)
+
+    android_jar_path = common.OPTIONS.android_jar_path
+    if not android_jar_path:
+      android_jar_path = os.path.join(os.environ.get(
+          'ANDROID_BUILD_TOP'), 'prebuilts/sdk/current/public/android.jar')
+      logger.warning('android_jar_path not found in options, falling back to'
+                     ' use %s', android_jar_path)
+
+    arguments_dict = {
+        'manifest': os.path.join(apex_dir, 'apex_manifest.pb'),
+        'build_info': os.path.join(apex_dir, 'apex_build_info.pb'),
+        'assets_dir': os.path.join(apex_dir, 'assets'),
+        'android_jar_path': android_jar_path,
+        'key': payload_key,
+        'pubkey': payload_public_key,
+    }
+    for filename in arguments_dict.values():
+      assert os.path.exists(filename), 'file {} not found'.format(filename)
+
+    # The repack process will add back these files later in the payload image.
+    for name in ['apex_manifest.pb', 'apex_manifest.json', 'lost+found']:
+      path = os.path.join(payload_dir, name)
+      if os.path.isfile(path):
+        os.remove(path)
+      elif os.path.isdir(path):
+        shutil.rmtree(path)
+
+    repacked_apex = common.MakeTempFile(suffix='.apex')
+    repack_cmd = ['apexer', '--force', '--include_build_info',
+                  '--do_not_check_keyname', '--apexer_tool_path',
+                  os.getenv('PATH')]
+    for key, val in arguments_dict.items():
+      repack_cmd.append('--' + key)
+      repack_cmd.append(val)
+    manifest_json = os.path.join(apex_dir, 'apex_manifest.json')
+    if os.path.exists(manifest_json):
+      repack_cmd.append('--manifest_json')
+      repack_cmd.append(manifest_json)
+    repack_cmd.append(payload_dir)
+    repack_cmd.append(repacked_apex)
+    common.RunAndCheckOutput(repack_cmd)
+
+    return repacked_apex
+
+
 def SignApexPayload(avbtool, payload_file, payload_key_path, payload_key_name,
                     algorithm, salt, no_hashtree, signing_args=None):
   """Signs a given payload_file with the payload key."""
@@ -155,7 +280,8 @@
 
 
 def SignApex(avbtool, apex_data, payload_key, container_key, container_pw,
-             codename_to_api_level_map, no_hashtree, signing_args=None):
+             apk_keys, codename_to_api_level_map,
+             no_hashtree, signing_args=None):
   """Signs the current APEX with the given payload/container keys.
 
   Args:
@@ -163,6 +289,7 @@
     payload_key: The path to payload signing key (w/ extension).
     container_key: The path to container signing key (w/o extension).
     container_pw: The matching password of the container_key, or None.
+    apk_keys: A dict that holds the signing keys for apk files.
     codename_to_api_level_map: A dict that maps from codename to API level.
     no_hashtree: Don't include hashtree in the signed APEX.
     signing_args: Additional args to be passed to the payload signer.
@@ -177,7 +304,15 @@
   APEX_PAYLOAD_IMAGE = 'apex_payload.img'
   APEX_PUBKEY = 'apex_pubkey'
 
-  # 1a. Extract and sign the APEX_PAYLOAD_IMAGE entry with the given
+  # 1. Extract the apex payload image and sign the containing apk files. Repack
+  # the apex file after signing.
+  payload_public_key = common.ExtractAvbPublicKey(avbtool, payload_key)
+  apk_signer = ApexApkSigner(apex_file, container_pw,
+                             codename_to_api_level_map)
+  apex_file = apk_signer.ProcessApexFile(apk_keys, payload_key,
+                                         payload_public_key)
+
+  # 2a. Extract and sign the APEX_PAYLOAD_IMAGE entry with the given
   # payload_key.
   payload_dir = common.MakeTempDir(prefix='apex-payload-')
   with zipfile.ZipFile(apex_file) as apex_fd:
@@ -195,8 +330,7 @@
       no_hashtree,
       signing_args)
 
-  # 1b. Update the embedded payload public key.
-  payload_public_key = common.ExtractAvbPublicKey(avbtool, payload_key)
+  # 2b. Update the embedded payload public key.
 
   common.ZipDelete(apex_file, APEX_PAYLOAD_IMAGE)
   if APEX_PUBKEY in zip_items:
@@ -206,11 +340,11 @@
   common.ZipWrite(apex_zip, payload_public_key, arcname=APEX_PUBKEY)
   common.ZipClose(apex_zip)
 
-  # 2. Align the files at page boundary (same as in apexer).
+  # 3. Align the files at page boundary (same as in apexer).
   aligned_apex = common.MakeTempFile(prefix='apex-container-', suffix='.apex')
   common.RunAndCheckOutput(['zipalign', '-f', '4096', apex_file, aligned_apex])
 
-  # 3. Sign the APEX container with container_key.
+  # 4. Sign the APEX container with container_key.
   signed_apex = common.MakeTempFile(prefix='apex-container-', suffix='.apex')
 
   # Specify the 4K alignment when calling SignApk.
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 6a6f119..2e235ee 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -69,6 +69,7 @@
     self.extra_signapk_args = []
     self.java_path = "java"  # Use the one on the path by default.
     self.java_args = ["-Xmx2048m"]  # The default JVM args.
+    self.android_jar_path = None
     self.public_key_suffix = ".x509.pem"
     self.private_key_suffix = ".pk8"
     # use otatools built boot_signer by default
@@ -1823,7 +1824,7 @@
         argv, "hvp:s:x:" + extra_opts,
         ["help", "verbose", "path=", "signapk_path=",
          "signapk_shared_library_path=", "extra_signapk_args=",
-         "java_path=", "java_args=", "public_key_suffix=",
+         "java_path=", "java_args=", "android_jar_path=", "public_key_suffix=",
          "private_key_suffix=", "boot_signer_path=", "boot_signer_args=",
          "verity_signer_path=", "verity_signer_args=", "device_specific=",
          "extra=", "logfile=", "aftl_server=", "aftl_key_path=",
@@ -1852,6 +1853,8 @@
       OPTIONS.java_path = a
     elif o in ("--java_args",):
       OPTIONS.java_args = shlex.split(a)
+    elif o in ("--android_jar_path",):
+      OPTIONS.android_jar_path = a
     elif o in ("--public_key_suffix",):
       OPTIONS.public_key_suffix = a
     elif o in ("--private_key_suffix",):
diff --git a/tools/releasetools/sign_apex.py b/tools/releasetools/sign_apex.py
index 4c0850c..b0128dc 100755
--- a/tools/releasetools/sign_apex.py
+++ b/tools/releasetools/sign_apex.py
@@ -31,6 +31,10 @@
   --payload_extra_args <args>
       Optional flag that specifies any extra args to be passed to payload signer
       (e.g. --payload_extra_args="--signing_helper_with_files /path/to/helper").
+
+  -e  (--extra_apks)  <name,name,...=key>
+      Add extra APK name/key pairs. This is useful to sign the apk files in the
+      apex payload image.
 """
 
 import logging
@@ -43,8 +47,8 @@
 logger = logging.getLogger(__name__)
 
 
-def SignApexFile(avbtool, apex_file, payload_key, container_key,
-                 no_hashtree, signing_args=None):
+def SignApexFile(avbtool, apex_file, payload_key, container_key, no_hashtree,
+                 apk_keys=None, signing_args=None):
   """Signs the given apex file."""
   with open(apex_file, 'rb') as input_fp:
     apex_data = input_fp.read()
@@ -57,6 +61,7 @@
       container_pw=None,
       codename_to_api_level_map=None,
       no_hashtree=no_hashtree,
+      apk_keys=apk_keys,
       signing_args=signing_args)
 
 
@@ -77,18 +82,26 @@
       options['payload_key'] = a
     elif o == '--payload_extra_args':
       options['payload_extra_args'] = a
+    elif o in ("-e", "--extra_apks"):
+      names, key = a.split("=")
+      names = names.split(",")
+      for n in names:
+        if 'extra_apks' not in options:
+          options['extra_apks'] = {}
+        options['extra_apks'].update({n: key})
     else:
       return False
     return True
 
   args = common.ParseOptions(
       argv, __doc__,
-      extra_opts='',
+      extra_opts='e:',
       extra_long_opts=[
           'avbtool=',
           'container_key=',
           'payload_extra_args=',
           'payload_key=',
+          'extra_apks=',
       ],
       extra_option_handler=option_handler)
 
@@ -105,6 +118,7 @@
       options['payload_key'],
       options['container_key'],
       no_hashtree=False,
+      apk_keys=options.get('extra_apks', {}),
       signing_args=options.get('payload_extra_args'))
   shutil.copyfile(signed_apex, args[1])
   logger.info("done.")
diff --git a/tools/releasetools/sign_target_files_apks.py b/tools/releasetools/sign_target_files_apks.py
index fffbace..cce771c 100755
--- a/tools/releasetools/sign_target_files_apks.py
+++ b/tools/releasetools/sign_target_files_apks.py
@@ -103,6 +103,9 @@
       Specify any additional args that are needed to AVB-sign the image
       (e.g. "--signing_helper /path/to/helper"). The args will be appended to
       the existing ones in info dict.
+
+  --android_jar_path <path>
+      Path to the android.jar to repack the apex file.
 """
 
 from __future__ import print_function
@@ -151,6 +154,7 @@
 OPTIONS.avb_keys = {}
 OPTIONS.avb_algorithms = {}
 OPTIONS.avb_extra_args = {}
+OPTIONS.android_jar_path = None
 
 
 AVB_FOOTER_ARGS_BY_PARTITION = {
@@ -492,6 +496,7 @@
             payload_key,
             container_key,
             key_passwords[container_key],
+            apk_keys,
             codename_to_api_level_map,
             no_hashtree=True,
             signing_args=OPTIONS.avb_extra_args.get('apex'))
@@ -1247,6 +1252,8 @@
   apex_keys_info = ReadApexKeysInfo(input_zip)
   apex_keys = GetApexKeys(apex_keys_info, apk_keys)
 
+  # TODO(xunchang) check for the apks inside the apex files, and abort early if
+  # the keys are not available.
   CheckApkAndApexKeysAvailable(
       input_zip,
       set(apk_keys.keys()) | set(apex_keys.keys()),
diff --git a/tools/releasetools/test_apex_utils.py b/tools/releasetools/test_apex_utils.py
index 5d4cc77..df61ac0 100644
--- a/tools/releasetools/test_apex_utils.py
+++ b/tools/releasetools/test_apex_utils.py
@@ -32,6 +32,8 @@
     # The default payload signing key.
     self.payload_key = os.path.join(self.testdata_dir, 'testkey.key')
 
+    common.OPTIONS.search_path = test_utils.get_search_path()
+
   @staticmethod
   def _GetTestPayload():
     payload_file = common.MakeTempFile(prefix='apex-', suffix='.img')
@@ -126,3 +128,30 @@
         payload_file,
         os.path.join(self.testdata_dir, 'testkey_with_passwd.key'),
         no_hashtree=True)
+
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_ApexApkSigner_noApkPresent(self):
+    apex_path = os.path.join(self.testdata_dir, 'foo.apex')
+    signer = apex_utils.ApexApkSigner(apex_path, None, None)
+    processed_apex = signer.ProcessApexFile({}, self.payload_key,
+                                            None)
+    self.assertEqual(apex_path, processed_apex)
+
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_ApexApkSigner_apkKeyNotPresent(self):
+    apex_path = os.path.join(self.testdata_dir, 'has_apk.apex')
+    signer = apex_utils.ApexApkSigner(apex_path, None, None)
+    self.assertRaises(apex_utils.ApexSigningError, signer.ProcessApexFile, {},
+                      self.payload_key, None)
+
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_ApexApkSigner_signApk(self):
+    apex_path = os.path.join(self.testdata_dir, 'has_apk.apex')
+    signer = apex_utils.ApexApkSigner(apex_path, None, None)
+    apk_keys = {'wifi-service-resources.apk': os.path.join(
+        self.testdata_dir, 'testkey')}
+
+    self.payload_key = os.path.join(self.testdata_dir, 'testkey_RSA4096.key')
+    payload_pubkey = common.ExtractAvbPublicKey('avbtool',
+                                                self.payload_key)
+    signer.ProcessApexFile(apk_keys, self.payload_key, payload_pubkey)
diff --git a/tools/releasetools/test_sign_apex.py b/tools/releasetools/test_sign_apex.py
index 79d1de4..82f5938 100644
--- a/tools/releasetools/test_sign_apex.py
+++ b/tools/releasetools/test_sign_apex.py
@@ -41,3 +41,19 @@
         container_key,
         False)
     self.assertTrue(os.path.exists(signed_foo_apex))
+
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_SignApexWithApk(self):
+    test_apex = os.path.join(self.testdata_dir, 'has_apk.apex')
+    payload_key = os.path.join(self.testdata_dir, 'testkey_RSA4096.key')
+    container_key = os.path.join(self.testdata_dir, 'testkey')
+    apk_keys = {'wifi-service-resources.apk': os.path.join(
+        self.testdata_dir, 'testkey')}
+    signed_test_apex = sign_apex.SignApexFile(
+        'avbtool',
+        test_apex,
+        payload_key,
+        container_key,
+        False,
+        apk_keys)
+    self.assertTrue(os.path.exists(signed_test_apex))
diff --git a/tools/releasetools/testdata/has_apk.apex b/tools/releasetools/testdata/has_apk.apex
new file mode 100644
index 0000000..12bbdf9
--- /dev/null
+++ b/tools/releasetools/testdata/has_apk.apex
Binary files differ