#!/bin/sh

set -eu

usage()
{
	cat << EOF

Usage: unmkinitramfs [-v] initramfs-file directory

Options:
  -v   Display verbose messages about extraction

See unmkinitramfs(8) for further details.

EOF
}

usage_error()
{
	usage >&2
	exit 2
}

# Extract a compressed cpio archive
xcpio()
{
	archive_uncomp="$1"
	archive="$2"
	dir="$3"
	shift 3

	{
		cat "$archive_uncomp"
		if gzip -t "$archive" >/dev/null 2>&1 ; then
			gzip -c -d "$archive"
		elif zstd -q -c -t "$archive" >/dev/null 2>&1 ; then
			zstd -q -c -d "$archive"
		elif xzcat -t "$archive" >/dev/null 2>&1 ; then
			xzcat "$archive"
		elif lz4cat -t < "$archive" >/dev/null 2>&1 ; then
			lz4cat "$archive"
		elif bzip2 -t "$archive" >/dev/null 2>&1 ; then
			bzip2 -c -d "$archive"
		elif lzop -t "$archive" >/dev/null 2>&1 ; then
			lzop -c -d "$archive"
		# Ignoring other data, which may be garbage at the end of the file
		fi
	} | (
		if [ -n "$dir" ]; then
			mkdir -p -- "$dir"
			cd -- "$dir"
		fi
		cpio "$@"
	)
}

# Read bytes out of a file, checking that they are valid hex digits
readhex()
{
	dd < "$1" bs=1 skip="$2" count="$3" 2> /dev/null | \
		LANG=C grep -E "^[0-9A-Fa-f]{$3}\$"
}

# Check for a zero byte in a file
checkzero()
{
	dd < "$1" bs=1 skip="$2" count=1 2> /dev/null | \
		LANG=C grep -q -z '^$'
}

# Split an initramfs into archives and run cpio/xcpio to extract them
splitinitramfs()
{
	initramfs="$1"
	dir="$2"
	shift 2

	# Ensure this exists so we can use it unconditionally later
	touch "$tempdir/main-uncomp.cpio"

	count=0
	start=0
	while true; do
		# There may be prepended uncompressed archives.  cpio
		# won't tell us the true size of these so we have to
		# parse the headers and padding ourselves.  This is
		# very roughly based on linux/lib/earlycpio.c
		end=$start
		while true; do
			headoff=$end
			magic="$(readhex "$initramfs" $end 6)" || break
			test "$magic" = 070701 || test "$magic" = 070702 || break
			namesize=$((0x$(readhex "$initramfs" $((end + 94)) 8)))
			filesize=$((0x$(readhex "$initramfs" $((end + 54)) 8)))
			nameoff=$((end + 110))
			end=$(((nameoff + namesize + 3) & ~3))
			end=$(((end + filesize + 3) & ~3))

			# Check for EOF marker.  Note that namesize
			# includes a null terminator.
			if [ $namesize = 11 ] \
			   && name="$(dd if="$initramfs" bs=1 skip=$nameoff count=$((namesize - 1)) 2> /dev/null)" \
			   && [ "$name" = 'TRAILER!!!' ]; then
				# There might be more zero padding
				# before the next archive, so read
				# through all of it.
				while checkzero "$initramfs" $end; do
					end=$((end + 4))
				done
				break
			fi
		done

		if [ $end -eq $start ]; then
			break
		fi

		# Check whether this should be treated as an "early"
		# or "main" initramfs.  Currently all filenames the
		# kernel looks for in an early initramfs begin with
		# kernel/ subdirectory, but we should never create
		# this in the main initramfs.
		if dd < "$initramfs" skip=$start count=$((end - start)) \
			iflag=skip_bytes,count_bytes 2> /dev/null |
		   cpio -i --list 2> /dev/null |
		   grep -q ^kernel/; then
			# Extract to early, early2, ... subdirectories
			count=$((count + 1))
			if [ $count -eq 1 ]; then
				subdir=early
			else
				subdir=early$count
			fi
			dd < "$initramfs" skip=$start count=$((end - start)) \
				iflag=skip_bytes,count_bytes 2> /dev/null |
			(
				if [ -n "$dir" ]; then
					mkdir -p -- "$dir/$subdir"
					cd -- "$dir/$subdir"
				fi
				cpio -i "$@"
			)
		else
			# Append to main-uncomp.cpio, excluding the
			# trailer so cpio won't stop before the
			# (de)compressed part.
			dd < "$initramfs" skip=$start \
				count=$((headoff - start)) \
				iflag=skip_bytes,count_bytes \
				>> "$tempdir/main-uncomp.cpio" 2> /dev/null
		fi

		start=$end
	done

	# Split out final archive if necessary
	if [ "$end" -gt 0 ]; then
		subarchive="$tempdir/main-comp.cpio"
		dd < "$initramfs" skip="$end" iflag=skip_bytes 2> /dev/null \
			> "$subarchive"
	else
		subarchive="$initramfs"
	fi

	# If we found an early initramfs, extract main initramfs to
	# main subdirectory.  Otherwise don't use a subdirectory (for
	# backward compatibility).
	if [ "$count" -gt 0 ]; then
		subdir=main
	else
		subdir=.
	fi

	xcpio "$tempdir/main-uncomp.cpio" "$subarchive" \
		"${dir:+$dir/$subdir}" -i "$@"
}

extract_with_3cpio()
{
	local initramfs="$1"
	local dir="$2"
	shift 2

	3cpio --examine "$initramfs" > "$tempdir/parts"
	if test "$(wc -l < "$tempdir/parts")" = 1; then
		3cpio --extract -C "$dir" "$@" "$initramfs"
		return
	fi

	local early_count=0
	local start=
	mkdir -p "$dir"
	printf "end\tend\n" >> "$tempdir/parts"
	while read -r end _; do
		[ "$end" != "end" ] || end=""
		if [ "$start" = "" ]; then
			start="$end"
			continue
		fi
		dd < "$initramfs" skip="$start" ${end:+count=$((end - start))} \
				iflag=skip_bytes,count_bytes > "$tempdir/cpio" 2> /dev/null
		if 3cpio --list "$tempdir/cpio" | grep -q ^kernel/; then
			# Extract to early, early2, ... subdirectories
			early_count=$((early_count + 1))
			if [ $early_count -eq 1 ]; then
				subdir=early
			else
				subdir=early$early_count
			fi
			3cpio --extract -C "$dir/$subdir" "$@" "$tempdir/cpio"
		else
			# Append to main.cpio
			cat "$tempdir/cpio" >> "$tempdir/main.cpio"
		fi
		start="$end"
	done < "$tempdir/parts"

	if [ "$early_count" -eq 0 ]; then
		maindir="$dir"
	else
		maindir="$dir/main"
	fi
	3cpio --extract -C "$maindir" "$@" "$tempdir/main.cpio"
}

OPTIONS=$(getopt -o hv --long help,list,verbose -n "$0" -- "$@") || usage_error

cpio_opts="--preserve-modification-time --no-absolute-filenames --quiet"
mode=extract
verbose=
expected_args=2
eval set -- "$OPTIONS"

while true; do
	case "$1" in
        -h|--help)
		usage
		exit 0
	;;
	--list)
		# For lsinitramfs
		mode=list
		cpio_opts="${cpio_opts:+${cpio_opts} --list}"
		expected_args=1
		shift
	;;
	-v|--verbose)
		verbose=1
		cpio_opts="${cpio_opts:+${cpio_opts} --verbose}"
		shift
	;;
	--)
		shift
		break
	;;
	*)
		echo "Internal error!" >&2
		exit 1
	esac
done

if [ $# -ne $expected_args ]; then
	usage_error
fi

tempdir="$(mktemp -d "${TMPDIR:-/var/tmp}/unmkinitramfs_XXXXXX")"
trap 'rm -rf "$tempdir"' EXIT

if command -v 3cpio >/dev/null 2>&1; then
	if test "$mode" = list; then
		3cpio --list ${verbose:+--verbose} "$1"
	else
		extract_with_3cpio "$1" "$2" ${verbose:+--verbose}
	fi
else
	# shellcheck disable=SC2086
	splitinitramfs "$1" "${2:-}" $cpio_opts
fi