#!/bin/ash
# Split Fatdog "humoungous initrd" into small initrd and base sfs
# Copyright James Budiono 2012, 2013, 2016-2019, 2022
# License: GNU GPL Version 3 or later
#
# $1-initrd, $2-outputdir, $3-tmpdir, $4-basesfs-default-path, $5-base-sfs
#
# Jan 2018 - step: support early microcode update; gettexting
# Sep 2022 - james: add -mini flag

### configuration
TEXTDOMAIN=fatdog OUTPUT_CHARSET=UTF-8
export TEXTDOMAIN OUTPUT_CHARSET
. gettext.sh
# modules required for initrd (apart from MIN_MODULES)
BASE_MODULES="${BASE_MODULES:-cryptoloop nbd cifs aufs usb-storage \
ohci uhci ehci xhci-pci xhci-hcd xhci-pci xhci-plat-hcd \
xts hid-generic dm-crypt dm-snapshot sg \
crypto_simd cryptd glue_helper algif_skcipher af_alg}"
MODULE_SFS=${MODULE_SFS:-kernel-modules.sfs}
BASE_SFS=${BASE_SFS:-fd64.sfs}
KERNEL_MODULE_COMPRESSION=${KERNEL_MODULE_COMPRESSION:-'-comp xz'}
BASESFS_COMPRESSION=${BASESFS_COMPRESSION:-'-comp gzip -Xcompression-level 4'}
[ "$BOOTSTATE_PATH" ] && . $BOOTSTATE_PATH   			# BASE_SFS_PATH
[ "$BASE_SFS_PATH" ] && BASE_SFS=${BASE_SFS_PATH##*/}	# updated by $5

### set runtime variables
SPLIT_MODE=${SPLIT_MODE:-micro} # mini or micro, default is micro for compatibility reasons
TMPDIR=/tmp
INITRD=
OUTPUTDIR=
BASE_SFS_DEFAULT_PATH=

##### helpers #####

### update BASE_MODULES with list of required modules and their dependencies
# input: $1-BASE_MODULES, $2-MODULES_DEP returns: updated BASE_MODULES
get_modules_dep() {
	local OLD_MODULES GREP_PARAM MODULES_DEP
	BASE_MODULES=$1 MODULES_DEP=$2
	while [ ! "$OLD_MODULES" = "$BASE_MODULES" ]; do
		OLD_MODULES=$BASE_MODULES
		get_grep_param $BASE_MODULES
		set -- $(grep -E "$GREP_PARAM" $MODULES_DEP)
		while [ $1 ]; do
			case $1 in
				*:) ;;
				*)	new=${1##*/}; new=${new%.ko}; 
					case $BASE_MODULES in 
						*$new*) ;;
						*) BASE_MODULES="$BASE_MODULES $new"
					esac ;;
			esac
			shift
		done
	done
}
# $*-modules, returns GREP_PARAM
get_grep_param() {
	GREP_PARAM=
	while [ $1 ]; do
		GREP_PARAM="$GREP_PARAM|/$1.ko:"
		shift
	done
	GREP_PARAM=${GREP_PARAM#|}
}

abort_abort_abort() {
	echo "$1"
	cd "$CWD"
	rm -rf $INITRD_EXTRACT
	exit 1
}

usage() {
	cat << EOF
$(gettext "Usage"): $0 [-mini|-micro] source-initrd output-dir [tmpdir] [basesfs-default-path] [base-sfs]

$(eval_gettext 'Split Fatdog standard initrd to "micro" or "mini" initrd.

[-mini]  remove basesfs, but keep all kernel modules.
[-micro] remove basesfs, keep only specified BASE_MODULES, and move
         the rest of the modules to an updated basesfs.

If none specified, "micro" is assumed unless SPLIT_MODE envvar exists,
then its value will be used.

[tmpdir] is the work directory. By default /tmp is used, but splitting
requires a lot of space. If your /tmp is small, it may overflow silently
and causes broken / un-bootable initrd.
In this case, it is advisable to use a work directory elsewhere,
especially when doing "micro" splitting.
[tmpdir] must be on a Linux filesystem.

[basesfs-default-path] The default basesfs location, to be inserted into
system-init. Default is blank, which means, "do not change anything".

[base-sfs] Name of the basesfs. Default is $BASE_SFS.')
EOF
	exit
}

### align sfs to 256K so it is usable by qemu directly
# $1-path to sfs
padsfs() {
	ORIGSIZE=$(stat -c %s "$1")
	BLOCKS256K=$(( ((ORIGSIZE/1024/256)+1) ))
	dd if=/dev/zero of="$1" bs=256K seek=$BLOCKS256K count=0
}

############ main ############

# 1. process args
[ $(id -u) -ne 0 ] && echo "$(gettext "You must be root to use this.")" && exit

case "$1" in
	-mini)     SPLIT_MODE=mini;  shift ;;
	-micro)    SPLIT_MODE=micro; shift ;;
	-h|--help) usage ;;
esac

INITRD="$1"
OUTPUTDIR="$2"
[ "$3" ] && TMPDIR="$(readlink -f "$3")"
BASE_SFS_DEFAULT_PATH="$4"
[ "$5" ] && BASE_SFS="${5##*/}"

# 2. initialise working variables
! [ "$INITRD" -a "$OUTPUTDIR" ] && usage
INITRD="$(readlink -f "$INITRD")"
OUTPUTDIR="$(readlink -f "$OUTPUTDIR")"
[ ! -e "$INITRD" ] && echo "$(gettext "initrd does not exist")" && exit
[ ! -e "$OUTPUTDIR" ] && mkdir -p "$OUTPUTDIR"
CWD="$(pwd)"

# 3. unpack initrd
INITRD_EXTRACT=$(mktemp -dp $TMPDIR initrd-XXXXXXXX)
cd $INITRD_EXTRACT

case $(file -b --mime-type "$INITRD") in
	*gzip) TYPE="gzip" && zcat "$INITRD" | cpio -i ;;
	*xz) TYPE="xz" && xzcat "$INITRD" | cpio -i ;;
	*bzip2) TYPE="bzip2" && bzcat "$INITRD" | cpio -i ;;
	*) TYPE="none" && cat "$INITRD" | cpio -i ;;
esac || exit

# 4. Check integrity - make sure both kernel modules and basesfs are there
[ -e "$INITRD_EXTRACT/$MODULE_SFS" ] || abort_abort_abort "$(eval_gettext "Cannot find \$MODULE_SFS in initrd, exiting.")"
[ -e "$INITRD_EXTRACT/$BASE_SFS" ] || abort_abort_abort "$(eval_gettext "Cannot find \$BASE_SFS in initrd, exiting.")"

####### MICRO SPLIT
if [ $SPLIT_MODE = "micro" ]; then

# A1. extract kernel modules
MODULE_EXTRACT=$(mktemp -dp $TMPDIR module-XXXXXXXX); rmdir $MODULE_EXTRACT
unsquashfs -d $MODULE_EXTRACT $INITRD_EXTRACT/$MODULE_SFS
rm $INITRD_EXTRACT/$MODULE_SFS

# A2. extract base sfs (can't append directly - too bad :( - means we need more space)
BASE_EXTRACT=$(mktemp -dp $TMPDIR base-XXXXXXXX); rmdir $BASE_EXTRACT
unsquashfs -d $BASE_EXTRACT $INITRD_EXTRACT/$BASE_SFS
rm $INITRD_EXTRACT/$BASE_SFS

# A3. merge them
cp -af $MODULE_EXTRACT/* $BASE_EXTRACT

# A4. depmod the new base sfs
KERNEL_VERSION=$(ls $MODULE_EXTRACT/lib/modules)
depmod -a -b $BASE_EXTRACT $KERNEL_VERSION

# A5. repack base sfs - straight to target
rm -f "$OUTPUTDIR/$BASE_SFS"
mksquashfs $BASE_EXTRACT "$OUTPUTDIR/$BASE_SFS" $BASESFS_COMPRESSION
[ -z "$NOPAD" ] && padsfs "$OUTPUTDIR/$BASE_SFS"
rm -rf $BASE_EXTRACT

# A6. delete all modules except what we need
# A6.1 get all modules and their dependencies
get_modules_dep "$BASE_MODULES" $(echo $MODULE_EXTRACT/lib/modules/*/modules.dep)
#echo $BASE_MODULES

# A6.2 convert this to find params
FIND_PARAM=""
for p in $BASE_MODULES; do
	FIND_PARAM="$FIND_PARAM -a ! -name $p.ko"
done
FIND_PARAM=${FIND_PARAM#*-a}

# A6.3 clean it - except the files we need
find $MODULE_EXTRACT/* \( -type f $FIND_PARAM -o -type l \) -delete
find $MODULE_EXTRACT/* -type d -empty -delete

# A7. depmod then repack module sfs
depmod -a -b $MODULE_EXTRACT $KERNEL_VERSION
mksquashfs $MODULE_EXTRACT $INITRD_EXTRACT/$MODULE_SFS $KERNEL_MODULE_COMPRESSION
rm -rf $MODULE_EXTRACT


####### MINI SPLIT
elif [ $SPLIT_MODE = "mini" ]; then

# B1. Simply basesfs to output dir
mv "$INITRD_EXTRACT/$BASE_SFS" "$OUTPUTDIR/$BASE_SFS"

####### UNKNOW SPLIT
else

abort_abort_abort "$(eval_gettext "SYSTEM error: split mode --> $SPLIT_MODE <-- is NOT supported!")"

fi


### SHARED finalisations

# 5. change default base_sfs from "initrd" to "local" so that it will search basesfs on local drives
#    while we're at it, increase the default delay to 5 seconds too
sed -i -e '
s/BASE_SFS_DEFAULT=initrd/BASE_SFS_DEFAULT=local/
s/DEFAULT_DEVICE_DELAY=0/DEFAULT_DEVICE_DELAY=5/
' $INITRD_EXTRACT/sbin/system-init 
[ $BASE_SFS_DEFAULT_PATH ] &&
sed -i -e "
s|BASE_SFS_DEFAULT_PATH=.*|BASE_SFS_DEFAULT_PATH=$BASE_SFS_DEFAULT_PATH|
" $INITRD_EXTRACT/sbin/system-init 

# 6. repack initrd - don't compress it due to early microcode update and fatdog*-installer.sh
cd $INITRD_EXTRACT
find . | cpio -o -H newc > "$OUTPUTDIR/initrd"

# 7. final cleanup
cd "$CWD"
rm -rf $INITRD_EXTRACT  
#rox -d "$OUTPUTDIR"
