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