scripts/make_fit.py - linux - Git at Google

 #!/usr/bin/env python3
 # SPDX-License-Identifier: GPL-2.0+
 #
 # Copyright 2024 Google LLC
 # Written by Simon Glass <sjg@chromium.org>
 #

 """Build a FIT containing a lot of devicetree files

 Usage:
     make_fit.py -A arm64 -n 'Linux-6.6' -O linux
         -o arch/arm64/boot/image.fit -k /tmp/kern/arch/arm64/boot/image.itk
         @arch/arm64/boot/dts/dtbs-list -E -c gzip

 Creates a FIT containing the supplied kernel and a set of devicetree files,
 either specified individually or listed in a file (with an '@' prefix).

 Use -E to generate an external FIT (where the data is placed after the
 FIT data structure). This allows parsing of the data without loading
 the entire FIT.

 Use -c to compress the data, using bzip2, gzip, lz4, lzma, lzo and
 zstd algorithms.

 Use -D to decompose "composite" DTBs into their base components and
 deduplicate the resulting base DTBs and DTB overlays. This requires the
 DTBs to be sourced from the kernel build directory, as the implementation
 looks at the .cmd files produced by the kernel build.

 The resulting FIT can be booted by bootloaders which support FIT, such
 as U-Boot, Linuxboot, Tianocore, etc.

 Note that this tool does not yet support adding a ramdisk / initrd.
 """

 import argparse
 import collections
 import os
 import subprocess
 import sys
 import tempfile
 import time

 import libfdt


 # Tool extension and the name of the command-line tools
 CompTool = collections.namedtuple('CompTool', 'ext,tools')

 COMP_TOOLS = {
     'bzip2': CompTool('.bz2', 'bzip2'),
     'gzip': CompTool('.gz', 'pigz,gzip'),
     'lz4': CompTool('.lz4', 'lz4'),
     'lzma': CompTool('.lzma', 'lzma'),
     'lzo': CompTool('.lzo', 'lzop'),
     'zstd': CompTool('.zstd', 'zstd'),
 }


 def parse_args():
     """Parse the program ArgumentParser

     Returns:
         Namespace object containing the arguments
     """
     epilog = 'Build a FIT from a directory tree containing .dtb files'
     parser = argparse.ArgumentParser(epilog=epilog, fromfile_prefix_chars='@')
     parser.add_argument('-A', '--arch', type=str, required=True,
           help='Specifies the architecture')
     parser.add_argument('-c', '--compress', type=str, default='none',
           help='Specifies the compression')
     parser.add_argument('-D', '--decompose-dtbs', action='store_true',
           help='Decompose composite DTBs into base DTB and overlays')
     parser.add_argument('-E', '--external', action='store_true',
           help='Convert the FIT to use external data')
     parser.add_argument('-n', '--name', type=str, required=True,
           help='Specifies the name')
     parser.add_argument('-o', '--output', type=str, required=True,
           help='Specifies the output file (.fit)')
     parser.add_argument('-O', '--os', type=str, required=True,
           help='Specifies the operating system')
     parser.add_argument('-k', '--kernel', type=str, required=True,
           help='Specifies the (uncompressed) kernel input file (.itk)')
     parser.add_argument('-v', '--verbose', action='store_true',
                         help='Enable verbose output')
     parser.add_argument('dtbs', type=str, nargs='*',
           help='Specifies the devicetree files to process')

     return parser.parse_args()


 def setup_fit(fsw, name):
     """Make a start on writing the FIT

     Outputs the root properties and the 'images' node

     Args:
         fsw (libfdt.FdtSw): Object to use for writing
         name (str): Name of kernel image
     """
     fsw.INC_SIZE = 65536
     fsw.finish_reservemap()
     fsw.begin_node('')
     fsw.property_string('description', f'{name} with devicetree set')
     fsw.property_u32('#address-cells', 1)

     fsw.property_u32('timestamp', int(time.time()))
     fsw.begin_node('images')


 def write_kernel(fsw, data, args):
     """Write out the kernel image

     Writes a kernel node along with the required properties

     Args:
         fsw (libfdt.FdtSw): Object to use for writing
         data (bytes): Data to write (possibly compressed)
         args (Namespace): Contains necessary strings:
             arch: FIT architecture, e.g. 'arm64'
             fit_os: Operating Systems, e.g. 'linux'
             name: Name of OS, e.g. 'Linux-6.6.0-rc7'
             compress: Compression algorithm to use, e.g. 'gzip'
     """
     with fsw.add_node('kernel'):
         fsw.property_string('description', args.name)
         fsw.property_string('type', 'kernel_noload')
         fsw.property_string('arch', args.arch)
         fsw.property_string('os', args.os)
         fsw.property_string('compression', args.compress)
         fsw.property('data', data)
         fsw.property_u32('load', 0)
         fsw.property_u32('entry', 0)


 def finish_fit(fsw, entries):
     """Finish the FIT ready for use

     Writes the /configurations node and subnodes

     Args:
         fsw (libfdt.FdtSw): Object to use for writing
         entries (list of tuple): List of configurations:
             str: Description of model
             str: Compatible stringlist
     """
     fsw.end_node()
     seq = 0
     with fsw.add_node('configurations'):
         for model, compat, files in entries:
             seq += 1
             with fsw.add_node(f'conf-{seq}'):
                 fsw.property('compatible', bytes(compat))
                 fsw.property_string('description', model)
                 fsw.property('fdt', bytes(''.join(f'fdt-{x}\x00' for x in files), "ascii"))
                 fsw.property_string('kernel', 'kernel')
     fsw.end_node()


 def compress_data(inf, compress):
     """Compress data using a selected algorithm

     Args:
         inf (IOBase): Filename containing the data to compress
         compress (str): Compression algorithm, e.g. 'gzip'

     Return:
         bytes: Compressed data
     """
     if compress == 'none':
         return inf.read()

     comp = COMP_TOOLS.get(compress)
     if not comp:
         raise ValueError(f"Unknown compression algorithm '{compress}'")

     with tempfile.NamedTemporaryFile() as comp_fname:
         with open(comp_fname.name, 'wb') as outf:
             done = False
             for tool in comp.tools.split(','):
                 try:
                     subprocess.call([tool, '-c'], stdin=inf, stdout=outf)
                     done = True
                     break
                 except FileNotFoundError:
                     pass
             if not done:
                 raise ValueError(f'Missing tool(s): {comp.tools}\n')
             with open(comp_fname.name, 'rb') as compf:
                 comp_data = compf.read()
     return comp_data


 def output_dtb(fsw, seq, fname, arch, compress):
     """Write out a single devicetree to the FIT

     Args:
         fsw (libfdt.FdtSw): Object to use for writing
         seq (int): Sequence number (1 for first)
         fname (str): Filename containing the DTB
         arch: FIT architecture, e.g. 'arm64'
         compress (str): Compressed algorithm, e.g. 'gzip'
     """
     with fsw.add_node(f'fdt-{seq}'):
         fsw.property_string('description', os.path.basename(fname))
         fsw.property_string('type', 'flat_dt')
         fsw.property_string('arch', arch)
         fsw.property_string('compression', compress)

         with open(fname, 'rb') as inf:
             compressed = compress_data(inf, compress)
         fsw.property('data', compressed)


 def process_dtb(fname, args):
     """Process an input DTB, decomposing it if requested and is possible

     Args:
         fname (str): Filename containing the DTB
         args (Namespace): Program arguments
     Returns:
         tuple:
             str: Model name string
             str: Root compatible string
             files: list of filenames corresponding to the DTB
     """
     # Get the compatible / model information
     with open(fname, 'rb') as inf:
         data = inf.read()
     fdt = libfdt.FdtRo(data)
     model = fdt.getprop(0, 'model').as_str()
     compat = fdt.getprop(0, 'compatible')

     if args.decompose_dtbs:
         # Check if the DTB needs to be decomposed
         path, basename = os.path.split(fname)
         cmd_fname = os.path.join(path, f'.{basename}.cmd')
         with open(cmd_fname, 'r', encoding='ascii') as inf:
             cmd = inf.read()

         if 'scripts/dtc/fdtoverlay' in cmd:
             # This depends on the structure of the composite DTB command
             files = cmd.split()
             files = files[files.index('-i') + 1:]
         else:
             files = [fname]
     else:
         files = [fname]

     return (model, compat, files)

 def build_fit(args):
     """Build the FIT from the provided files and arguments

     Args:
         args (Namespace): Program arguments

     Returns:
         tuple:
             bytes: FIT data
             int: Number of configurations generated
             size: Total uncompressed size of data
     """
     seq = 0
     size = 0
     fsw = libfdt.FdtSw()
     setup_fit(fsw, args.name)
     entries = []
     fdts = {}

     # Handle the kernel
     with open(args.kernel, 'rb') as inf:
         comp_data = compress_data(inf, args.compress)
     size += os.path.getsize(args.kernel)
     write_kernel(fsw, comp_data, args)

     for fname in args.dtbs:
         # Ignore non-DTB (*.dtb) files
         if os.path.splitext(fname)[1] != '.dtb':
             continue

         (model, compat, files) = process_dtb(fname, args)

         for fn in files:
             if fn not in fdts:
                 seq += 1
                 size += os.path.getsize(fn)
                 output_dtb(fsw, seq, fn, args.arch, args.compress)
                 fdts[fn] = seq

         files_seq = [fdts[fn] for fn in files]

         entries.append([model, compat, files_seq])

     finish_fit(fsw, entries)

     # Include the kernel itself in the returned file count
     return fsw.as_fdt().as_bytearray(), seq + 1, size


 def run_make_fit():
     """Run the tool's main logic"""
     args = parse_args()

     out_data, count, size = build_fit(args)
     with open(args.output, 'wb') as outf:
         outf.write(out_data)

     ext_fit_size = None
     if args.external:
         mkimage = os.environ.get('MKIMAGE', 'mkimage')
         subprocess.check_call([mkimage, '-E', '-F', args.output],
                               stdout=subprocess.DEVNULL)

         with open(args.output, 'rb') as inf:
             data = inf.read()
         ext_fit = libfdt.FdtRo(data)
         ext_fit_size = ext_fit.totalsize()

     if args.verbose:
         comp_size = len(out_data)
         print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB',
               end='')
         if ext_fit_size:
             print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB',
                   end='')
         print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB')


 if __name__ == "__main__":
     sys.exit(run_make_fit())
	#!/usr/bin/env python3
	# SPDX-License-Identifier: GPL-2.0+
	#
	# Copyright 2024 Google LLC
	# Written by Simon Glass <sjg@chromium.org>
	#

	"""Build a FIT containing a lot of devicetree files

	Usage:
	make_fit.py -A arm64 -n 'Linux-6.6' -O linux
	-o arch/arm64/boot/image.fit -k /tmp/kern/arch/arm64/boot/image.itk
	@arch/arm64/boot/dts/dtbs-list -E -c gzip

	Creates a FIT containing the supplied kernel and a set of devicetree files,
	either specified individually or listed in a file (with an '@' prefix).

	Use -E to generate an external FIT (where the data is placed after the
	FIT data structure). This allows parsing of the data without loading
	the entire FIT.

	Use -c to compress the data, using bzip2, gzip, lz4, lzma, lzo and
	zstd algorithms.

	Use -D to decompose "composite" DTBs into their base components and
	deduplicate the resulting base DTBs and DTB overlays. This requires the
	DTBs to be sourced from the kernel build directory, as the implementation
	looks at the .cmd files produced by the kernel build.

	The resulting FIT can be booted by bootloaders which support FIT, such
	as U-Boot, Linuxboot, Tianocore, etc.

	Note that this tool does not yet support adding a ramdisk / initrd.
	"""

	import argparse
	import collections
	import os
	import subprocess
	import sys
	import tempfile
	import time

	import libfdt


	# Tool extension and the name of the command-line tools
	CompTool = collections.namedtuple('CompTool', 'ext,tools')

	COMP_TOOLS = {
	'bzip2': CompTool('.bz2', 'bzip2'),
	'gzip': CompTool('.gz', 'pigz,gzip'),
	'lz4': CompTool('.lz4', 'lz4'),
	'lzma': CompTool('.lzma', 'lzma'),
	'lzo': CompTool('.lzo', 'lzop'),
	'zstd': CompTool('.zstd', 'zstd'),
	}


	def parse_args():
	"""Parse the program ArgumentParser

	Returns:
	Namespace object containing the arguments
	"""
	epilog = 'Build a FIT from a directory tree containing .dtb files'
	parser = argparse.ArgumentParser(epilog=epilog, fromfile_prefix_chars='@')
	parser.add_argument('-A', '--arch', type=str, required=True,
	help='Specifies the architecture')
	parser.add_argument('-c', '--compress', type=str, default='none',
	help='Specifies the compression')
	parser.add_argument('-D', '--decompose-dtbs', action='store_true',
	help='Decompose composite DTBs into base DTB and overlays')
	parser.add_argument('-E', '--external', action='store_true',
	help='Convert the FIT to use external data')
	parser.add_argument('-n', '--name', type=str, required=True,
	help='Specifies the name')
	parser.add_argument('-o', '--output', type=str, required=True,
	help='Specifies the output file (.fit)')
	parser.add_argument('-O', '--os', type=str, required=True,
	help='Specifies the operating system')
	parser.add_argument('-k', '--kernel', type=str, required=True,
	help='Specifies the (uncompressed) kernel input file (.itk)')
	parser.add_argument('-v', '--verbose', action='store_true',
	help='Enable verbose output')
	parser.add_argument('dtbs', type=str, nargs='*',
	help='Specifies the devicetree files to process')

	return parser.parse_args()


	def setup_fit(fsw, name):
	"""Make a start on writing the FIT

	Outputs the root properties and the 'images' node

	Args:
	fsw (libfdt.FdtSw): Object to use for writing
	name (str): Name of kernel image
	"""
	fsw.INC_SIZE = 65536
	fsw.finish_reservemap()
	fsw.begin_node('')
	fsw.property_string('description', f'{name} with devicetree set')
	fsw.property_u32('#address-cells', 1)

	fsw.property_u32('timestamp', int(time.time()))
	fsw.begin_node('images')


	def write_kernel(fsw, data, args):
	"""Write out the kernel image

	Writes a kernel node along with the required properties

	Args:
	fsw (libfdt.FdtSw): Object to use for writing
	data (bytes): Data to write (possibly compressed)
	args (Namespace): Contains necessary strings:
	arch: FIT architecture, e.g. 'arm64'
	fit_os: Operating Systems, e.g. 'linux'
	name: Name of OS, e.g. 'Linux-6.6.0-rc7'
	compress: Compression algorithm to use, e.g. 'gzip'
	"""
	with fsw.add_node('kernel'):
	fsw.property_string('description', args.name)
	fsw.property_string('type', 'kernel_noload')
	fsw.property_string('arch', args.arch)
	fsw.property_string('os', args.os)
	fsw.property_string('compression', args.compress)
	fsw.property('data', data)
	fsw.property_u32('load', 0)
	fsw.property_u32('entry', 0)


	def finish_fit(fsw, entries):
	"""Finish the FIT ready for use

	Writes the /configurations node and subnodes

	Args:
	fsw (libfdt.FdtSw): Object to use for writing
	entries (list of tuple): List of configurations:
	str: Description of model
	str: Compatible stringlist
	"""
	fsw.end_node()
	seq = 0
	with fsw.add_node('configurations'):
	for model, compat, files in entries:
	seq += 1
	with fsw.add_node(f'conf-{seq}'):
	fsw.property('compatible', bytes(compat))
	fsw.property_string('description', model)
	fsw.property('fdt', bytes(''.join(f'fdt-{x}\x00' for x in files), "ascii"))
	fsw.property_string('kernel', 'kernel')
	fsw.end_node()


	def compress_data(inf, compress):
	"""Compress data using a selected algorithm

	Args:
	inf (IOBase): Filename containing the data to compress
	compress (str): Compression algorithm, e.g. 'gzip'

	Return:
	bytes: Compressed data
	"""
	if compress == 'none':
	return inf.read()

	comp = COMP_TOOLS.get(compress)
	if not comp:
	raise ValueError(f"Unknown compression algorithm '{compress}'")

	with tempfile.NamedTemporaryFile() as comp_fname:
	with open(comp_fname.name, 'wb') as outf:
	done = False
	for tool in comp.tools.split(','):
	try:
	subprocess.call([tool, '-c'], stdin=inf, stdout=outf)
	done = True
	break
	except FileNotFoundError:
	pass
	if not done:
	raise ValueError(f'Missing tool(s): {comp.tools}\n')
	with open(comp_fname.name, 'rb') as compf:
	comp_data = compf.read()
	return comp_data


	def output_dtb(fsw, seq, fname, arch, compress):
	"""Write out a single devicetree to the FIT

	Args:
	fsw (libfdt.FdtSw): Object to use for writing
	seq (int): Sequence number (1 for first)
	fname (str): Filename containing the DTB
	arch: FIT architecture, e.g. 'arm64'
	compress (str): Compressed algorithm, e.g. 'gzip'
	"""
	with fsw.add_node(f'fdt-{seq}'):
	fsw.property_string('description', os.path.basename(fname))
	fsw.property_string('type', 'flat_dt')
	fsw.property_string('arch', arch)
	fsw.property_string('compression', compress)

	with open(fname, 'rb') as inf:
	compressed = compress_data(inf, compress)
	fsw.property('data', compressed)


	def process_dtb(fname, args):
	"""Process an input DTB, decomposing it if requested and is possible

	Args:
	fname (str): Filename containing the DTB
	args (Namespace): Program arguments
	Returns:
	tuple:
	str: Model name string
	str: Root compatible string
	files: list of filenames corresponding to the DTB
	"""
	# Get the compatible / model information
	with open(fname, 'rb') as inf:
	data = inf.read()
	fdt = libfdt.FdtRo(data)
	model = fdt.getprop(0, 'model').as_str()
	compat = fdt.getprop(0, 'compatible')

	if args.decompose_dtbs:
	# Check if the DTB needs to be decomposed
	path, basename = os.path.split(fname)
	cmd_fname = os.path.join(path, f'.{basename}.cmd')
	with open(cmd_fname, 'r', encoding='ascii') as inf:
	cmd = inf.read()

	if 'scripts/dtc/fdtoverlay' in cmd:
	# This depends on the structure of the composite DTB command
	files = cmd.split()
	files = files[files.index('-i') + 1:]
	else:
	files = [fname]
	else:
	files = [fname]

	return (model, compat, files)

	def build_fit(args):
	"""Build the FIT from the provided files and arguments

	Args:
	args (Namespace): Program arguments

	Returns:
	tuple:
	bytes: FIT data
	int: Number of configurations generated
	size: Total uncompressed size of data
	"""
	seq = 0
	size = 0
	fsw = libfdt.FdtSw()
	setup_fit(fsw, args.name)
	entries = []
	fdts = {}

	# Handle the kernel
	with open(args.kernel, 'rb') as inf:
	comp_data = compress_data(inf, args.compress)
	size += os.path.getsize(args.kernel)
	write_kernel(fsw, comp_data, args)

	for fname in args.dtbs:
	# Ignore non-DTB (*.dtb) files
	if os.path.splitext(fname)[1] != '.dtb':
	continue

	(model, compat, files) = process_dtb(fname, args)

	for fn in files:
	if fn not in fdts:
	seq += 1
	size += os.path.getsize(fn)
	output_dtb(fsw, seq, fn, args.arch, args.compress)
	fdts[fn] = seq

	files_seq = [fdts[fn] for fn in files]

	entries.append([model, compat, files_seq])

	finish_fit(fsw, entries)

	# Include the kernel itself in the returned file count
	return fsw.as_fdt().as_bytearray(), seq + 1, size


	def run_make_fit():
	"""Run the tool's main logic"""
	args = parse_args()

	out_data, count, size = build_fit(args)
	with open(args.output, 'wb') as outf:
	outf.write(out_data)

	ext_fit_size = None
	if args.external:
	mkimage = os.environ.get('MKIMAGE', 'mkimage')
	subprocess.check_call([mkimage, '-E', '-F', args.output],
	stdout=subprocess.DEVNULL)

	with open(args.output, 'rb') as inf:
	data = inf.read()
	ext_fit = libfdt.FdtRo(data)
	ext_fit_size = ext_fit.totalsize()

	if args.verbose:
	comp_size = len(out_data)
	print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB',
	end='')
	if ext_fit_size:
	print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB',
	end='')
	print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB')


	if __name__ == "__main__":
	sys.exit(run_make_fit())