xref: /linux/scripts/make_fit.py (revision a1ff5a7d78a036d6c2178ee5acd6ba4946243800) 
1#!/usr/bin/env python3
2# SPDX-License-Identifier: GPL-2.0+
3#
4# Copyright 2024 Google LLC
5# Written by Simon Glass <sjg@chromium.org>
6#
7
8"""Build a FIT containing a lot of devicetree files
9
10Usage:
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
15Creates a FIT containing the supplied kernel and a set of devicetree files,
16either specified individually or listed in a file (with an '@' prefix).
17
18Use -E to generate an external FIT (where the data is placed after the
19FIT data structure). This allows parsing of the data without loading
20the entire FIT.
21
22Use -c to compress the data, using bzip2, gzip, lz4, lzma, lzo and
23zstd algorithms.
24
25Use -D to decompose "composite" DTBs into their base components and
26deduplicate the resulting base DTBs and DTB overlays. This requires the
27DTBs to be sourced from the kernel build directory, as the implementation
28looks at the .cmd files produced by the kernel build.
29
30The resulting FIT can be booted by bootloaders which support FIT, such
31as U-Boot, Linuxboot, Tianocore, etc.
32
33Note that this tool does not yet support adding a ramdisk / initrd.
34"""
35
36import argparse
37import collections
38import os
39import subprocess
40import sys
41import tempfile
42import time
43
44import libfdt
45
46
47# Tool extension and the name of the command-line tools
48CompTool = collections.namedtuple('CompTool', 'ext,tools')
49
50COMP_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
60def 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
92def 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
111def 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
136def 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
160def 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
194def 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
215def 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
252def 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
301def 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
330if __name__ == "__main__":
331    sys.exit(run_make_fit())
332