Add script for mixing shared system images with a super image

Bug: 289369510
Test: ./development/gsi/repack_super_image/mix_ssi_with_device_image.py \
      --device-image-files aosp_cf_arm64_phone-img.zip \
      --ssi-files ssi-target_files.zip \
      --misc-info misc_info.txt \
      --ota-tools otatools.zip \
      --output-dir output_dir
Change-Id: I0c16f2e13036a5d66e96309c268214231d56d90b

gsi/repack_super_image/mix_ssi_with_device_image.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+#
+# Copyright 2023 - The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Mix shared system images with a super image and disable vbmeta.
+
+Example:
+./development/gsi/repack_super_image/mix_ssi_with_device_image.py \
+    --device-image-files aosp_cf_arm64_phone-img.zip \
+    --ssi-files ssi-target_files.zip \
+    --misc-info misc_info.txt \
+    --ota-tools otatools.zip \
+    --output-dir ./output
+
+or
+
+ANDROID_PRODUCT_OUT=out/target/product/vsoc_arm64 \
+    ANDROID_HOST_OUT=out/host/linux-x86/ \
+    ./development/gsi/repack_super_image/mix_ssi_with_device_image.py \
+    --ssi-files out/target/product/ssi/IMAGES/ \
+    --output-dir ./output
+"""
+
+import argparse
+import os
+import shutil
+import subprocess
+import tempfile
+import zipfile
+
+import repack_super_image
+
+
+SSI_PARTITIONS = ["system", "system_ext", "product"]
+
+
+def add_arguments(parser):
+  parser.add_argument("--device-image-files",
+                      default=os.getenv("ANDROID_PRODUCT_OUT"),
+                      help="The path to the img zip or directory containing "
+                           "device images to be mixed with the given SSI. "
+                           "Caution: If this is a directory, super.img and "
+                           "vbmeta.img under this directory will be modified. "
+                           "Default: $ANDROID_PRODUCT_OUT")
+  parser.add_argument("--ssi-files", required=True,
+                      help="The path to the target_files zip or directory "
+                           "containing shared system images for mixing.")
+  parser.add_argument("--ota-tools",
+                      default=os.getenv("ANDROID_HOST_OUT"),
+                      help="The path to the device OTA tools zip or directory "
+                           "for mixing images. "
+                           "Default: $ANDROID_HOST_OUT")
+  parser.add_argument("--misc-info",
+                      help="The device misc_info.txt for mixing images. "
+                           "Default: {args.device_image_files}/misc_info.txt.")
+  parser.add_argument("--output-dir",
+                      help="The output directory for the mixed image. "
+                           "Default: {args.device_image_files}.")
+
+
+def unzip_ssi_images(ssi_target_files, output_dir):
+  """Unzip shared system images from the target files zipfile."""
+  if not os.path.exists(ssi_target_files):
+    raise FileNotFoundError(f"{ssi_target_files} does not exist.")
+  with zipfile.ZipFile(ssi_target_files) as ssi_zip:
+    for part_img in SSI_PARTITIONS:
+      try:
+        ssi_zip.extract(f"IMAGES/{part_img}.img", output_dir)
+      except KeyError:
+        pass
+  return os.path.join(output_dir, "IMAGES")
+
+
+def unzip_super_images(device_img_artifact, output_dir):
+  """Unzip super.img from the device image artifact zipfile."""
+  if not os.path.exists(device_img_artifact):
+    raise FileNotFoundError(f"{device_img_artifact} does not exist.")
+  with zipfile.ZipFile(device_img_artifact) as device_img_zip:
+    device_img_zip.extract("super.img", output_dir)
+  return os.path.join(output_dir, "super.img")
+
+
+def collect_ssi(ssi_dir):
+  """Collect system, system_ext and product images for mixing."""
+  ssi_imgs = dict()
+  for part_img in SSI_PARTITIONS:
+    img_path = os.path.join(ssi_dir, f"{part_img}.img")
+    if os.path.exists(img_path):
+      ssi_imgs[part_img] = img_path
+    else:
+      # system.img is mandatory, while other images are optional in SSI
+      if part_img == "system":
+        raise FileNotFoundError(f"{img_path} does not exist.")
+      else:
+        print(f"No {part_img}.img")
+        ssi_imgs[part_img] = ""
+  return ssi_imgs
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  add_arguments(parser)
+  args = parser.parse_args()
+
+  if not args.device_image_files:
+    raise ValueError("device image path is not set.")
+
+  output_dir = args.output_dir if args.output_dir else args.device_image_files
+  if not os.path.isdir(output_dir):
+    raise ValueError(f"output directory {output_dir} is not valid.")
+  print(f"Output directory {output_dir}")
+
+  temp_dirs = []
+  try:
+    if os.path.isdir(args.ssi_files):
+      ssi_dir = args.ssi_files
+    else:
+      temp_dir = tempfile.mkdtemp(prefix="ssi_")
+      temp_dirs.append(temp_dir)
+      ssi_dir = unzip_ssi_images(args.ssi_files, temp_dir)
+
+    device_misc_info = args.misc_info
+    if os.path.isdir(args.device_image_files):
+      device_image_dir = args.device_image_files
+      super_img = os.path.join(device_image_dir, "super.img")
+      if not device_misc_info:
+        device_misc_info = os.path.join(device_image_dir, "misc_info.txt")
+    else:
+      super_img = unzip_super_images(args.device_image_files, output_dir)
+
+    if not device_misc_info or not os.path.exists(device_misc_info):
+      raise FileNotFoundError(f"misc_info {device_misc_info} does not exist.")
+    if not os.path.exists(super_img):
+      raise FileNotFoundError(f"{super_img} does not exist.")
+
+    if not args.ota_tools or not os.path.exists(args.ota_tools):
+      raise FileNotFoundError(f"otatools {args.ota_tools} does not exist.")
+    if os.path.isdir(args.ota_tools):
+      ota_tools_dir = args.ota_tools
+    else:
+      print("Unzip OTA tools.")
+      ota_tools_dir = tempfile.mkdtemp(prefix="ota_tools")
+      temp_dirs.append(ota_tools_dir)
+      with zipfile.ZipFile(args.ota_tools) as ota_tools_zip:
+        repack_super_image.unzip_ota_tools(ota_tools_zip, ota_tools_dir)
+
+    mix_part_imgs = collect_ssi(ssi_dir)
+    output_super_img = os.path.join(output_dir, "super.img")
+    repack_super_image.repack_super_image(ota_tools_dir, device_misc_info,
+                                          super_img, mix_part_imgs,
+                                          output_super_img)
+    print(f"Created mixed super.img at {output_super_img}")
+
+    avbtool_path = os.path.join(ota_tools_dir, "bin", "avbtool")
+    vbmeta_img = os.path.join(output_dir, "vbmeta.img")
+    subprocess.check_call([avbtool_path, "make_vbmeta_image",
+                           "--flag", "2", "--output", vbmeta_img])
+    print(f"Created vbmeta.img at {vbmeta_img}")
+  finally:
+    for temp_dir in temp_dirs:
+      shutil.rmtree(temp_dir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+  main()

gsi/repack_super_image/repack_super_image.py

@@ -118,7 +118,7 @@ def rewrite_misc_info(args_part_imgs, unpacked_part_imgs, lpmake_path,
 
 
 def repack_super_image(ota_tools_dir, misc_info_path, super_img_path,
-                       part_imgs):
+                       part_imgs, output_path):
   temp_dirs = []
   temp_files = []
 
@@ -168,7 +168,7 @@ def repack_super_image(ota_tools_dir, misc_info_path, super_img_path,
     build_super_image_path = os.path.join(ota_tools_dir, BIN_DIR_NAME,
                                           "build_super_image")
     subprocess.check_call([build_super_image_path, misc_info_file_path,
-                           super_img_path])
+                           output_path])
   finally:
     for temp_dir in temp_dirs:
       shutil.rmtree(temp_dir, ignore_errors=True)
@@ -232,7 +232,7 @@ def main():
       ota_tools_dir = temp_ota_tools_dir
 
     repack_super_image(ota_tools_dir, args.misc_info, args.super_img,
-                       args_part_imgs)
+                       args_part_imgs, args.super_img)
   finally:
     if temp_ota_tools_dir:
       shutil.rmtree(temp_ota_tools_dir)