xref: /linux/scripts/make_fit.py (revision 4f9786035f9e519db41375818e1d0b5f20da2f10) 
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        try:
283            (model, compat, files) = process_dtb(fname, args)
284        except Exception as e:
285            sys.stderr.write(f"Error processing {fname}:\n")
286            raise e
287
288        for fn in files:
289            if fn not in fdts:
290                seq += 1
291                size += os.path.getsize(fn)
292                output_dtb(fsw, seq, fn, args.arch, args.compress)
293                fdts[fn] = seq
294
295        files_seq = [fdts[fn] for fn in files]
296
297        entries.append([model, compat, files_seq])
298
299    finish_fit(fsw, entries)
300
301    # Include the kernel itself in the returned file count
302    return fsw.as_fdt().as_bytearray(), seq + 1, size
303
304
305def run_make_fit():
306    """Run the tool's main logic"""
307    args = parse_args()
308
309    out_data, count, size = build_fit(args)
310    with open(args.output, 'wb') as outf:
311        outf.write(out_data)
312
313    ext_fit_size = None
314    if args.external:
315        mkimage = os.environ.get('MKIMAGE', 'mkimage')
316        subprocess.check_call([mkimage, '-E', '-F', args.output],
317                              stdout=subprocess.DEVNULL)
318
319        with open(args.output, 'rb') as inf:
320            data = inf.read()
321        ext_fit = libfdt.FdtRo(data)
322        ext_fit_size = ext_fit.totalsize()
323
324    if args.verbose:
325        comp_size = len(out_data)
326        print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB',
327              end='')
328        if ext_fit_size:
329            print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB',
330                  end='')
331        print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB')
332
333
334if __name__ == "__main__":
335    sys.exit(run_make_fit())
336