1#!/usr/bin/python3 2# 3# Copyright (C) 2024 The Android Open Source Project 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); you may not 6# use this file except in compliance with the License. You may obtain a copy of 7# 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, WITHOUT 13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 14# License for the specific language governing permissions and limitations under 15# the License. 16 17import argparse 18import dataclasses 19import os 20import shlex 21import subprocess 22import tempfile 23from typing import List 24 25from build_chd_utils import copy_files, unzip_otatools 26 27"""Builds a vendor_boot-chd_debug.img. 28 29The vendor_boot-chd_debug.img is built by adding those CHD specific debugging 30files to a Cuttlefish's vendor_boot-debug.img, using a new ramdisk fragment. 31 32Test command: 33python3 tools/treble/cuttlefish/build_chd_debug_ramdisk.py \ 34 $ANDROID_PRODUCT_OUT/vendor_boot-debug.img \ 35 -o $ANDROID_PRODUCT_OUT/vendor_boot-chd_debug.img \ 36 --otatools_zip $ANDROID_PRODUCT_OUT/otatools.zip \ 37 --add_file chd_debug.prop:adb_debug.prop 38""" 39 40# The value of ramdisk type needs to be synchronized with 41# `system/tools/mkbootimg/mkbootimg.py`. We choose `_PLATFORM` here because the 42# CHD debug ramdisk will be used in normal boot (not for _RECOVERY or _DLKM). 43_VENDOR_RAMDISK_TYPE_PLATFORM = '1' 44 45 46def _parse_args() -> argparse.Namespace: 47 """Parse the arguments for building the chd debug ramdisk. 48 49 Returns: 50 An object of the parsed arguments. 51 """ 52 parser = argparse.ArgumentParser() 53 parser.add_argument('input_img', 54 help='The input Cuttlefish vendor boot debug image.') 55 parser.add_argument('--output_img', '-o', required=True, 56 help='The output CHD vendor boot debug image.') 57 parser.add_argument('--otatools_zip', required=True, 58 help='Path to the otatools.zip.') 59 parser.add_argument('--add_file', action='append', default=[], 60 help='The file to be added to the CHD debug ramdisk. ' 61 'The format is <src path>:<dst path>.') 62 return parser.parse_args() 63 64 65@dataclasses.dataclass 66class ImageOptions: 67 """The options for building the CHD vendor boot debug image. 68 69 Attributes: 70 input_image: path of the input vendor boot debug image. 71 output_image: path of the output CHD vendor boot debug image. 72 otatools_dir: path of the otatools directory. 73 temp_dir: path of the temporary directory for ramdisk filesystem. 74 files_to_add: a list of files to be added in the debug ramdisk, where a 75 pair defines the src and dst path of each file. 76 files_to_remove: a list of files to be removed from the input vendor boot 77 debug image. 78 """ 79 input_image: str 80 output_image: str 81 otatools_dir: str 82 temp_dir: str 83 files_to_add: List[str] = dataclasses.field(default_factory=list) 84 files_to_remove: List[str] = dataclasses.field(default_factory=list) 85 86 87@dataclasses.dataclass 88class BootImage: 89 """Provide some functions to modify a boot image. 90 91 Attributes: 92 bootimg: path of the input boot image to be modified. 93 bootimg_dir: path of a temporary directory that would be used to extract 94 the input boot image. 95 unpack_bootimg_bin: path of the `unpack_bootimg` executable. 96 mkbootfs_bin: path of the `mkbootfs` executable. 97 mkbootimg_bin: path of the `mkbootimg` executable. 98 lz4_bin: path of the `lz4` executable. 99 toybox_bin: path of the `toybox` executable. 100 bootimg_args: the arguments that were used to build this boot image. 101 """ 102 bootimg: str 103 bootimg_dir: str 104 unpack_bootimg_bin: str 105 mkbootfs_bin: str 106 mkbootimg_bin: str 107 lz4_bin: str 108 toybox_bin: str 109 bootimg_args: List[str] = dataclasses.field(default_factory=list) 110 111 def _get_ramdisk_fragments(self) -> List[str]: 112 """Get the path to all ramdisk fragments at `self.bootimg_dir`.""" 113 return [os.path.join(self.bootimg_dir, file) 114 for file in os.listdir(self.bootimg_dir) 115 if file.startswith('vendor_ramdisk')] 116 117 def _compress_ramdisk(self, root_dir: str, ramdisk_file: str) -> None: 118 """Compress all the files under `root_dir` to generate `ramdisk_file`. 119 120 Args: 121 root_dir: root directory of the ramdisk content. 122 ramdisk_file: path of the output ramdisk file. 123 """ 124 mkbootfs_cmd = [self.mkbootfs_bin, root_dir] 125 mkbootfs_result = subprocess.run( 126 mkbootfs_cmd, check=True, capture_output=True) 127 compress_cmd = [self.lz4_bin, '-l', '-12', '--favor-decSpeed'] 128 with open(ramdisk_file, 'w') as o: 129 subprocess.run( 130 compress_cmd, check=True, input=mkbootfs_result.stdout, stdout=o) 131 132 def _decompress_ramdisk(self, ramdisk_file: str, output_dir: str) -> str: 133 """Decompress `ramdisk_file` to a new file at `output_dir`. 134 135 Args: 136 ramdisk_file: path of the ramdisk file to be decompressed. 137 output_dir: path of the output directory. 138 139 Returns: 140 Path of the uncompressed ramdisk. 141 """ 142 if not os.path.exists(output_dir): 143 raise FileNotFoundError(f'Decompress output {output_dir} does not exist') 144 uncompressed_ramdisk = os.path.join(output_dir, 'uncompressed_ramdisk') 145 decompress_cmd = [self.lz4_bin, '-d', ramdisk_file, uncompressed_ramdisk] 146 subprocess.run(decompress_cmd, check=True) 147 return uncompressed_ramdisk 148 149 def _extract_ramdisk(self, ramdisk_file: str, root_dir: str) -> None: 150 """Extract the files from a uncompressed ramdisk to `root_dir`. 151 152 Args: 153 ramdisk_file: path of the ramdisk file to be extracted. 154 root_dir: path of the extracted ramdisk root directory. 155 """ 156 # Use `toybox cpio` instead of `cpio` to avoid invoking cpio from the host 157 # environment. 158 extract_cmd = [self.toybox_bin, 'cpio', '-i', '-F', ramdisk_file] 159 subprocess.run(extract_cmd, cwd=root_dir, check=True) 160 161 def unpack(self) -> None: 162 """Unpack the boot.img and capture the bootimg arguments.""" 163 if self.bootimg_args: 164 raise RuntimeError(f'cannot unpack {self.bootimg} twice') 165 print(f'Unpacking {self.bootimg} to {self.bootimg_dir}') 166 unpack_cmd = [ 167 self.unpack_bootimg_bin, 168 '--boot_img', self.bootimg, 169 '--out', self.bootimg_dir, 170 '--format', 'mkbootimg' 171 ] 172 unpack_result = subprocess.run(unpack_cmd, check=True, 173 capture_output=True, encoding='utf-8') 174 self.bootimg_args = shlex.split(unpack_result.stdout) 175 176 def add_ramdisk(self, ramdisk_root: str) -> None: 177 """Add a new ramdisk fragment and update the bootimg arguments. 178 179 Args: 180 ramdisk_root: path of the root directory which contains the content of 181 the new ramdisk fragment. 182 """ 183 # Name the new ramdisk using the smallest unused index. 184 ramdisk_fragments = self._get_ramdisk_fragments() 185 new_ramdisk_name = f'vendor_ramdisk{len(ramdisk_fragments):02d}' 186 new_ramdisk_file = os.path.join(self.bootimg_dir, new_ramdisk_name) 187 if os.path.exists(new_ramdisk_file): 188 raise FileExistsError(f'{new_ramdisk_file} already exists') 189 print(f'Adding a new vendor ramdisk fragment {new_ramdisk_file}') 190 self._compress_ramdisk(ramdisk_root, new_ramdisk_file) 191 192 # Update the bootimg arguments to include the new ramdisk file. 193 self.bootimg_args.extend([ 194 '--ramdisk_type', _VENDOR_RAMDISK_TYPE_PLATFORM, 195 '--ramdisk_name', 'chd', 196 '--vendor_ramdisk_fragment', new_ramdisk_file 197 ]) 198 199 def remove_file(self, file_name: str) -> None: 200 """Remove `file_name` from all the existing ramdisk fragments. 201 202 Args: 203 file_name: path of the file to be removed, relative to the ramdisk root 204 directory. 205 206 Raises: 207 FileNotFoundError if `file_name` cannot be found in any of the ramdisk 208 fragments. 209 """ 210 ramdisk_fragments = self._get_ramdisk_fragments() 211 is_removed = False 212 for ramdisk in ramdisk_fragments: 213 print(f'Attempting to remove {file_name} from {ramdisk}') 214 with tempfile.TemporaryDirectory() as temp_dir: 215 uncompressed_ramdisk = self._decompress_ramdisk(ramdisk, temp_dir) 216 extracted_ramdisk_dir = os.path.join(temp_dir, 'extracted_ramdisk') 217 os.mkdir(extracted_ramdisk_dir) 218 self._extract_ramdisk(uncompressed_ramdisk, extracted_ramdisk_dir) 219 file_path = os.path.join(extracted_ramdisk_dir, file_name) 220 if os.path.exists(file_path): 221 os.remove(file_path) 222 is_removed = True 223 print(f'{file_path} was removed') 224 self._compress_ramdisk(extracted_ramdisk_dir, ramdisk) 225 226 if not is_removed: 227 raise FileNotFoundError( 228 f'cannot remove {file_name} from {ramdisk_fragments}' 229 ) 230 231 def pack(self, output_img: str) -> None: 232 """Pack the boot.img using `self.bootimg_args`. 233 234 Args: 235 output_img: path of the output boot image. 236 """ 237 print(f'Packing {output_img} with args: {self.bootimg_args}') 238 mkbootimg_cmd = [ 239 self.mkbootimg_bin, '--vendor_boot', output_img 240 ] + self.bootimg_args 241 subprocess.check_call(mkbootimg_cmd) 242 243 244def _prepare_env(otatools_dir: str) -> List[str]: 245 """Get the executable path of the required otatools. 246 247 We need `unpack_bootimg`, `mkbootfs`, `mkbootimg`, `lz4` and `toybox` for 248 building CHD debug ramdisk. This function returns the path to the above tools 249 in order. 250 251 Args: 252 otatools_dir: path of the otatools directory. 253 254 Raises: 255 FileNotFoundError if any required otatool does not exist. 256 """ 257 tools_path = [] 258 for tool in ['unpack_bootimg', 'mkbootfs', 'mkbootimg', 'lz4', 'toybox']: 259 tool_path = os.path.join(otatools_dir, 'bin', tool) 260 if not os.path.exists(tool_path): 261 raise FileNotFoundError(f'otatool {tool_path} does not exist') 262 tools_path.append(tool_path) 263 return tools_path 264 265 266def build_chd_debug_ramdisk(options: ImageOptions) -> None: 267 """Build a new vendor boot debug image. 268 269 1. If `options.files_to_remove` present, remove these files from all the 270 existing ramdisk fragments. 271 2. If `options.files_to_add` present, create a new ramdisk fragment which 272 adds these files, and add this new fragment into the input image. 273 274 Args: 275 options: a `ImageOptions` object which specifies the options for building 276 a CHD vendor boot debug image. 277 278 Raises: 279 FileExistsError if having duplicated ramdisk fragments. 280 FileNotFoundError if any required otatool does not exist or if 281 `options.files_to_remove` is not present in any of the ramdisk fragments 282 of `input_image`. 283 """ 284 unpack_bootimg, mkbootfs, mkbootimg, lz4, toybox = _prepare_env( 285 options.otatools_dir) 286 bootimg = BootImage( 287 bootimg=options.input_image, 288 bootimg_dir=os.path.join(options.temp_dir, 'bootimg'), 289 unpack_bootimg_bin=unpack_bootimg, 290 mkbootfs_bin=mkbootfs, 291 mkbootimg_bin=mkbootimg, 292 lz4_bin=lz4, 293 toybox_bin=toybox) 294 bootimg.unpack() 295 296 for f in options.files_to_remove: 297 bootimg.remove_file(f) 298 299 if options.files_to_add: 300 print(f'Adding {options.files_to_add} to {options.input_image}') 301 new_ramdisk_fragment = os.path.join(options.temp_dir, 302 'new_ramdisk_fragment') 303 os.mkdir(new_ramdisk_fragment) 304 copy_files(options.files_to_add, new_ramdisk_fragment) 305 bootimg.add_ramdisk(new_ramdisk_fragment) 306 307 bootimg.pack(options.output_image) 308 309 310def main(temp_dir: str) -> None: 311 args = _parse_args() 312 otatools_dir = os.path.join(temp_dir, 'otatools') 313 unzip_otatools(args.otatools_zip, otatools_dir, [ 314 'bin/unpack_bootimg', 'bin/mkbootfs', 'bin/mkbootimg', 'bin/lz4', 315 'bin/toybox', 'lib64/*' 316 ]) 317 options = ImageOptions( 318 input_image=args.input_img, 319 output_image=args.output_img, 320 otatools_dir=otatools_dir, 321 temp_dir=temp_dir, 322 files_to_add=args.add_file) 323 build_chd_debug_ramdisk(options) 324 325 326if __name__ == '__main__': 327 with tempfile.TemporaryDirectory() as temp_dir: 328 main(temp_dir) 329