initrd-builder

#!/usr/bin/env python3
# Populate an initrd directory with programs and their libraries
# https://www.kernel.org/doc/Documentation/early-userspace/buffer-format.txt
import sys
import re
import os
import cpiofile
import traceback
import argparse
import subprocess
from tempfile import NamedTemporaryFile

bindir = "bin"
libdir = "lib64"
verbose = False
cwd = '.'
no_strip = False

depsfile = None
deps = {}
cpio = cpiofile.CPIO()

def system(*s):
	# do not close file descriptors, which will allow
	# communication from sub-make invocations to the make
	# that invoked us.  allow failures
	return subprocess.run(s, close_fds=False, capture_output=True)

def die(s):
	print(s, file=sys.stderr)
	exit(1)

def filemode(filename):
	try:
		s = os.stat(filename)
		return s.st_mode
	except Exception as e:
		print(e, file=sys.stderr)
		return None
def exists(f):
	try:
		os.stat(f)
		return True
	except Exception as e:
		return False
def readfile(filename):
	with open(filename, "rb") as f:
		return f.read()

def deps_add(filename):
	#print(f"{filename=} {filename in deps}")
	if filename in deps:
		return False
	deps[filename] = 1
	if depsfile:
		print("\t" + filename + " \\", file=depsfile)
	return True

### try stripping executables and libraries before loading them
def try_strip(filename):
	data = readfile(filename)
	if no_strip \
	or filename.endswith(".ko") \
	or data[0:4] != b'\x7fELF':
		return data

	try:
		with NamedTemporaryFile() as tmp:
			rc = system("strip", "-o", tmp.name, filename)
			if rc.returncode == 0:
				# strip succeeded, use it instead
				new_data = readfile(tmp.name)
				#print(f"{filename}: {len(data)} -> {len(new_data)}")
				return new_data
	except Exception as e:
		# something went wrong with strip, use the file as is
		pass

	return data

def envsubst(match):
	varname = match.group(1)
	value = os.environ.get(varname, None)
	#print(f"$({varname})='{value}'", file=sys.stderr)
	if not value:
		raise ValueError(f"{varname}: undefined variable")
	return value

def process_line(line):
	# do any environment substitution
	line = re.sub(r"\$\((.*?)\)", envsubst, line)

	words = re.split(r'\s+', line)
	if words[0] == "symlink":
		# the filename is the second part, the data in that
		# file is the path to the real file
		cpio.symlink(words[2], words[1])
		return
	if words[0] == "mkdir":
		mode = 0o777
		if len(words) > 2:
			mode = int(words[2], 8)
		cpio.mkdir(words[1], mode)
		return
	if words[0] == "mknod":
		# mknod path [c|b] major minor
		cpio.mknod(words[1], words[2], words[3], words[4])
		return
	if words[0] == "exec":
		# command line stuff to be read by init
		# exec /init.d/10-module cmd [args...]
		cpio.add(words[1], data=words[2:])
		return

	# normal file
	filename = words[0]
	dest = bindir
	if len(words) > 1:
		dest = words[1]

	# - marks optional files; do not die if they do not exist
	# ! marks forced files; do not die if their dependencies do not exist
	optional = False
	force = False
	if filename[0] == '-':
		optional = True
		filename = filename[1:]
	if filename[0] == '!':
		force = True
		filename = filename[1:]

	# if a relative file name is provided, update it with the cwd
	if filename[0] != '/':
		filename = os.path.join(cwd, filename)

	if optional and not exists(filename):
		print(filename + ": skipping")
		return

	size = 0

	if filename.endswith('/'):
		# recursive directory copy
		olddir = os.path.normpath(filename)

		if len(words) > 2:
			pattern = words[2]
		else:
			# skip editor temp files .*.swp
			pattern = r'\.[^.]*\.sw[po]$'

		for root, dirs, files in os.walk(olddir):
			# strip the non-relative path part
			# is there a better way to do this?
			newroot = os.path.join(dest, os.path.relpath(root, start=olddir))
			cpio.mkdir(newroot)
			for f in files:
				oldname = os.path.join(root, f)
				newname = os.path.join(newroot, f)

				# filter any that match the exclusion pattern
				if re.search(pattern, oldname):
					continue

				data = readfile(oldname)
				cpio.add(newname, data=data)

				if deps_add(oldname):
					size += len(data)

		if verbose:
			print("%s: %d (recursive)" % (olddir, size))


		# ldd is *not* done on recursive directories
		return

	# strip the local path
	dest_filename = os.path.basename(filename)

	data = try_strip(filename)
	cpio.add(dest + "/" + dest_filename, data=data, mode=filemode(filename))

	if deps_add(filename):
		size += len(data)

	file_deps = system("ldd", filename)
	if file_deps.returncode != 0 \
	or file_deps.stderr == b'\tnot a dynamic executable\n':
		if verbose:
			print("%s: %d" % (filename, size))
		return

	dep_size = 0
	for dep in file_deps.stdout.split(b'\n'):
		dep = dep.decode('utf-8')
		if re.search(r'not found', dep):
			print("MISSING LIBRARY: ", dep, file=sys.stderr)
			# do not die on optional missing dependencies
			if force:
				continue

		dep = re.sub(r'^.* => ', '', dep)
		dep = re.sub(r'^\s*', '', dep)
		lib = dep.split(' ')[0]

		if lib == '' or re.match(r'^linux-vdso', lib):
			continue
		
		# strip the local path and put the library in the libdir
		dest_lib = os.path.basename(lib)
		data = try_strip(lib)
		cpio.add(libdir + "/" + dest_lib, data=data)
		if deps_add(lib):
			dep_size += len(data)

	if verbose:
		print("%s: %d (deps %d)" % (filename, size, dep_size))

parser = argparse.ArgumentParser()
parser.add_argument('files',
	nargs='+',
	help="Configuration files to parse")
parser.add_argument('-o', '--output',
	dest='cpio', type=str,
	default='-',
	help="Output of the cpio file")
parser.add_argument( '-v', '--verbose',
	dest='verbose', action='store_true',
	help="Log every file included")
parser.add_argument( '-c', '--xz',
	dest='xz', action='store_true',
	help="Compress the initrd with xz")
parser.add_argument('-r', '--relative',
	dest='relative', action='store_true',
	help="Use the path of the config for relative path files")
parser.add_argument('-M', '--deps',
	dest='deps', type=str,
	help="Output dependencies for Makefile rebuilds")
parser.add_argument('--no-strip',
	dest='no_strip', action='store_true',
	help="Do not strip libraries and executables")
parser.add_argument('--add',
	dest='extra', type=str, action='append',
	help="Extra binaries to be included")

args = parser.parse_args()
verbose = args.verbose
#cpio.verbose = verbose
initrd_filename = args.cpio
no_strip = args.no_strip

if args.deps:
	depsfile = open(args.deps, "w")
	print("# Autogenerated dependencies for " + initrd_filename, file=depsfile)
	print(initrd_filename + ": \\", file=depsfile)


for filename in args.files:
	if args.relative:
		cwd = os.path.split(filename)[0]
	if cwd == '':
		cwd = '.'

	with open(filename, "r") as f:
		linenum = 0
		for line in f.readlines():
			linenum += 1

			# strip any commands and skip blank lines
			line = re.sub(r'\n$', '', line)
			line = re.sub(r'^\s*', '', line)
			line = re.sub(r'\s*#.*$','', line)
			if len(line) == 0:
				continue

			#print("%s:%d: %s" % (filename, linenum, line))

			try:
				process_line(line)
			except Exception as e:
				print("%s:%d: failed '%s' (cwd=%s)" % (filename, linenum, line, cwd), e, file=sys.stderr)
				print(traceback.format_exc())
				exit(1)

# restore the cwd to the current directory
cwd = '.'
if args.extra:
	for line in args.extra:
		try:
			process_line(line)
		except Exception as e:
			print("%s:%d: failed '%s'" % (filename, linenum, line), e, file=sys.stderr)
			print(traceback.format_exc())
			exit(1)

if verbose:
	print("**** all files added")

compressed = args.xz or initrd_filename.endswith('.xz')
image = cpio.tobytes(compressed)

if verbose:
	print("**** writing cpio image ")

if initrd_filename == '-':
	sys.stdout.buffer.write(image)
else:
	with open(initrd_filename, "wb") as initrd:
		initrd.write(image)