17a23b027SSimon Glass#!/usr/bin/env python3 27a23b027SSimon Glass# SPDX-License-Identifier: GPL-2.0+ 37a23b027SSimon Glass# 47a23b027SSimon Glass# Copyright 2024 Google LLC 57a23b027SSimon Glass# Written by Simon Glass <sjg@chromium.org> 67a23b027SSimon Glass# 77a23b027SSimon Glass 87a23b027SSimon Glass"""Build a FIT containing a lot of devicetree files 97a23b027SSimon Glass 107a23b027SSimon GlassUsage: 117a23b027SSimon Glass make_fit.py -A arm64 -n 'Linux-6.6' -O linux 127a23b027SSimon Glass -o arch/arm64/boot/image.fit -k /tmp/kern/arch/arm64/boot/image.itk 137a23b027SSimon Glass @arch/arm64/boot/dts/dtbs-list -E -c gzip 147a23b027SSimon Glass 157a23b027SSimon GlassCreates a FIT containing the supplied kernel and a set of devicetree files, 167a23b027SSimon Glasseither specified individually or listed in a file (with an '@' prefix). 177a23b027SSimon Glass 187a23b027SSimon GlassUse -E to generate an external FIT (where the data is placed after the 197a23b027SSimon GlassFIT data structure). This allows parsing of the data without loading 207a23b027SSimon Glassthe entire FIT. 217a23b027SSimon Glass 227a23b027SSimon GlassUse -c to compress the data, using bzip2, gzip, lz4, lzma, lzo and 237a23b027SSimon Glasszstd algorithms. 247a23b027SSimon Glass 25*17c31adeSChen-Yu TsaiUse -D to decompose "composite" DTBs into their base components and 26*17c31adeSChen-Yu Tsaideduplicate the resulting base DTBs and DTB overlays. This requires the 27*17c31adeSChen-Yu TsaiDTBs to be sourced from the kernel build directory, as the implementation 28*17c31adeSChen-Yu Tsailooks at the .cmd files produced by the kernel build. 29*17c31adeSChen-Yu Tsai 307a23b027SSimon GlassThe resulting FIT can be booted by bootloaders which support FIT, such 317a23b027SSimon Glassas U-Boot, Linuxboot, Tianocore, etc. 327a23b027SSimon Glass 337a23b027SSimon GlassNote that this tool does not yet support adding a ramdisk / initrd. 347a23b027SSimon Glass""" 357a23b027SSimon Glass 367a23b027SSimon Glassimport argparse 377a23b027SSimon Glassimport collections 387a23b027SSimon Glassimport os 397a23b027SSimon Glassimport subprocess 407a23b027SSimon Glassimport sys 417a23b027SSimon Glassimport tempfile 427a23b027SSimon Glassimport time 437a23b027SSimon Glass 447a23b027SSimon Glassimport libfdt 457a23b027SSimon Glass 467a23b027SSimon Glass 477a23b027SSimon Glass# Tool extension and the name of the command-line tools 487a23b027SSimon GlassCompTool = collections.namedtuple('CompTool', 'ext,tools') 497a23b027SSimon Glass 507a23b027SSimon GlassCOMP_TOOLS = { 517a23b027SSimon Glass 'bzip2': CompTool('.bz2', 'bzip2'), 527a23b027SSimon Glass 'gzip': CompTool('.gz', 'pigz,gzip'), 537a23b027SSimon Glass 'lz4': CompTool('.lz4', 'lz4'), 547a23b027SSimon Glass 'lzma': CompTool('.lzma', 'lzma'), 557a23b027SSimon Glass 'lzo': CompTool('.lzo', 'lzop'), 567a23b027SSimon Glass 'zstd': CompTool('.zstd', 'zstd'), 577a23b027SSimon Glass} 587a23b027SSimon Glass 597a23b027SSimon Glass 607a23b027SSimon Glassdef parse_args(): 617a23b027SSimon Glass """Parse the program ArgumentParser 627a23b027SSimon Glass 637a23b027SSimon Glass Returns: 647a23b027SSimon Glass Namespace object containing the arguments 657a23b027SSimon Glass """ 667a23b027SSimon Glass epilog = 'Build a FIT from a directory tree containing .dtb files' 677a23b027SSimon Glass parser = argparse.ArgumentParser(epilog=epilog, fromfile_prefix_chars='@') 687a23b027SSimon Glass parser.add_argument('-A', '--arch', type=str, required=True, 697a23b027SSimon Glass help='Specifies the architecture') 707a23b027SSimon Glass parser.add_argument('-c', '--compress', type=str, default='none', 717a23b027SSimon Glass help='Specifies the compression') 72*17c31adeSChen-Yu Tsai parser.add_argument('-D', '--decompose-dtbs', action='store_true', 73*17c31adeSChen-Yu Tsai help='Decompose composite DTBs into base DTB and overlays') 747a23b027SSimon Glass parser.add_argument('-E', '--external', action='store_true', 757a23b027SSimon Glass help='Convert the FIT to use external data') 767a23b027SSimon Glass parser.add_argument('-n', '--name', type=str, required=True, 777a23b027SSimon Glass help='Specifies the name') 787a23b027SSimon Glass parser.add_argument('-o', '--output', type=str, required=True, 797a23b027SSimon Glass help='Specifies the output file (.fit)') 807a23b027SSimon Glass parser.add_argument('-O', '--os', type=str, required=True, 817a23b027SSimon Glass help='Specifies the operating system') 827a23b027SSimon Glass parser.add_argument('-k', '--kernel', type=str, required=True, 837a23b027SSimon Glass help='Specifies the (uncompressed) kernel input file (.itk)') 847a23b027SSimon Glass parser.add_argument('-v', '--verbose', action='store_true', 857a23b027SSimon Glass help='Enable verbose output') 867a23b027SSimon Glass parser.add_argument('dtbs', type=str, nargs='*', 877a23b027SSimon Glass help='Specifies the devicetree files to process') 887a23b027SSimon Glass 897a23b027SSimon Glass return parser.parse_args() 907a23b027SSimon Glass 917a23b027SSimon Glass 927a23b027SSimon Glassdef setup_fit(fsw, name): 937a23b027SSimon Glass """Make a start on writing the FIT 947a23b027SSimon Glass 957a23b027SSimon Glass Outputs the root properties and the 'images' node 967a23b027SSimon Glass 977a23b027SSimon Glass Args: 987a23b027SSimon Glass fsw (libfdt.FdtSw): Object to use for writing 997a23b027SSimon Glass name (str): Name of kernel image 1007a23b027SSimon Glass """ 1017a23b027SSimon Glass fsw.INC_SIZE = 65536 1027a23b027SSimon Glass fsw.finish_reservemap() 1037a23b027SSimon Glass fsw.begin_node('') 1047a23b027SSimon Glass fsw.property_string('description', f'{name} with devicetree set') 1057a23b027SSimon Glass fsw.property_u32('#address-cells', 1) 1067a23b027SSimon Glass 1077a23b027SSimon Glass fsw.property_u32('timestamp', int(time.time())) 1087a23b027SSimon Glass fsw.begin_node('images') 1097a23b027SSimon Glass 1107a23b027SSimon Glass 1117a23b027SSimon Glassdef write_kernel(fsw, data, args): 1127a23b027SSimon Glass """Write out the kernel image 1137a23b027SSimon Glass 1147a23b027SSimon Glass Writes a kernel node along with the required properties 1157a23b027SSimon Glass 1167a23b027SSimon Glass Args: 1177a23b027SSimon Glass fsw (libfdt.FdtSw): Object to use for writing 1187a23b027SSimon Glass data (bytes): Data to write (possibly compressed) 1197a23b027SSimon Glass args (Namespace): Contains necessary strings: 1207a23b027SSimon Glass arch: FIT architecture, e.g. 'arm64' 1217a23b027SSimon Glass fit_os: Operating Systems, e.g. 'linux' 1227a23b027SSimon Glass name: Name of OS, e.g. 'Linux-6.6.0-rc7' 1237a23b027SSimon Glass compress: Compression algorithm to use, e.g. 'gzip' 1247a23b027SSimon Glass """ 1257a23b027SSimon Glass with fsw.add_node('kernel'): 1267a23b027SSimon Glass fsw.property_string('description', args.name) 1277a23b027SSimon Glass fsw.property_string('type', 'kernel_noload') 1287a23b027SSimon Glass fsw.property_string('arch', args.arch) 1297a23b027SSimon Glass fsw.property_string('os', args.os) 1307a23b027SSimon Glass fsw.property_string('compression', args.compress) 1317a23b027SSimon Glass fsw.property('data', data) 1327a23b027SSimon Glass fsw.property_u32('load', 0) 1337a23b027SSimon Glass fsw.property_u32('entry', 0) 1347a23b027SSimon Glass 1357a23b027SSimon Glass 1367a23b027SSimon Glassdef finish_fit(fsw, entries): 1377a23b027SSimon Glass """Finish the FIT ready for use 1387a23b027SSimon Glass 1397a23b027SSimon Glass Writes the /configurations node and subnodes 1407a23b027SSimon Glass 1417a23b027SSimon Glass Args: 1427a23b027SSimon Glass fsw (libfdt.FdtSw): Object to use for writing 1437a23b027SSimon Glass entries (list of tuple): List of configurations: 1447a23b027SSimon Glass str: Description of model 1457a23b027SSimon Glass str: Compatible stringlist 1467a23b027SSimon Glass """ 1477a23b027SSimon Glass fsw.end_node() 1487a23b027SSimon Glass seq = 0 1497a23b027SSimon Glass with fsw.add_node('configurations'): 150*17c31adeSChen-Yu Tsai for model, compat, files in entries: 1517a23b027SSimon Glass seq += 1 1527a23b027SSimon Glass with fsw.add_node(f'conf-{seq}'): 1537a23b027SSimon Glass fsw.property('compatible', bytes(compat)) 1547a23b027SSimon Glass fsw.property_string('description', model) 155*17c31adeSChen-Yu Tsai fsw.property('fdt', bytes(''.join(f'fdt-{x}\x00' for x in files), "ascii")) 1567a23b027SSimon Glass fsw.property_string('kernel', 'kernel') 1577a23b027SSimon Glass fsw.end_node() 1587a23b027SSimon Glass 1597a23b027SSimon Glass 1607a23b027SSimon Glassdef compress_data(inf, compress): 1617a23b027SSimon Glass """Compress data using a selected algorithm 1627a23b027SSimon Glass 1637a23b027SSimon Glass Args: 1647a23b027SSimon Glass inf (IOBase): Filename containing the data to compress 1657a23b027SSimon Glass compress (str): Compression algorithm, e.g. 'gzip' 1667a23b027SSimon Glass 1677a23b027SSimon Glass Return: 1687a23b027SSimon Glass bytes: Compressed data 1697a23b027SSimon Glass """ 1707a23b027SSimon Glass if compress == 'none': 1717a23b027SSimon Glass return inf.read() 1727a23b027SSimon Glass 1737a23b027SSimon Glass comp = COMP_TOOLS.get(compress) 1747a23b027SSimon Glass if not comp: 1757a23b027SSimon Glass raise ValueError(f"Unknown compression algorithm '{compress}'") 1767a23b027SSimon Glass 1777a23b027SSimon Glass with tempfile.NamedTemporaryFile() as comp_fname: 1787a23b027SSimon Glass with open(comp_fname.name, 'wb') as outf: 1797a23b027SSimon Glass done = False 1807a23b027SSimon Glass for tool in comp.tools.split(','): 1817a23b027SSimon Glass try: 1827a23b027SSimon Glass subprocess.call([tool, '-c'], stdin=inf, stdout=outf) 1837a23b027SSimon Glass done = True 1847a23b027SSimon Glass break 1857a23b027SSimon Glass except FileNotFoundError: 1867a23b027SSimon Glass pass 1877a23b027SSimon Glass if not done: 1887a23b027SSimon Glass raise ValueError(f'Missing tool(s): {comp.tools}\n') 1897a23b027SSimon Glass with open(comp_fname.name, 'rb') as compf: 1907a23b027SSimon Glass comp_data = compf.read() 1917a23b027SSimon Glass return comp_data 1927a23b027SSimon Glass 1937a23b027SSimon Glass 1947a23b027SSimon Glassdef output_dtb(fsw, seq, fname, arch, compress): 1957a23b027SSimon Glass """Write out a single devicetree to the FIT 1967a23b027SSimon Glass 1977a23b027SSimon Glass Args: 1987a23b027SSimon Glass fsw (libfdt.FdtSw): Object to use for writing 1997a23b027SSimon Glass seq (int): Sequence number (1 for first) 200e06a698aSChen-Yu Tsai fname (str): Filename containing the DTB 2017a23b027SSimon Glass arch: FIT architecture, e.g. 'arm64' 2027a23b027SSimon Glass compress (str): Compressed algorithm, e.g. 'gzip' 2037a23b027SSimon Glass """ 2047a23b027SSimon Glass with fsw.add_node(f'fdt-{seq}'): 205*17c31adeSChen-Yu Tsai fsw.property_string('description', os.path.basename(fname)) 2067a23b027SSimon Glass fsw.property_string('type', 'flat_dt') 2077a23b027SSimon Glass fsw.property_string('arch', arch) 2087a23b027SSimon Glass fsw.property_string('compression', compress) 2097a23b027SSimon Glass 2107a23b027SSimon Glass with open(fname, 'rb') as inf: 2117a23b027SSimon Glass compressed = compress_data(inf, compress) 2127a23b027SSimon Glass fsw.property('data', compressed) 2137a23b027SSimon Glass 2147a23b027SSimon Glass 215*17c31adeSChen-Yu Tsaidef process_dtb(fname, args): 216*17c31adeSChen-Yu Tsai """Process an input DTB, decomposing it if requested and is possible 217*17c31adeSChen-Yu Tsai 218*17c31adeSChen-Yu Tsai Args: 219*17c31adeSChen-Yu Tsai fname (str): Filename containing the DTB 220*17c31adeSChen-Yu Tsai args (Namespace): Program arguments 221*17c31adeSChen-Yu Tsai Returns: 222*17c31adeSChen-Yu Tsai tuple: 223*17c31adeSChen-Yu Tsai str: Model name string 224*17c31adeSChen-Yu Tsai str: Root compatible string 225*17c31adeSChen-Yu Tsai files: list of filenames corresponding to the DTB 226*17c31adeSChen-Yu Tsai """ 227*17c31adeSChen-Yu Tsai # Get the compatible / model information 228*17c31adeSChen-Yu Tsai with open(fname, 'rb') as inf: 229*17c31adeSChen-Yu Tsai data = inf.read() 230*17c31adeSChen-Yu Tsai fdt = libfdt.FdtRo(data) 231*17c31adeSChen-Yu Tsai model = fdt.getprop(0, 'model').as_str() 232*17c31adeSChen-Yu Tsai compat = fdt.getprop(0, 'compatible') 233*17c31adeSChen-Yu Tsai 234*17c31adeSChen-Yu Tsai if args.decompose_dtbs: 235*17c31adeSChen-Yu Tsai # Check if the DTB needs to be decomposed 236*17c31adeSChen-Yu Tsai path, basename = os.path.split(fname) 237*17c31adeSChen-Yu Tsai cmd_fname = os.path.join(path, f'.{basename}.cmd') 238*17c31adeSChen-Yu Tsai with open(cmd_fname, 'r', encoding='ascii') as inf: 239*17c31adeSChen-Yu Tsai cmd = inf.read() 240*17c31adeSChen-Yu Tsai 241*17c31adeSChen-Yu Tsai if 'scripts/dtc/fdtoverlay' in cmd: 242*17c31adeSChen-Yu Tsai # This depends on the structure of the composite DTB command 243*17c31adeSChen-Yu Tsai files = cmd.split() 244*17c31adeSChen-Yu Tsai files = files[files.index('-i') + 1:] 245*17c31adeSChen-Yu Tsai else: 246*17c31adeSChen-Yu Tsai files = [fname] 247*17c31adeSChen-Yu Tsai else: 248*17c31adeSChen-Yu Tsai files = [fname] 249*17c31adeSChen-Yu Tsai 250*17c31adeSChen-Yu Tsai return (model, compat, files) 251*17c31adeSChen-Yu Tsai 2527a23b027SSimon Glassdef build_fit(args): 2537a23b027SSimon Glass """Build the FIT from the provided files and arguments 2547a23b027SSimon Glass 2557a23b027SSimon Glass Args: 2567a23b027SSimon Glass args (Namespace): Program arguments 2577a23b027SSimon Glass 2587a23b027SSimon Glass Returns: 2597a23b027SSimon Glass tuple: 2607a23b027SSimon Glass bytes: FIT data 2617a23b027SSimon Glass int: Number of configurations generated 2627a23b027SSimon Glass size: Total uncompressed size of data 2637a23b027SSimon Glass """ 2647a23b027SSimon Glass seq = 0 2657a23b027SSimon Glass size = 0 2667a23b027SSimon Glass fsw = libfdt.FdtSw() 2677a23b027SSimon Glass setup_fit(fsw, args.name) 2687a23b027SSimon Glass entries = [] 269*17c31adeSChen-Yu Tsai fdts = {} 2707a23b027SSimon Glass 2717a23b027SSimon Glass # Handle the kernel 2727a23b027SSimon Glass with open(args.kernel, 'rb') as inf: 2737a23b027SSimon Glass comp_data = compress_data(inf, args.compress) 2747a23b027SSimon Glass size += os.path.getsize(args.kernel) 2757a23b027SSimon Glass write_kernel(fsw, comp_data, args) 2767a23b027SSimon Glass 2777a23b027SSimon Glass for fname in args.dtbs: 278*17c31adeSChen-Yu Tsai # Ignore non-DTB (*.dtb) files 279*17c31adeSChen-Yu Tsai if os.path.splitext(fname)[1] != '.dtb': 280*17c31adeSChen-Yu Tsai continue 281*17c31adeSChen-Yu Tsai 282*17c31adeSChen-Yu Tsai (model, compat, files) = process_dtb(fname, args) 283*17c31adeSChen-Yu Tsai 284*17c31adeSChen-Yu Tsai for fn in files: 285*17c31adeSChen-Yu Tsai if fn not in fdts: 2867a23b027SSimon Glass seq += 1 287*17c31adeSChen-Yu Tsai size += os.path.getsize(fn) 288*17c31adeSChen-Yu Tsai output_dtb(fsw, seq, fn, args.arch, args.compress) 289*17c31adeSChen-Yu Tsai fdts[fn] = seq 290*17c31adeSChen-Yu Tsai 291*17c31adeSChen-Yu Tsai files_seq = [fdts[fn] for fn in files] 292*17c31adeSChen-Yu Tsai 293*17c31adeSChen-Yu Tsai entries.append([model, compat, files_seq]) 2947a23b027SSimon Glass 2957a23b027SSimon Glass finish_fit(fsw, entries) 2967a23b027SSimon Glass 2977a23b027SSimon Glass # Include the kernel itself in the returned file count 2987a23b027SSimon Glass return fsw.as_fdt().as_bytearray(), seq + 1, size 2997a23b027SSimon Glass 3007a23b027SSimon Glass 3017a23b027SSimon Glassdef run_make_fit(): 3027a23b027SSimon Glass """Run the tool's main logic""" 3037a23b027SSimon Glass args = parse_args() 3047a23b027SSimon Glass 3057a23b027SSimon Glass out_data, count, size = build_fit(args) 3067a23b027SSimon Glass with open(args.output, 'wb') as outf: 3077a23b027SSimon Glass outf.write(out_data) 3087a23b027SSimon Glass 3097a23b027SSimon Glass ext_fit_size = None 3107a23b027SSimon Glass if args.external: 3117a23b027SSimon Glass mkimage = os.environ.get('MKIMAGE', 'mkimage') 3127a23b027SSimon Glass subprocess.check_call([mkimage, '-E', '-F', args.output], 3137a23b027SSimon Glass stdout=subprocess.DEVNULL) 3147a23b027SSimon Glass 3157a23b027SSimon Glass with open(args.output, 'rb') as inf: 3167a23b027SSimon Glass data = inf.read() 3177a23b027SSimon Glass ext_fit = libfdt.FdtRo(data) 3187a23b027SSimon Glass ext_fit_size = ext_fit.totalsize() 3197a23b027SSimon Glass 3207a23b027SSimon Glass if args.verbose: 3217a23b027SSimon Glass comp_size = len(out_data) 3227a23b027SSimon Glass print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB', 3237a23b027SSimon Glass end='') 3247a23b027SSimon Glass if ext_fit_size: 3257a23b027SSimon Glass print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB', 3267a23b027SSimon Glass end='') 3277a23b027SSimon Glass print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB') 3287a23b027SSimon Glass 3297a23b027SSimon Glass 3307a23b027SSimon Glassif __name__ == "__main__": 3317a23b027SSimon Glass sys.exit(run_make_fit()) 332