1#!/usr/bin/env python3
2#
3# Copyright 2020 - The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""This script is for test infrastructure to mix images in a super image."""
18
19import argparse
20import os
21import shutil
22import stat
23import subprocess
24import tempfile
25import zipfile
26
27
28# The file extension of the unpacked images.
29IMG_FILE_EXT = ".img"
30
31# The directory containing executable files in OTA tools zip.
32BIN_DIR_NAME = "bin"
33
34
35def existing_abs_path(path):
36  """Validates that a path exists and returns the absolute path."""
37  abs_path = os.path.abspath(path)
38  if not os.path.exists(abs_path):
39    raise ValueError(path + " does not exist.")
40  return abs_path
41
42
43def partition_image(part_img):
44  """Splits a string into a pair of strings by "="."""
45  part, sep, img = part_img.partition("=")
46  if not part or not sep:
47    raise ValueError(part_img + " is not in the format of "
48                     "PARITITON_NAME=IMAGE_PATH.")
49  return part, (existing_abs_path(img) if img else "")
50
51
52def unzip_ota_tools(ota_tools_zip, output_dir):
53  """Unzips OTA tools and sets the files in bin/ to be executable."""
54  ota_tools_zip.extractall(output_dir)
55  permissions = (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH |
56                 stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)
57  for root_dir, dir_names, file_names in os.walk(
58      os.path.join(output_dir, BIN_DIR_NAME)):
59    for file_name in file_names:
60      file_path = os.path.join(root_dir, file_name)
61      file_stat = os.stat(file_path)
62      os.chmod(file_path, file_stat.st_mode | permissions)
63
64
65def is_sparse_image(image_path):
66  """Checks whether a file is a sparse image."""
67  with open(image_path, "rb") as image_file:
68    return image_file.read(4) == b"\x3a\xff\x26\xed"
69
70
71def rewrite_misc_info(args_part_imgs, unpacked_part_imgs, lpmake_path,
72                      input_file, output_file):
73  """Changes the lpmake path and image paths in a misc info file.
74
75  Args:
76    args_part_imgs: A dict of {partition_name: image_path} that the user
77                    intends to substitute. The partition_names must not have
78                    slot suffixes.
79    unpacked_part_imgs: A dict of {partition_name: image_path} unpacked from
80                        the input super image. The partition_names must have
81                        slot suffixes if the misc info enables virtual_ab.
82    lpmake_path: The path to the lpmake binary.
83    input_file: The input misc info file object.
84    output_file: The output misc info file object.
85
86  Returns:
87    The list of the partition names without slot suffixes.
88  """
89  virtual_ab = False
90  partition_names = ()
91  for line in input_file:
92    split_line = line.strip().split("=", 1)
93    if len(split_line) < 2:
94      split_line = (split_line[0], "")
95    if split_line[0] == "dynamic_partition_list":
96      partition_names = split_line[1].split()
97    elif split_line[0] == "lpmake":
98      output_file.write("lpmake=%s\n" % lpmake_path)
99      continue
100    elif split_line[0].endswith("_image"):
101      continue
102    elif split_line[0] == "virtual_ab" and split_line[1] == "true":
103      virtual_ab = True
104    output_file.write(line)
105
106  for partition_name in partition_names:
107    img_path = args_part_imgs.get(partition_name)
108    if img_path is None:
109      # _a is the active slot for the super images built from source.
110      img_path = unpacked_part_imgs.get(partition_name + "_a" if virtual_ab
111                                        else partition_name)
112    if img_path is None:
113      raise KeyError("No image for " + partition_name + " partition.")
114    if img_path != "":
115      output_file.write("%s_image=%s\n" % (partition_name, img_path))
116
117  return partition_names
118
119
120def main():
121  parser = argparse.ArgumentParser(
122    description="This script is for test infrastructure to mix images in a "
123                "super image.")
124
125  parser.add_argument("--temp-dir",
126                      default=tempfile.gettempdir(),
127                      type=existing_abs_path,
128                      help="The directory where this script creates "
129                           "temporary files.")
130  parser.add_argument("--ota-tools",
131                      required=True,
132                      type=existing_abs_path,
133                      help="The path to the zip or directory containing OTA "
134                           "tools.")
135  parser.add_argument("--misc-info",
136                      required=True,
137                      type=existing_abs_path,
138                      help="The path to the misc info file.")
139  parser.add_argument("super_img",
140                      metavar="SUPER_IMG",
141                      type=existing_abs_path,
142                      help="The path to the super image to be repacked.")
143  parser.add_argument("part_imgs",
144                      metavar="PART_IMG",
145                      nargs="*",
146                      type=partition_image,
147                      help="The partition and the image that will be added "
148                           "to the super image. The format is "
149                           "PARITITON_NAME=IMAGE_PATH. PARTITION_NAME must "
150                           "not have slot suffix. If IMAGE_PATH is empty, the "
151                           "partition will be resized to 0.")
152  args = parser.parse_args()
153
154  # Convert the args.part_imgs to a dictionary.
155  args_part_imgs = dict()
156  for part, img in args.part_imgs:
157    if part in args_part_imgs:
158      raise ValueError(part + " partition is repeated.")
159    args_part_imgs[part] = img
160
161  if args.temp_dir:
162    tempfile.tempdir = args.temp_dir
163
164  temp_dirs = []
165  temp_files = []
166  try:
167    if os.path.isdir(args.ota_tools):
168      ota_tools_dir = args.ota_tools
169    else:
170      print("Unzip OTA tools.")
171      temp_ota_tools_dir = tempfile.mkdtemp(prefix="ota_tools")
172      temp_dirs.append(temp_ota_tools_dir)
173      with zipfile.ZipFile(args.ota_tools) as ota_tools_zip:
174        unzip_ota_tools(ota_tools_zip, temp_ota_tools_dir)
175      ota_tools_dir = temp_ota_tools_dir
176
177    if not is_sparse_image(args.super_img):
178      super_img_path = args.super_img
179    else:
180      print("Convert to unsparsed super image.")
181      simg2img_path = os.path.join(ota_tools_dir, BIN_DIR_NAME, "simg2img")
182      with tempfile.NamedTemporaryFile(
183          mode="wb", prefix="super", suffix=".img",
184          delete=False) as raw_super_img:
185        temp_files.append(raw_super_img.name)
186        super_img_path = raw_super_img.name
187        subprocess.check_call([simg2img_path, args.super_img, super_img_path])
188
189    print("Unpack super image.")
190    unpacked_part_imgs = dict()
191    lpunpack_path = os.path.join(ota_tools_dir, BIN_DIR_NAME, "lpunpack")
192    unpack_dir = tempfile.mkdtemp(prefix="lpunpack")
193    temp_dirs.append(unpack_dir)
194    subprocess.check_call([lpunpack_path, super_img_path, unpack_dir])
195    for file_name in os.listdir(unpack_dir):
196      if file_name.endswith(IMG_FILE_EXT):
197        part = file_name[:-len(IMG_FILE_EXT)]
198        unpacked_part_imgs[part] = os.path.join(unpack_dir, file_name)
199
200    print("Create temporary misc info.")
201    lpmake_path = os.path.join(ota_tools_dir, BIN_DIR_NAME, "lpmake")
202    with tempfile.NamedTemporaryFile(
203        mode="w", prefix="misc_info", suffix=".txt",
204        delete=False) as misc_info_file:
205      temp_files.append(misc_info_file.name)
206      misc_info_file_path = misc_info_file.name
207      with open(args.misc_info, "r") as misc_info:
208        part_list = rewrite_misc_info(args_part_imgs, unpacked_part_imgs,
209                                      lpmake_path, misc_info, misc_info_file)
210
211    # Check that all input partitions are in the partition list.
212    parts_not_found = args_part_imgs.keys() - set(part_list)
213    if parts_not_found:
214      raise ValueError("Cannot find partitions in misc info: " +
215                       " ".join(parts_not_found))
216
217    print("Build super image.")
218    build_super_image_path = os.path.join(ota_tools_dir, BIN_DIR_NAME,
219                                          "build_super_image")
220    subprocess.check_call([build_super_image_path, misc_info_file_path,
221                           args.super_img])
222  finally:
223    for temp_dir in temp_dirs:
224      shutil.rmtree(temp_dir, ignore_errors=True)
225    for temp_file in temp_files:
226      os.remove(temp_file)
227
228
229if __name__ == "__main__":
230  main()
231