xref: /aosp_15_r20/tools/treble/gki/repack_gki_lib.py (revision 105f628577ac4ba0e277a494fbb614ed8c12a994) 
1"""Helper library for repacking GKI boot images."""
2import os
3import shutil
4import subprocess
5import tempfile
6
7from treble.fetcher import fetcher_lib
8
9
10def fetch_bootimg(client, out_dir, build_id, kernel_version, target):
11  """Fetches boot.img artifacts from a given build ID."""
12  fetcher_lib.fetch_artifacts(
13      client=client,
14      build_id=build_id,
15      target=target,
16      pattern=r'(gki_.*-img-.*\.zip|gki_.*-target_files-.*\.zip|otatools.zip)'
17      .format(version=kernel_version),
18      out_dir=out_dir)
19
20
21def fetch_kernel(client, out_dir, build_id, kernel_target, kernel_debug_target):
22  """Fetches kernel artifacts from a given build ID."""
23  kernel_dir = os.path.join(out_dir, 'kernel')
24  kernel_debug_dir = os.path.join(out_dir, 'kernel_debug')
25  os.makedirs(kernel_dir)
26  os.makedirs(kernel_debug_dir)
27
28  fetcher_lib.fetch_artifacts(
29      client=client,
30      build_id=build_id,
31      target=kernel_target,
32      pattern=r'(Image|Image.lz4|System\.map|abi_symbollist|vmlinux|vmlinux.symvers)',
33      out_dir=kernel_dir)
34  fetcher_lib.fetch_artifacts(
35      client=client,
36      build_id=build_id,
37      target=kernel_debug_target,
38      pattern=r'(Image|Image.lz4|System\.map|abi-generated.xml|abi-full-generated.xml|vmlinux|vymlinx.symvers)',
39      out_dir=kernel_debug_dir)
40
41  print('Compressing kernels')
42
43  def compress_kernel(kernel_path):
44    zipped_kernel_path = os.path.join(os.path.dirname(kernel_path), 'Image.gz')
45    with open(zipped_kernel_path, 'wb') as zipped_kernel:
46      cmd = [
47          'gzip',
48          '-nc',
49          kernel_path,
50      ]
51      print(' '.join(cmd))
52      subprocess.check_call(cmd, stdout=zipped_kernel)
53
54  compress_kernel(os.path.join(kernel_dir, 'Image'))
55  compress_kernel(os.path.join(kernel_debug_dir, 'Image'))
56
57  return kernel_dir, kernel_debug_dir
58
59
60def _replace_kernel(bootimg_path, kernel_path):
61  """Unpacks a boot.img, replaces the kernel, then repacks."""
62  with tempfile.TemporaryDirectory() as unpack_dir:
63    print('Unpacking bootimg %s' % bootimg_path)
64    cmd = [
65        'out/host/linux-x86/bin/unpack_bootimg',
66        '--boot_img',
67        bootimg_path,
68        '--out',
69        unpack_dir,
70        '--format',
71        'mkbootimg',
72    ]
73    print(' '.join(cmd))
74    mkbootimg_args = subprocess.check_output(cmd).decode('utf-8').split(' ')
75    print('Copying kernel %s' % kernel_path)
76    shutil.copy(kernel_path, os.path.join(unpack_dir, 'kernel'))
77    print('Repacking with mkbootimg')
78    cmd = [
79        'out/host/linux-x86/bin/mkbootimg',
80        '--output',
81        bootimg_path,
82    ] + mkbootimg_args
83    print(' '.join(cmd))
84    subprocess.check_call(cmd)
85
86
87def repack_bootimgs(bootimg_dir, kernel_dir, kernel_debug_dir):
88  """Repacks all boot images in a given dir using the provided kernels."""
89  for bootimg_path in os.listdir(bootimg_dir):
90    bootimg_path = os.path.join(bootimg_dir, bootimg_path)
91    if not bootimg_path.endswith('.img'):
92      continue
93
94    kernel_name = 'Image'
95    if '-gz' in bootimg_path:
96      kernel_name = 'Image.gz'
97    elif '-lz4' in bootimg_path:
98      kernel_name = 'Image.lz4'
99
100    kernel_path = os.path.join(kernel_dir, kernel_name)
101    if bootimg_path.endswith('-allsyms.img'):
102      kernel_path = os.path.join(kernel_debug_dir, kernel_name)
103
104    _replace_kernel(bootimg_path, kernel_path)
105
106
107def repack_img_zip(img_zip_path, kernel_dir, kernel_debug_dir, kernel_version):
108  """Repacks boot images within an img.zip archive."""
109  with tempfile.TemporaryDirectory() as unzip_dir:
110    pattern = 'boot-{}*'.format(kernel_version)
111    print('Unzipping %s to repack bootimgs' % img_zip_path)
112    cmd = [
113        'unzip',
114        '-d',
115        unzip_dir,
116        img_zip_path,
117        pattern,
118    ]
119    print(' '.join(cmd))
120    subprocess.check_call(cmd)
121    repack_bootimgs(unzip_dir, kernel_dir, kernel_debug_dir)
122    cmd = [
123        'zip',
124        img_zip_path,
125        pattern,
126    ]
127    print(' '.join(cmd))
128    subprocess.check_call(cmd, cwd=unzip_dir)
129
130
131def replace_target_files_zip_kernels(target_files_zip_path, kernel_out_dir,
132                                     kernel_version):
133  """Replaces the BOOT/kernel-* kernels within a target_files.zip archive."""
134  with tempfile.TemporaryDirectory() as unzip_dir:
135    pattern = 'BOOT/kernel-{}*'.format(kernel_version)
136    print(
137        'Unzipping %s to replace kernels in preparation for signing' %
138        target_files_zip_path,)
139    cmd = [
140        'unzip',
141        '-d',
142        unzip_dir,
143        target_files_zip_path,
144        pattern,
145    ]
146    print(' '.join(cmd))
147    subprocess.check_call(cmd)
148    for kernel in os.listdir(kernel_out_dir):
149      if kernel.startswith('kernel-{}'.format(kernel_version)):
150        print('Copying %s' % kernel)
151        shutil.copy(
152            os.path.join(kernel_out_dir, kernel),
153            os.path.join(unzip_dir, 'BOOT'))
154    cmd = [
155        'zip',
156        target_files_zip_path,
157        pattern,
158    ]
159    print(' '.join(cmd))
160    subprocess.check_call(cmd, cwd=unzip_dir)
161