#!/usr/bin/env bash

# Alternative to unmkinitramfs. Extract initrd with one or more cpio archives
# Newer unmkinitramfs extracts a cpio containing compressed modules into the main cpio.
# This script extracts each cpio separately, mirroring the original layout.
# Support initrd with up to 3 uncompressed and 1 compressed cpio.

ARCHIVE=$1

if [ -z "$ARCHIVE" ]; then
	echo "USAGE: extract_initrd /path/to/original/initrd"
	exit 1
else 
	echo "Archive is $ARCHIVE"
fi

if [ -d "./extracted" ]; then rm -rf ./extracted; fi

#### Set some functions to use later #############

get_compression () {
	COMPRESSION=$(file $ARCHIVE | egrep -o "Zstandard|XZ|gzip|ASCII")
	echo "Archive is $COMPRESSION"
}


decompress  () {

	if [ "$COMPRESSION" = "XZ" ]; then
		mv $ARCHIVE $ARCHIVE.xz
		unxz $ARCHIVE.xz
		#$ARCHIVE is already set. Move on to extract_archive.
		return

	elif [ "$COMPRESSION" =  "Zstandard" ]; then
		mv $ARCHIVE $ARCHIVE.zst
		unzstd $ARCHIVE.zst
		return

	elif [ "$COMPRESSION" =  "gzip" ]; then
		mv $ARCHIVE $ARCHIVE.gz
		gunzip $ARCHIVE.gz
		return
	else
		echo "Error.. Unsupported compression format $COMPRESSION"
		exit 1
	fi
}


extract_archive () { 

	cpio -idv <../../$ARCHIVE &> /tmp/blocknum
	# Get the block count number
	sed -i '$!d' /tmp/blocknum
	sed -i 's: .*::' /tmp/blocknum
}


#### Sort archive 1 (it may be compressed) ########

# The blocknum is for cpio to skip, first cpio skips zero
echo "0" >/tmp/blocknum
SKIP=$(cat /tmp/blocknum)

# make a copy, skip the blocknum
dd if=$ARCHIVE skip=$SKIP of=cpio1.img

# Update $ARCHIVE to work on
ARCHIVE=cpio1.img

get_compression

if [ "$COMPRESSION" = "ASCII" ]; then

	mkdir -p extracted/early
	pushd extracted/early
	extract_archive
	popd
else
	decompress
	mkdir -p extracted/main
	pushd extracted/main
	extract_archive
	popd
	rm -f cpio*
	rm -f /tmp/blocknum
	# There was only one main cpio. Move it out of the subdir.
	cp -a extracted/main/* extracted/
	rm -rf extracted/main/
	exit 0
fi

#### Sort archive 2 (it may be compressed) ########

SKIP=$(cat /tmp/blocknum)
dd if=$ARCHIVE skip=$SKIP of=cpio2.img
ARCHIVE=cpio2.img
#rm -f cpio1.img

get_compression

if [ "$COMPRESSION" = "ASCII" ]; then

	mkdir -p extracted/early2
	pushd extracted/early2
	extract_archive
	popd
else
	decompress
	mkdir -p extracted/main
	pushd extracted/main
	extract_archive
	popd
	rm -f cpio*
	rm -f /tmp/blocknum
	exit 0
fi

#### Sort archive 3 (it may be compressed) ########

SKIP=$(cat /tmp/blocknum)
dd if=$ARCHIVE skip=$SKIP of=cpio3.img
ARCHIVE=cpio3.img
get_compression

if [ "$COMPRESSION" = "ASCII" ]; then

	mkdir -p extracted/early3
	pushd extracted/early3
	extract_archive
	popd
else
	decompress
	mkdir -p extracted/main
	pushd extracted/main
	extract_archive
	popd
	rm -f cpio*
	rm -f /tmp/blocknum
	exit 0
fi

#### Sort archive 4 (it will be compressed) ######

SKIP=$(cat /tmp/blocknum)
dd if=$ARCHIVE skip=$SKIP of=cpio4.img
ARCHIVE=cpio4.img
get_compression

if [ "$COMPRESSION" = "ASCII" ]; then

	pushd extracted/main
	extract_archive
	popd
else
	decompress
	mkdir -p extracted/main
	pushd extracted/main
	extract_archive
	popd
	rm -f cpio*
	rm -f /tmp/blocknum
	exit 0
fi

###################################################

exit 0
