xref: /linux/scripts/make_fit.py (revision f898c16a0624e7f2dcb0b1cda6916c9be6489197) 
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
25The resulting FIT can be booted by bootloaders which support FIT, such
26as U-Boot, Linuxboot, Tianocore, etc.
27
28Note that this tool does not yet support adding a ramdisk / initrd.
29"""
30
31import argparse
32import collections
33import os
34import subprocess
35import sys
36import tempfile
37import time
38
39import libfdt
40
41
42# Tool extension and the name of the command-line tools
43CompTool = collections.namedtuple('CompTool', 'ext,tools')
44
45COMP_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
55def 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
85def 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
104def 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
129def 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
153def 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
187def 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
221def 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
259def 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
288if __name__ == "__main__":
289    sys.exit(run_make_fit())
290