Handles capex in signing script

Today, the signing script simply ignores capex files, because it
looks for hardcoded '.apex' suffix. Add support to handle capex
as well.

Bug: 190574334
Test: Sign a target file locally
Change-Id: I3085ca7b0396a4fbf1b220f7de44d4eafb60c3d8
diff --git a/tools/releasetools/sign_target_files_apks.py b/tools/releasetools/sign_target_files_apks.py
index dd2de36..2859948 100755
--- a/tools/releasetools/sign_target_files_apks.py
+++ b/tools/releasetools/sign_target_files_apks.py
@@ -217,6 +217,18 @@
     raise RuntimeError("Missing {} in AVB_FOOTER_ARGS".format(partition))
 
 
+def IsApexFile(filename):
+  return filename.endswith(".apex") or filename.endswith(".capex")
+
+
+def GetApexFilename(filename):
+  name = os.path.basename(filename)
+  # Replace the suffix for compressed apex
+  if name.endswith(".capex"):
+    return name.replace(".capex", ".apex")
+  return name
+
+
 def GetApkCerts(certmap):
   # apply the key remapping to the contents of the file
   for apk, cert in certmap.items():
@@ -356,8 +368,8 @@
   unknown_files = []
   for info in input_tf_zip.infolist():
     # Handle APEXes on all partitions
-    if info.filename.endswith('.apex'):
-      name = os.path.basename(info.filename)
+    if IsApexFile(info.filename):
+      name = GetApexFilename(info.filename)
       if name not in known_keys:
         unknown_files.append(name)
       continue
@@ -388,10 +400,11 @@
 
   invalid_apexes = []
   for info in input_tf_zip.infolist():
-    if not info.filename.endswith('.apex'):
+    if not IsApexFile(info.filename):
       continue
 
-    name = os.path.basename(info.filename)
+    name = GetApexFilename(info.filename)
+
     (payload_key, container_key) = apex_keys[name]
     if ((payload_key in common.SPECIAL_CERT_STRINGS and
          container_key not in common.SPECIAL_CERT_STRINGS) or
@@ -541,8 +554,9 @@
         common.ZipWriteStr(output_tf_zip, out_info, data)
 
     # Sign bundled APEX files on all partitions
-    elif filename.endswith(".apex"):
-      name = os.path.basename(filename)
+    elif IsApexFile(filename):
+      name = GetApexFilename(filename)
+
       payload_key, container_key = apex_keys[name]
 
       # We've asserted not having a case with only one of them PRESIGNED.