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 The resulting FIT can be booted by bootloaders which support FIT, such
  26 as U-Boot, Linuxboot, Tianocore, etc.
  27
  28 Note that this tool does not yet support adding a ramdisk / initrd.
  29 """
  30
  31 import argparse
  32 import collections
  33 import os
  34 import subprocess
  35 import sys
  36 import tempfile
  37 import time
  38
  39 import libfdt
  40
  41
  42 # Tool extension and the name of the command-line tools
  43 CompTool = collections.namedtuple('CompTool', 'ext,tools')
  44
  45 COMP_TOOLS = {
  46     'bzip2': CompTool('.bz2', 'bzip2'),
  47     'gzip': CompTool('.gz', 'pigz,gzip'),
  48     'lz4': CompTool('.lz4', 'lz4'),
  49     'lzma': CompTool('.lzma', 'lzma'),
  50     'lzo': CompTool('.lzo', 'lzop'),
  51     'zstd': CompTool('.zstd', 'zstd'),
  52 }
  53
  54
  55 def parse_args():
  56     """Parse the program ArgumentParser
  57
  58     Returns:
  59         Namespace object containing the arguments
  60     """
  61     epilog = 'Build a FIT from a directory tree containing .dtb files'
  62     parser = argparse.ArgumentParser(epilog=epilog, fromfile_prefix_chars='@')
  63     parser.add_argument('-A', '--arch', type=str, required=True,
  64           help='Specifies the architecture')
  65     parser.add_argument('-c', '--compress', type=str, default='none',
  66           help='Specifies the compression')
  67     parser.add_argument('-E', '--external', action='store_true',
  68           help='Convert the FIT to use external data')
  69     parser.add_argument('-n', '--name', type=str, required=True,
  70           help='Specifies the name')
  71     parser.add_argument('-o', '--output', type=str, required=True,
  72           help='Specifies the output file (.fit)')
  73     parser.add_argument('-O', '--os', type=str, required=True,
  74           help='Specifies the operating system')
  75     parser.add_argument('-k', '--kernel', type=str, required=True,
  76           help='Specifies the (uncompressed) kernel input file (.itk)')
  77     parser.add_argument('-v', '--verbose', action='store_true',
  78                         help='Enable verbose output')
  79     parser.add_argument('dtbs', type=str, nargs='*',
  80           help='Specifies the devicetree files to process')
  81
  82     return parser.parse_args()
  83
  84
  85 def setup_fit(fsw, name):
  86     """Make a start on writing the FIT
  87
  88     Outputs the root properties and the 'images' node
  89
  90     Args:
  91         fsw (libfdt.FdtSw): Object to use for writing
  92         name (str): Name of kernel image
  93     """
  94     fsw.INC_SIZE = 65536
  95     fsw.finish_reservemap()
  96     fsw.begin_node('')
  97     fsw.property_string('description', f'{name} with devicetree set')
  98     fsw.property_u32('#address-cells', 1)
  99
 100     fsw.property_u32('timestamp', int(time.time()))
 101     fsw.begin_node('images')
 102
 103
 104 def write_kernel(fsw, data, args):
 105     """Write out the kernel image
 106
 107     Writes a kernel node along with the required properties
 108
 109     Args:
 110         fsw (libfdt.FdtSw): Object to use for writing
 111         data (bytes): Data to write (possibly compressed)
 112         args (Namespace): Contains necessary strings:
 113             arch: FIT architecture, e.g. 'arm64'
 114             fit_os: Operating Systems, e.g. 'linux'
 115             name: Name of OS, e.g. 'Linux-6.6.0-rc7'
 116             compress: Compression algorithm to use, e.g. 'gzip'
 117     """
 118     with fsw.add_node('kernel'):
 119         fsw.property_string('description', args.name)
 120         fsw.property_string('type', 'kernel_noload')
 121         fsw.property_string('arch', args.arch)
 122         fsw.property_string('os', args.os)
 123         fsw.property_string('compression', args.compress)
 124         fsw.property('data', data)
 125         fsw.property_u32('load', 0)
 126         fsw.property_u32('entry', 0)
 127
 128
 129 def finish_fit(fsw, entries):
 130     """Finish the FIT ready for use
 131
 132     Writes the /configurations node and subnodes
 133
 134     Args:
 135         fsw (libfdt.FdtSw): Object to use for writing
 136         entries (list of tuple): List of configurations:
 137             str: Description of model
 138             str: Compatible stringlist
 139     """
 140     fsw.end_node()
 141     seq = 0
 142     with fsw.add_node('configurations'):
 143         for model, compat in entries:
 144             seq += 1
 145             with fsw.add_node(f'conf-{seq}'):
 146                 fsw.property('compatible', bytes(compat))
 147                 fsw.property_string('description', model)
 148                 fsw.property_string('fdt', f'fdt-{seq}')
 149                 fsw.property_string('kernel', 'kernel')
 150     fsw.end_node()
 151
 152
 153 def compress_data(inf, compress):
 154     """Compress data using a selected algorithm
 155
 156     Args:
 157         inf (IOBase): Filename containing the data to compress
 158         compress (str): Compression algorithm, e.g. 'gzip'
 159
 160     Return:
 161         bytes: Compressed data
 162     """
 163     if compress == 'none':
 164         return inf.read()
 165
 166     comp = COMP_TOOLS.get(compress)
 167     if not comp:
 168         raise ValueError(f"Unknown compression algorithm '{compress}'")
 169
 170     with tempfile.NamedTemporaryFile() as comp_fname:
 171         with open(comp_fname.name, 'wb') as outf:
 172             done = False
 173             for tool in comp.tools.split(','):
 174                 try:
 175                     subprocess.call([tool, '-c'], stdin=inf, stdout=outf)
 176                     done = True
 177                     break
 178                 except FileNotFoundError:
 179                     pass
 180             if not done:
 181                 raise ValueError(f'Missing tool(s): {comp.tools}\n')
 182             with open(comp_fname.name, 'rb') as compf:
 183                 comp_data = compf.read()
 184     return comp_data
 185
 186
 187 def output_dtb(fsw, seq, fname, arch, compress):
 188     """Write out a single devicetree to the FIT
 189
 190     Args:
 191         fsw (libfdt.FdtSw): Object to use for writing
 192         seq (int): Sequence number (1 for first)
 193         fname (str): Filename containing the DTB
 194         arch: FIT architecture, e.g. 'arm64'
 195         compress (str): Compressed algorithm, e.g. 'gzip'
 196
 197     Returns:
 198         tuple:
 199             str: Model name
 200             bytes: Compatible stringlist
 201     """
 202     with fsw.add_node(f'fdt-{seq}'):
 203         # Get the compatible / model information
 204         with open(fname, 'rb') as inf:
 205             data = inf.read()
 206         fdt = libfdt.FdtRo(data)
 207         model = fdt.getprop(0, 'model').as_str()
 208         compat = fdt.getprop(0, 'compatible')
 209
 210         fsw.property_string('description', model)
 211         fsw.property_string('type', 'flat_dt')
 212         fsw.property_string('arch', arch)
 213         fsw.property_string('compression', compress)
 214
 215         with open(fname, 'rb') as inf:
 216             compressed = compress_data(inf, compress)
 217         fsw.property('data', compressed)
 218     return model, compat
 219
 220
 221 def build_fit(args):
 222     """Build the FIT from the provided files and arguments
 223
 224     Args:
 225         args (Namespace): Program arguments
 226
 227     Returns:
 228         tuple:
 229             bytes: FIT data
 230             int: Number of configurations generated
 231             size: Total uncompressed size of data
 232     """
 233     seq = 0
 234     size = 0
 235     fsw = libfdt.FdtSw()
 236     setup_fit(fsw, args.name)
 237     entries = []
 238
 239     # Handle the kernel
 240     with open(args.kernel, 'rb') as inf:
 241         comp_data = compress_data(inf, args.compress)
 242     size += os.path.getsize(args.kernel)
 243     write_kernel(fsw, comp_data, args)
 244
 245     for fname in args.dtbs:
 246         # Ignore overlay (.dtbo) files
 247         if os.path.splitext(fname)[1] == '.dtb':
 248             seq += 1
 249             size += os.path.getsize(fname)
 250             model, compat = output_dtb(fsw, seq, fname, args.arch, args.compress)
 251             entries.append([model, compat])
 252
 253     finish_fit(fsw, entries)
 254
 255     # Include the kernel itself in the returned file count
 256     return fsw.as_fdt().as_bytearray(), seq + 1, size
 257
 258
 259 def run_make_fit():
 260     """Run the tool's main logic"""
 261     args = parse_args()
 262
 263     out_data, count, size = build_fit(args)
 264     with open(args.output, 'wb') as outf:
 265         outf.write(out_data)
 266
 267     ext_fit_size = None
 268     if args.external:
 269         mkimage = os.environ.get('MKIMAGE', 'mkimage')
 270         subprocess.check_call([mkimage, '-E', '-F', args.output],
 271                               stdout=subprocess.DEVNULL)
 272
 273         with open(args.output, 'rb') as inf:
 274             data = inf.read()
 275         ext_fit = libfdt.FdtRo(data)
 276         ext_fit_size = ext_fit.totalsize()
 277
 278     if args.verbose:
 279         comp_size = len(out_data)
 280         print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB',
 281               end='')
 282         if ext_fit_size:
 283             print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB',
 284                   end='')
 285         print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB')
 286
 287
 288 if __name__ == "__main__":
 289     sys.exit(run_make_fit())