scripts/make_fit.py

   1 #!/usr/bin/env python3
   2 # SPDX-License-Identifier: GPL-2.0+
   3 #
   4 # Copyright 2024 Google LLC
   5 # Written by Simon Glass <[email protected]>
   6 #
   7
   8 """Build a FIT containing a lot of devicetree files
   9
  10 Usage:
  11     make_fit.py -A arm64 -n 'Linux-6.6' -O linux
  12         -o arch/arm64/boot/image.fit -k /tmp/kern/arch/arm64/boot/image.itk
  13         @arch/arm64/boot/dts/dtbs-list -E -c gzip
  14
  15 Creates a FIT containing the supplied kernel and a set of devicetree files,
  16 either specified individually or listed in a file (with an '@' prefix).
  17
  18 Use -E to generate an external FIT (where the data is placed after the
  19 FIT data structure). This allows parsing of the data without loading
  20 the entire FIT.
  21
  22 Use -c to compress the data, using bzip2, gzip, lz4, lzma, lzo and
  23 zstd algorithms.
  24
  25 Use -D to decompose "composite" DTBs into their base components and
  26 deduplicate the resulting base DTBs and DTB overlays. This requires the
  27 DTBs to be sourced from the kernel build directory, as the implementation
  28 looks at the .cmd files produced by the kernel build.
  29
  30 The resulting FIT can be booted by bootloaders which support FIT, such
  31 as U-Boot, Linuxboot, Tianocore, etc.
  32
  33 Note that this tool does not yet support adding a ramdisk / initrd.
  34 """
  35
  36 import argparse
  37 import collections
  38 import os
  39 import subprocess
  40 import sys
  41 import tempfile
  42 import time
  43
  44 import libfdt
  45
  46
  47 # Tool extension and the name of the command-line tools
  48 CompTool = collections.namedtuple('CompTool', 'ext,tools')
  49
  50 COMP_TOOLS = {
  51     'bzip2': CompTool('.bz2', 'bzip2'),
  52     'gzip': CompTool('.gz', 'pigz,gzip'),
  53     'lz4': CompTool('.lz4', 'lz4'),
  54     'lzma': CompTool('.lzma', 'lzma'),
  55     'lzo': CompTool('.lzo', 'lzop'),
  56     'zstd': CompTool('.zstd', 'zstd'),
  57 }
  58
  59
  60 def parse_args():
  61     """Parse the program ArgumentParser
  62
  63     Returns:
  64         Namespace object containing the arguments
  65     """
  66     epilog = 'Build a FIT from a directory tree containing .dtb files'
  67     parser = argparse.ArgumentParser(epilog=epilog, fromfile_prefix_chars='@')
  68     parser.add_argument('-A', '--arch', type=str, required=True,
  69           help='Specifies the architecture')
  70     parser.add_argument('-c', '--compress', type=str, default='none',
  71           help='Specifies the compression')
  72     parser.add_argument('-D', '--decompose-dtbs', action='store_true',
  73           help='Decompose composite DTBs into base DTB and overlays')
  74     parser.add_argument('-E', '--external', action='store_true',
  75           help='Convert the FIT to use external data')
  76     parser.add_argument('-n', '--name', type=str, required=True,
  77           help='Specifies the name')
  78     parser.add_argument('-o', '--output', type=str, required=True,
  79           help='Specifies the output file (.fit)')
  80     parser.add_argument('-O', '--os', type=str, required=True,
  81           help='Specifies the operating system')
  82     parser.add_argument('-k', '--kernel', type=str, required=True,
  83           help='Specifies the (uncompressed) kernel input file (.itk)')
  84     parser.add_argument('-v', '--verbose', action='store_true',
  85                         help='Enable verbose output')
  86     parser.add_argument('dtbs', type=str, nargs='*',
  87           help='Specifies the devicetree files to process')
  88
  89     return parser.parse_args()
  90
  91
  92 def setup_fit(fsw, name):
  93     """Make a start on writing the FIT
  94
  95     Outputs the root properties and the 'images' node
  96
  97     Args:
  98         fsw (libfdt.FdtSw): Object to use for writing
  99         name (str): Name of kernel image
 100     """
 101     fsw.INC_SIZE = 65536
 102     fsw.finish_reservemap()
 103     fsw.begin_node('')
 104     fsw.property_string('description', f'{name} with devicetree set')
 105     fsw.property_u32('#address-cells', 1)
 106
 107     fsw.property_u32('timestamp', int(time.time()))
 108     fsw.begin_node('images')
 109
 110
 111 def write_kernel(fsw, data, args):
 112     """Write out the kernel image
 113
 114     Writes a kernel node along with the required properties
 115
 116     Args:
 117         fsw (libfdt.FdtSw): Object to use for writing
 118         data (bytes): Data to write (possibly compressed)
 119         args (Namespace): Contains necessary strings:
 120             arch: FIT architecture, e.g. 'arm64'
 121             fit_os: Operating Systems, e.g. 'linux'
 122             name: Name of OS, e.g. 'Linux-6.6.0-rc7'
 123             compress: Compression algorithm to use, e.g. 'gzip'
 124     """
 125     with fsw.add_node('kernel'):
 126         fsw.property_string('description', args.name)
 127         fsw.property_string('type', 'kernel_noload')
 128         fsw.property_string('arch', args.arch)
 129         fsw.property_string('os', args.os)
 130         fsw.property_string('compression', args.compress)
 131         fsw.property('data', data)
 132         fsw.property_u32('load', 0)
 133         fsw.property_u32('entry', 0)
 134
 135
 136 def finish_fit(fsw, entries):
 137     """Finish the FIT ready for use
 138
 139     Writes the /configurations node and subnodes
 140
 141     Args:
 142         fsw (libfdt.FdtSw): Object to use for writing
 143         entries (list of tuple): List of configurations:
 144             str: Description of model
 145             str: Compatible stringlist
 146     """
 147     fsw.end_node()
 148     seq = 0
 149     with fsw.add_node('configurations'):
 150         for model, compat, files in entries:
 151             seq += 1
 152             with fsw.add_node(f'conf-{seq}'):
 153                 fsw.property('compatible', bytes(compat))
 154                 fsw.property_string('description', model)
 155                 fsw.property('fdt', bytes(''.join(f'fdt-{x}\x00' for x in files), "ascii"))
 156                 fsw.property_string('kernel', 'kernel')
 157     fsw.end_node()
 158
 159
 160 def compress_data(inf, compress):
 161     """Compress data using a selected algorithm
 162
 163     Args:
 164         inf (IOBase): Filename containing the data to compress
 165         compress (str): Compression algorithm, e.g. 'gzip'
 166
 167     Return:
 168         bytes: Compressed data
 169     """
 170     if compress == 'none':
 171         return inf.read()
 172
 173     comp = COMP_TOOLS.get(compress)
 174     if not comp:
 175         raise ValueError(f"Unknown compression algorithm '{compress}'")
 176
 177     with tempfile.NamedTemporaryFile() as comp_fname:
 178         with open(comp_fname.name, 'wb') as outf:
 179             done = False
 180             for tool in comp.tools.split(','):
 181                 try:
 182                     subprocess.call([tool, '-c'], stdin=inf, stdout=outf)
 183                     done = True
 184                     break
 185                 except FileNotFoundError:
 186                     pass
 187             if not done:
 188                 raise ValueError(f'Missing tool(s): {comp.tools}\n')
 189             with open(comp_fname.name, 'rb') as compf:
 190                 comp_data = compf.read()
 191     return comp_data
 192
 193
 194 def output_dtb(fsw, seq, fname, arch, compress):
 195     """Write out a single devicetree to the FIT
 196
 197     Args:
 198         fsw (libfdt.FdtSw): Object to use for writing
 199         seq (int): Sequence number (1 for first)
 200         fname (str): Filename containing the DTB
 201         arch: FIT architecture, e.g. 'arm64'
 202         compress (str): Compressed algorithm, e.g. 'gzip'
 203     """
 204     with fsw.add_node(f'fdt-{seq}'):
 205         fsw.property_string('description', os.path.basename(fname))
 206         fsw.property_string('type', 'flat_dt')
 207         fsw.property_string('arch', arch)
 208         fsw.property_string('compression', compress)
 209
 210         with open(fname, 'rb') as inf:
 211             compressed = compress_data(inf, compress)
 212         fsw.property('data', compressed)
 213
 214
 215 def process_dtb(fname, args):
 216     """Process an input DTB, decomposing it if requested and is possible
 217
 218     Args:
 219         fname (str): Filename containing the DTB
 220         args (Namespace): Program arguments
 221     Returns:
 222         tuple:
 223             str: Model name string
 224             str: Root compatible string
 225             files: list of filenames corresponding to the DTB
 226     """
 227     # Get the compatible / model information
 228     with open(fname, 'rb') as inf:
 229         data = inf.read()
 230     fdt = libfdt.FdtRo(data)
 231     model = fdt.getprop(0, 'model').as_str()
 232     compat = fdt.getprop(0, 'compatible')
 233
 234     if args.decompose_dtbs:
 235         # Check if the DTB needs to be decomposed
 236         path, basename = os.path.split(fname)
 237         cmd_fname = os.path.join(path, f'.{basename}.cmd')
 238         with open(cmd_fname, 'r', encoding='ascii') as inf:
 239             cmd = inf.read()
 240
 241         if 'scripts/dtc/fdtoverlay' in cmd:
 242             # This depends on the structure of the composite DTB command
 243             files = cmd.split()
 244             files = files[files.index('-i') + 1:]
 245         else:
 246             files = [fname]
 247     else:
 248         files = [fname]
 249
 250     return (model, compat, files)
 251
 252 def build_fit(args):
 253     """Build the FIT from the provided files and arguments
 254
 255     Args:
 256         args (Namespace): Program arguments
 257
 258     Returns:
 259         tuple:
 260             bytes: FIT data
 261             int: Number of configurations generated
 262             size: Total uncompressed size of data
 263     """
 264     seq = 0
 265     size = 0
 266     fsw = libfdt.FdtSw()
 267     setup_fit(fsw, args.name)
 268     entries = []
 269     fdts = {}
 270
 271     # Handle the kernel
 272     with open(args.kernel, 'rb') as inf:
 273         comp_data = compress_data(inf, args.compress)
 274     size += os.path.getsize(args.kernel)
 275     write_kernel(fsw, comp_data, args)
 276
 277     for fname in args.dtbs:
 278         # Ignore non-DTB (*.dtb) files
 279         if os.path.splitext(fname)[1] != '.dtb':
 280             continue
 281
 282         (model, compat, files) = process_dtb(fname, args)
 283
 284         for fn in files:
 285             if fn not in fdts:
 286                 seq += 1
 287                 size += os.path.getsize(fn)
 288                 output_dtb(fsw, seq, fn, args.arch, args.compress)
 289                 fdts[fn] = seq
 290
 291         files_seq = [fdts[fn] for fn in files]
 292
 293         entries.append([model, compat, files_seq])
 294
 295     finish_fit(fsw, entries)
 296
 297     # Include the kernel itself in the returned file count
 298     return fsw.as_fdt().as_bytearray(), seq + 1, size
 299
 300
 301 def run_make_fit():
 302     """Run the tool's main logic"""
 303     args = parse_args()
 304
 305     out_data, count, size = build_fit(args)
 306     with open(args.output, 'wb') as outf:
 307         outf.write(out_data)
 308
 309     ext_fit_size = None
 310     if args.external:
 311         mkimage = os.environ.get('MKIMAGE', 'mkimage')
 312         subprocess.check_call([mkimage, '-E', '-F', args.output],
 313                               stdout=subprocess.DEVNULL)
 314
 315         with open(args.output, 'rb') as inf:
 316             data = inf.read()
 317         ext_fit = libfdt.FdtRo(data)
 318         ext_fit_size = ext_fit.totalsize()
 319
 320     if args.verbose:
 321         comp_size = len(out_data)
 322         print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB',
 323               end='')
 324         if ext_fit_size:
 325             print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB',
 326                   end='')
 327         print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB')
 328
 329
 330 if __name__ == "__main__":
 331     sys.exit(run_make_fit())