• Home
  • History
  • Annotate
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 
19 import argparse
20 import os
21 import shutil
22 import stat
23 import subprocess
24 import tempfile
25 import zipfile
26 
27 
28 # The file extension of the unpacked images.
29 IMG_FILE_EXT = ".img"
30 
31 # The directory containing executable files in OTA tools zip.
32 BIN_DIR_NAME = "bin"
33 
34 
35 def 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 
43 def 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 
52 def 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 
65 def 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 
71 def 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 
120 def repack_super_image(ota_tools_dir, misc_info_path, super_img_path,
121                        part_imgs, output_path):
122   temp_dirs = []
123   temp_files = []
124 
125   try:
126     if not is_sparse_image(super_img_path):
127       raw_super_img = super_img_path
128     else:
129       print("Convert to unsparsed super image.")
130       simg2img_path = os.path.join(ota_tools_dir, BIN_DIR_NAME, "simg2img")
131       with tempfile.NamedTemporaryFile(
132           mode="wb", prefix="super", suffix=".img",
133           delete=False) as raw_super_img_file:
134         temp_files.append(raw_super_img_file.name)
135         raw_super_img = raw_super_img_file.name
136         subprocess.check_call([
137             simg2img_path, super_img_path, raw_super_img])
138 
139     print("Unpack super image.")
140     unpacked_part_imgs = dict()
141     lpunpack_path = os.path.join(ota_tools_dir, BIN_DIR_NAME, "lpunpack")
142     unpack_dir = tempfile.mkdtemp(prefix="lpunpack")
143     temp_dirs.append(unpack_dir)
144     subprocess.check_call([lpunpack_path, raw_super_img, unpack_dir])
145     for file_name in os.listdir(unpack_dir):
146       if file_name.endswith(IMG_FILE_EXT):
147         part = file_name[:-len(IMG_FILE_EXT)]
148         unpacked_part_imgs[part] = os.path.join(unpack_dir, file_name)
149 
150     print("Create temporary misc info.")
151     lpmake_path = os.path.join(ota_tools_dir, BIN_DIR_NAME, "lpmake")
152     with tempfile.NamedTemporaryFile(
153         mode="w", prefix="misc_info", suffix=".txt",
154         delete=False) as misc_info_file:
155       temp_files.append(misc_info_file.name)
156       misc_info_file_path = misc_info_file.name
157       with open(misc_info_path, "r") as misc_info:
158         part_list = rewrite_misc_info(part_imgs, unpacked_part_imgs,
159                                       lpmake_path, misc_info, misc_info_file)
160 
161     # Check that all input partitions are in the partition list.
162     parts_not_found = part_imgs.keys() - set(part_list)
163     if parts_not_found:
164       raise ValueError("Cannot find partitions in misc info: " +
165                        " ".join(parts_not_found))
166 
167     print("Build super image.")
168     build_super_image_path = os.path.join(ota_tools_dir, BIN_DIR_NAME,
169                                           "build_super_image")
170     subprocess.check_call([build_super_image_path, misc_info_file_path,
171                            output_path])
172   finally:
173     for temp_dir in temp_dirs:
174       shutil.rmtree(temp_dir, ignore_errors=True)
175     for temp_file in temp_files:
176       os.remove(temp_file)
177 
178 
179 def main():
180   parser = argparse.ArgumentParser(
181     description="This script is for test infrastructure to mix images in a "
182                 "super image.")
183 
184   parser.add_argument("--temp-dir",
185                       default=tempfile.gettempdir(),
186                       type=existing_abs_path,
187                       help="The directory where this script creates "
188                            "temporary files.")
189   parser.add_argument("--ota-tools",
190                       required=True,
191                       type=existing_abs_path,
192                       help="The path to the zip or directory containing OTA "
193                            "tools.")
194   parser.add_argument("--misc-info",
195                       required=True,
196                       type=existing_abs_path,
197                       help="The path to the misc info file.")
198   parser.add_argument("super_img",
199                       metavar="SUPER_IMG",
200                       type=existing_abs_path,
201                       help="The path to the super image to be repacked.")
202   parser.add_argument("part_imgs",
203                       metavar="PART_IMG",
204                       nargs="*",
205                       type=partition_image,
206                       help="The partition and the image that will be added "
207                            "to the super image. The format is "
208                            "PARITITON_NAME=IMAGE_PATH. PARTITION_NAME must "
209                            "not have slot suffix. If IMAGE_PATH is empty, the "
210                            "partition will be resized to 0.")
211   args = parser.parse_args()
212 
213   # Convert the args.part_imgs to a dictionary.
214   args_part_imgs = dict()
215   for part, img in args.part_imgs:
216     if part in args_part_imgs:
217       raise ValueError(part + " partition is repeated.")
218     args_part_imgs[part] = img
219 
220   if args.temp_dir:
221     tempfile.tempdir = args.temp_dir
222 
223   temp_ota_tools_dir = None
224   try:
225     if os.path.isdir(args.ota_tools):
226       ota_tools_dir = args.ota_tools
227     else:
228       print("Unzip OTA tools.")
229       temp_ota_tools_dir = tempfile.mkdtemp(prefix="ota_tools")
230       with zipfile.ZipFile(args.ota_tools) as ota_tools_zip:
231         unzip_ota_tools(ota_tools_zip, temp_ota_tools_dir)
232       ota_tools_dir = temp_ota_tools_dir
233 
234     repack_super_image(ota_tools_dir, args.misc_info, args.super_img,
235                        args_part_imgs, args.super_img)
236   finally:
237     if temp_ota_tools_dir:
238       shutil.rmtree(temp_ota_tools_dir)
239 
240 
241 if __name__ == "__main__":
242   main()
243