#!/bin/bash
# Create ext4 image out of file tree
# Copyright (c) 2022-2023 Vitaly Chikunov <vt@altlinux.org>
# Inspired by the work of Leonid Krivoshein <klark@altlinux.org> on usermode-fs-tools.

set -efu
export DEBUGFS_PAGER=__none__
export LC_ALL=C
TMP=$(mktemp -d)
# shellcheck disable=SC2064
trap "rm -rf '$TMP'" 0

show_usage()
{
	echo "Usage: $0 [OPTIONS...] IMAGE"
	echo "Create ext4 IMAGE from rootfs (or DIR)."
	echo
	echo "  --append       Append files to existing raw IMAGE."
	echo "  --exclude=     Mask to exclude from dir (default: /tmp/*)"
	echo "  --fs=          Specify fs type: ext4 (default), ext3, ext2."
	echo "  --mkfs=        Specify additional augments for mkfs."
	echo "  --qcow2        After creation convert IMAGE to qcow2 format"
	echo "  --part         Create disk partition (gpt or mbr)."
	echo "  --realuids     Set UID/GID to real values instead of root:root."
	echo "                 Default when libfakeroot is loaded (as is under rooter)."
	echo "  --size=        Size of IMAGE (Kb) (default: double of du of DIR)"
	echo "  --source=      Source directory to copy into IMAGE (default: /)"
	echo "  --special      Copy special files (block, character devices, pipes)"
	echo "  --volume=      EXT4 volume name (default: basename of IMAGE)"
	echo "  -v, --verbose  Increase verbosity"
	echo
	echo "Informational options:"
	echo "  --cat=FILE     Cat FILE from IMAGE"
	echo "  --cp A B       Copy file/dir 'A' from the IMAGE into file/dir 'B'"
	echo "  --ls=DIR       Do not create an IMAGE but 'ls -l DIR' in it."
	echo "  --stat=FILE    Run stat on FILE in IMAGE"
	exit "$@"
}

warning()
{
	local w prefix="$(basename "$0"): "
	for w do
		echo >&2 "$prefix$w"
		prefix=
	done
}
fatal()
{
	warning "$@"
	exit 1
}

unset exclude
unset src
size=
image=
declare -i verbose=0
fs_type=ext4
create=y
mkfs=
unset al_all ls_path cat_path stat_path qcow2 copy_mode part extract

# Seems to be true under rooter and scriptlets.
[[ "${LD_PRELOAD-}" =~ libfakeroot.so ]] && realuids=y || realuids=n

for opt; do # opt
	shift
	arg=${opt##*=}
	case $opt in
		--source=*)     src+=( $arg ) ;;
		--exclude=*)    exclude+=( "$arg" ) ;;
		--volume=*)     volume=$arg ;;
		--size=*)       size=$arg ;;
		--fs=*)         fs_type=$arg ;;
		--ls=*)         ls_path=$arg ;;
		--ls)           ls_all=y ;;
		--cat=*)        cat_path=$arg ;;
		--cp)           copy_mode=y ;;
		--stat=*)       stat_path=$arg ;;
		--qcow2)        qcow2=y ;;
		--special)      special=y ;;
		--realuids)     realuids=y ;;
		--no-realuids)  unset realuids ;;
		--append)       unset create ;;
		--part)         part= ;;
		--part=*)       part=$arg ;;
		--extract)      extract=y ;;
		--mkfs=*)       mkfs+=" $arg" ;;
		-v | --verbose) verbose+=1 ;;
		-h | --help)    show_usage ;;
		-*)             fatal "Error: Unknown option $opt" ;;
		*)              set -- "$@" "$opt" ;;
	esac
done

[ -n "${1-}" ] || { warning "Specifying IMAGE is required." ""; show_usage 1; }
image="$1"; shift

[ $verbose -gt 0 ] && SETX="set -x" || SETX=:
[ $verbose -gt 0 ] && VERBOSE=-v || VERBOSE=
[ $verbose -gt 0 ] && QUIET= || QUIET=-q
[ $verbose -gt 1 ] && NULL="1" || NULL=/dev/null

lsdir()
{
	set -- "$@"
	while [ $# -gt 0 ]; do
		local d=$1
		shift
		echo "ls -l \"$d\"" >&3
		local a dirs=()
		while IFS= read -r a <&4; do
			[ -n "$a" ] || break
			[[ $a =~ ^\ *[0-9]+\ +([0-9]+)\ \([0-9]+\)\ +[0-9]+\ +[0-9]+\ +[0-9]+\ .{17}\ (.*)$ ]] || continue
			md=${BASH_REMATCH[1]}
			na=${BASH_REMATCH[2]}
			a="${a%"$na"}"
			case "$na" in . | .. | '') continue ;; esac
			case "$md" in 4????) na+='/'; dirs+=( "$d$na" ) ;; esac
			printf "%s%s\n" "$a" "$d$na"
		done
		set -- "$@" "${dirs[@]}"
	done
}

is_filesystem() {
	local TYPE=
	local $(/sbin/blkid -o export -- "$image")
	if [ -v extract ]; then
		[ -n "${PTTYPE-}" ] || fatal "'$image' does not have PTTYPE."
		local newimage=$image.extract
		trap "rm $VERBOSE -- ${newimage@Q}" 0
		($SETX; dd if="$image" of="$newimage" conv=sparse skip=2048)
		# `qemu-img dd` does not support conv=sparse.
		image=$newimage
		local $(/sbin/blkid -o export -- "$image")
		[[ "$TYPE" =~ ^ext[234]$ ]] || fatal "Extracted image from '$image' does not have ext4."
	elif [[ ! "$TYPE" =~ ^ext[234]$ ]]; then
		fatal "'$image' is not raw ext4 filesystem image." \
			"Use --extract for low speed fs extraction."
	fi
}
if [ -v ls_path ]; then
	is_filesystem
	/sbin/debugfs "$image" -R "ls -l $ls_path"
	exit
elif [ -v ls_all ]; then
	is_filesystem
	# Avoid wrapping of commands echo by readline
	export COLUMNS=9999
	mkfifo "$TMP/in" "$TMP/out"
	/sbin/debugfs "$image" <"$TMP/in" >"$TMP/out" &
	exec 3>"$TMP/in" 4<"$TMP/out"
	lsdir /
	exit
elif [ -v cat_path ]; then
	is_filesystem
	/sbin/debugfs "$image" -R "cat $cat_path" 2>/dev/null
	exit
elif [ -v stat_path ]; then
	is_filesystem
	/sbin/debugfs "$image" -R "stat $stat_path"
	exit
elif [ -v copy_mode ]; then
	is_filesystem
	[ $# -gt 2 ] && fatal "Too much arguments for cp mode"
	info=$(/sbin/debugfs "$image" -R "stat $1" 2>&1 | head -2)
	if [[ "$info" =~ File.not.found ]]; then
		echo "Error: $info" >&2
		exit 1
	elif [[ "$info" =~ Type:.directory ]]; then
		mkdir -p "$2"
		($SETX; /sbin/debugfs "$image" -R "rdump \"$1\" \"$2\"")
	else
		($SETX; /sbin/debugfs "$image" -R "dump \"$1\" \"$2\"")
	fi
	exit
fi
[ $# -gt 0 ] && fatal "Unknown argument(s): $*"

if [ ! -v src ]; then
	src=( "/" )
fi
if [ ! -v exclude ]; then
	# /tmp is on the same tmpfs so -xdev will not work.
	exclude=( "/tmp/*" )
fi

if [ -v create ]; then
	# shellcheck disable=SC2068
	total=$(set -f; du -scx ${exclude[@]/#/--exclude } "${src[@]}" 2>/dev/null | tail -1)
	total=${total%%[[:space:]]*}
	[ $verbose -gt 0 ] && echo >&2 "Source directory='${src[*]}' size=$total kb"
	unit=1024
	shopt -s nocasematch
	case "$size" in
		*B) unit=1/$unit ;;
		*K | *[0-9]) unit=1 ;;
		*M) unit=$((unit*unit)) ;;
		*G) unit=$((unit*unit*unit)) ;;
		*T) unit=$((unit*unit*unit*unit)) ;;
		*[A-Z]) fatal "Unknown unit size: $size" ;;
	esac
	shopt -u nocasematch
	sz=${size%[BKMGTbkmgt]}
	case "$sz" in
		+*)   size=$((total + ${sz#+} * $unit)) ;;
		*%)   size=$((total * ${sz%'%'} / 100)) ;;
		'*'*) size=$((total * ${sz#'*'})) ;;
		'')   size=$((total * 3)) ;;
		*)    size=$((sz * $unit)) ;;
	esac
	[ "$size" -lt 10000 ] && size="10000"
	unset unit sz
	if [ "${realuids-}" = 'n' ]; then
		unset realuids
		warning "libfakeroot is not detected so realuids mode is off, use --realuids to enable."
	fi

	[ $verbose -gt 0 ] && echo >&2 "Creating ${fs_type:-ext4} image '$image' size=$size" \
		${realuids+(realuids)} ${special+(special)} ${volume+volume=$volume}
	rm -f "$image"
	[ -v volume ] || volume="$(basename "$image")"
	# When verity enabled on ppc64 block size shall be equal to page size (which is 64K there).
	if [[ "$mkfs" =~ verity ]]; then
		bs=$(getconf PAGE_SIZE)
		if [ "$bs" != 4096 ]; then
		       # Last option takes precedence, thus prepend our blocksize setting so user can override it.
		       mkfs="-b $bs $mkfs"
		fi
		unset bs
	fi
	# mke2fs does not create ext4 by default (even if `-t ext4` is passed),
	# requiring appropriate additional options (like enabling journal, etc),
	# but mkfs.ext4 does the right thing.
	($SETX; /sbin/mkfs.$fs_type ${VERBOSE:--q} -F -L "$volume" -m0 $mkfs -- "$image" "$size")
fi
[ -f "$image" ] || fatal "Raw image '$image' not found."

set_metadata()
{
	local t=$1 md=$2 us=$3 gr=$4 ts=$5 f=$6
	local ty

	if [ -v realuids ]; then
		[ "$us" -eq 0 ] || printf 'sif "%s" uid %s\n' "$f" "$us"
		[ "$gr" -eq 0 ] || printf 'sif "%s" gid %s\n' "$f" "$gr"
	fi
	printf 'sif "%s" mtime @%s\n' "$f" "${ts%.*}"
	case "$t" in
		s) ty=014 ;;
		l) ty=012 ;;
		f) ty=010 ;;
		b) ty=006 ;;
		d) ty=004 ;;
		c) ty=002 ;;
		p) ty=001 ;;
	esac
	[ "$ty" = "010" ] || printf 'sif "%s" mode %03o%04o\n' "$f" "$ty" "0$md"
}

generate_debugfs_commands() {
	local dir=$1

	>"$TMP/hardlinks"
	declare -i hl
	( set -f
	  $SETX
	  # shellcheck disable=SC2068
	  find "$dir" -xdev \
		 ${exclude[@]/#/-not -path } \
		 -not -samefile "$image" \
		 -not -samefile "$dir" \
		 -printf '%D:%i %n %y %m %U %G %T@ %P\n'
	) | {
	printf 'lcd "%s"\n' "$dir"
	while read -r di hl t md us gr ts f; do
		f="${f//[[:cntrl:]]/?}"
		case "$t" in
		   f)	if [ "$hl" -gt 1 ]; then
				echo "$di $hl $t $md $us $gr $ts $f" >>"$TMP/hardlinks"
				continue
			fi
			printf 'write "%s" "%s"\n' "$f" "$f"
			;;
		   d)	printf 'mkdir "%s"\n' "$f"
			;;
		   l)	printf 'symlink "%s" "%s"\n' "$f" "$(readlink -n -- "$dir/$f")"
			;;
		   *)	if [ -v special ]; then
				case "$t" in
					c | b) minmaj="$(stat -c '%Hr %Lr' -- "$dir/$f")" ;;
					s) continue ;; # No support in debugfs(8).
					*) continue ;;
				esac
				printf 'cd "%s"\nmknod "%s" %s %s\ncd /\n' \
					"$(dirname "$f")" \
					"$(basename "$f")" \
					"$t" \
					"$minmaj"
			else
				continue
			fi
			;;
		esac
		set_metadata "$t" "$md" "$us" "$gr" "$ts" "$f"
	done
	sort "$TMP/hardlinks" \
	| while read -r di hl t md us gr ts f; do
		if [ "$di" != "${pdi-}" ]; then
			printf 'write "%s" "%s"\n' "$f" "$f"
			set_metadata "$t" "$md" "$us" "$gr" "$ts" "$f"
			printf 'sif "%s" links_count %d\n' "$f" "$hl"
			pdi="$di"
			pf="$f"
		else
			printf 'ln "%s" "%s"\n' "$pf" "$f"
		fi
	done
	}
}

e2pkg=$(rpm -q --qf '%{VERSION}' e2fsprogs)
e2ok='1.46.2'
if printf '%s\n%s\n' "$e2pkg" "$e2ok" | sort -CV; then
	echo >&2 "WARNING: e2fsprogs version below $e2ok is known to fail to create correct ext4 image!"
fi

for dir in "${src[@]}"; do
	generate_debugfs_commands "$dir" | ($SETX; /sbin/debugfs -f- -w "$image") >&"$NULL"
done

($SETX; /sbin/e2fsck $VERBOSE -p "$image")

if [ -v part ]; then
	if ! fallocate --insert-range --length 1MiB "$image" 2>/dev/null; then
		# FALLOC_FL_INSERT_RANGE is not implemented on tmpfs.
		mv "$image" "$image.part"
		($SETX; dd if="$image.part" of="$image" conv=sparse obs=1024 seek=1024 >&"$NULL")
		rm "$image.part"
	fi
	if [ "${part,,}" = "gpt" ]; then
		# Allocate 1MiB space for Secondary GPT.
		# sfdisk will truncate requested size= to 1MiB unit requiring amount to
		# be padded, so just make it '+' (maximum) and add 2MiB to the disk image.
		($SETX; truncate -s +2MiB "$image")
		# Note: sfdisk will not verify that requested size= will fit into
		# resulting partition size.
	fi
	# "The default start offset for the first partition is 1 MiB."
	($SETX; echo "size=+, type=L, bootable" | sfdisk --color=never ${part:+-X $part} $QUIET "$image")
fi

if [ -v qcow2 ]; then
	mv -f "$image" "$image.raw"
	($SETX; qemu-img convert "$image.raw" "$image" -O qcow2)
	rm "$image.raw"
fi
