create

#!/bin/bash
function configure() {
	VOLID="LINUX_RECOVERY_IMAGE"
	ROOT="root"
	ROOTFS="rootfs.img"
	ISO="out.iso"
	INSTALL_LANG=en_AU.UTF-8
	INSTALL_LANG_ENC=UTF-8
	INSTALL_KEYMAP=us
	INSTALL_TZ=Australia/Melbourne
	INSTALL_FONT=ter-v22n
	INSTALL_HOSTNAME=linux-recovery
	INSTALL_USER=user
}

true && return 1 2>/dev/null # exit if sourced
cd -- "$PWD" || exit "$?"
IFS=$'\n'
set -o errexit -o nounset -o pipefail || exit 1
function print_color() {
	local color
	color="$1"
	shift
	printf "\x1b[38;5;${color}m::\x1b[0m %s\n" "$@"
}
function print() {
	print_color 14 "$@"
	return 0
}
function debug() {
	print_color 13 "$@"
	return 0
}
function warn() {
	print_color 11 "$@" >&2
	return 0
}
function error() {
	print_color 9 "$@" >&2
	return 1
}
function write() {
	local file
	file="$1"
	shift
	debug "Writing to $file"
	printf "%s\n" "$@" >"$file"
}
function check_command_installed() {
	OUT=0
	for i in "$@"; do
		if ! command -v "$i" &>/dev/null; then
			error "$i could not be found" || true
			OUT=1
		fi
	done
	return "$OUT"
}
function check_root() {
	if [[ $EUID -ne 0 ]]; then
		error "This script must be run as root"
	fi
}
function confirm() {
	local REPLY
	local -a args
	local last
	args=("$@")
	last="$(("${#args[@]}" - 1))"
	args["$last"]="${args["$last"]} [y/N]"

	print "${args[@]}"
	read -n 1 -r
	printf "\n"
	if [[ "$REPLY" =~ ^[Yy]$ ]]; then return 0; fi
	error "Aborted"
}
function create_symlink() {
	local target
	local file
	target="$1"
	file="$2"

	if [[ -L "$file" ]]; then
		unlink -- "$file"
	fi
	if [[ -e "$file" ]]; then
		error "$file: file already exists as another type"
	fi
	ln -s -- "$target" "$file"
	debug "Symlink created: $file -> $target"
}
function create_dir() {
	local dir
	for dir in "$@"; do
		if [[ -d "$dir" ]]; then
			debug "Directory already exists: $dir"
			return 0
		fi
		if [[ -e "$dir" ]]; then
			error "$dir: file already exists as another type"
		fi
		mkdir -- "$dir"
		debug "Directory created: $dir"
	done
}
function copy_file() {
	local src
	local dest
	src="$1"
	dest="$2"
	if [[ -e "$dest" ]]; then
		unlink -- "$dest"
	fi
	cp -- "$src" "$dest"
	debug "File copied: $src -> $dest"
}
function copy_dir() {
	local src
	local dest
	src="$1"
	dest="$2"
	create_dir "$dest"
	cp --recursive -- "$src" "$dest"
	debug "Directory copied: $src -> $dest"
}
function delete_file() {
	local file
	for file in "$@"; do
		if [[ -e "$file" ]]; then
			unlink -- "$file"
			debug "File deleted: $file"
		fi
	done
}
function ensure_exists() {
	local file
	for file in "$@"; do
		if [[ ! -e "$file" ]]; then
			error "$file: file doesn't exist"
		fi
	done
}
function ensure_does_not_exist() {
	local file
	for file in "$@"; do
		if [[ -e "$file" ]]; then
			error "$file: file already exists"
		fi
	done
}
function install_packages() {
	local root
	local -A packages
	local -A force_packages
	root="$1"
	shift
	if [[ -z "$root" || ! -d "$root" ]]; then
		error "Invalid root directory"
	fi
	packages=()
	force_packages=()
	for i in "$@"; do
		if [[ -z "$i" ]]; then
			continue
		fi
		if [[ "${i:0:1}" == "!" ]]; then
			force_packages["${i:1}"]=1
			packages["${i:1}"]=1
			continue
		else
			force_packages["$i"]=0
			packages["$i"]=1
		fi
	done
	IFS=$'\n'
	if [[ -d "$root"/var/lib/pacman ]]; then
		while IFS= read -r line; do
			if [[ "${force_packages[$line]}" == "0" ]]; then
				unset -v "packages[$line]"
			fi
		done < <(pacman --root "$root" -Qq - <<<"${!packages[*]}" 2>/dev/null)
	fi
	if [[ ${#packages[@]} -eq 0 ]]; then
		debug "Packages are already installed"
		return 0
	fi
	debug "Installing ${#packages[@]} packages..."
	basestrap -cK -- "$root" - <<<"${!packages[*]}"
}
function finalise_chroot() {
	[[ -z "$1" ]] && return 1
	print "Killing processes in chroot"
	IFS=$'\n'
	while IFS= read -r pid; do
		debug "Killing process: $pid"
		kill -- "$pid"
	done < <(lsof -t -- "$1")
	if [[ -n "$(lsof -t -- "$1")" ]]; then
		error "Not all processes were killed in chroot, killing remaining processes..."
		sleep 2
		while IFS= read -r pid; do
			debug "Force killing process: $pid"
			kill -9 -- "$pid"
		done < <(lsof -t -- "$1")
		return 1
	fi
	print "Unmounting chroot if already mounted"
	umount --verbose --recursive -- "$1/dev/shm" "$1/proc" "$1/sys" "$1/dev" "$1/run" "$1/tmp" "$1" || true
}
function initialise_chroot() {
	[[ -z "$1" ]] && return 1
	finalise_chroot "$1"
	print "Initialising chroot"
	mount --verbose --types proc /proc -- "$1/proc"
	mount --verbose --rbind /sys -- "$1/sys"
	mount --verbose --make-rslave -- "$1/sys"
	mount --verbose --rbind /dev -- "$1/dev"
	mount --verbose --make-rslave -- "$1/dev"
	mount --verbose --types tmpfs tmpfs -- "$1/run"
	mount --verbose --types tmpfs tmpfs -- "$1/tmp"
	if [[ -L "/dev/shm" ]]; then
		unlink -- "/dev/shm"
		mkdir -- "/dev/shm"
		mount --verbose --types tmpfs --options nosuid,nodev,noexec shm /dev/shm
		chmod 1777 /dev/shm /run/shm
	fi
}
function chr() {
	local dir
	dir="$1"
	shift
	env --ignore-environment chroot -- "$dir" "$@"
}
function register_service() {
	local dir
	dir="$1"
	shift
	for service in "$@"; do
		create_symlink ../"$service" "$dir"/etc/dinit.d/boot.d/"$service"
	done
}
function check_installed() {
	check_command_installed mksquashfs chroot grub-install mkdir ln cp mv unlink basestrap pacman printf grub-mkrescue awk mkfs.ext4 tar mount umount resize2fs true false env lsof kill read
}
function pacman_install() {
	pacman -S artools-base coreutils e2fsprogs gawk grub lsof pacman squashfs-tools tar util-linux
}
check_root

configure
ROOT="$(realpath -- "$ROOT")"
ISO="$(realpath -- "$ISO")"
BOOT="$ROOT/boot"
ROOTFS="$(realpath -- "$ROOTFS")"
function clean() {
	check_installed || return "$?"
	finalise_chroot "$ROOT" || return "$?"
	delete_file "$ISO"
	delete_file "$ROOTFS"
	if [[ -d "$ROOT" ]]; then
		confirm "Recursively delete root directory at '$ROOT'?" && rm --one-file-system --recursive --verbose -- "$ROOT"
	fi
}
function make() {
	check_installed || return "$?"
	local -a PACKAGES

	ensure_exists grub.cfg update-dracut

	confirm "Script settings: " \
		"  Root directory is at '$ROOT'." \
		"  Root filesystem image is at '$ROOTFS'." \
		"  Output ISO file will be saved to '$ISO'." \
		"  ISO label is '$VOLID'." \
		"Recovery environment options:" \
		"  Language is '$INSTALL_LANG'." \
		"  Language encoding is '$INSTALL_LANG_ENC'." \
		"  Keymap is '$INSTALL_KEYMAP'." \
		"  Timezone is '$INSTALL_TZ'." \
		"  Hostname is '$INSTALL_HOSTNAME'." \
		"  User is '$INSTALL_USER'." \
		"  Font is '$INSTALL_FONT'." \
		"Change these settings in the script if needed." \
		"Continue?"

	PACKAGES=(
		base base-devel dinit-base dinit dinit-rc
		dbus-dinit seatd-dinit qemu-guest-agent-dinit spice-vdagent-dinit networkmanager-dinit gdm-dinit bluez-dinit cups-dinit openssh-dinit ufw-dinit cronie-dinit
		nano neovim git curl wget sudo mandoc man-pages less ripgrep bat
		linux linux-firmware intel-ucode amd-ucode efibootmgr grub dracut terminus-font squashfs-tools device-mapper xorg-xhost usbutils
		gnome-session gnome-backgrounds gnome-calculator gnome-characters gnome-clocks gnome-color-manager gnome-terminal gnome-control-center gnome-font-viewer gnome-keyring gnome-menus gnome-settings-daemon gnome-shell gnome-shell-extensions gnome-system-monitor gnome-text-editor gnome-user-docs gnome-user-share gnome-tweaks
		gvfs loupe nautilus simple-scan snapshot sushi evince
		xdg-desktop-portal-gnome xdg-user-dirs-gtk firefox baobab
		gparted dosfstools jfsutils f2fs-tools btrfs-progs exfatprogs ntfs-3g reiserfsprogs udftools xfsprogs nilfs-utils polkit gpart mtools
		lvm2 mdadm cryptsetup dmraid
		clonezilla grsync rsync
		memtest86+ dmidecode stress stress-ng sysbench
		fzf jq psmisc zip unzip unrar p7zip lzop zstd lz4 xz bzip2 gzip tar
		iputils bind nmap net-tools traceroute whois tcpdump iptraf-ng
		testdisk ddrescue lsof extundelete smartmontools r-linux foremost btop rsync rclone rkhunter clamav clamtk chntpw gnome-disk-utility
		ttf-ubuntu-font-family gnu-free-fonts ttf-fira-code ttf-fira-mono ttf-fira-sans ttf-roboto ttf-roboto-mono
		artools-base gufw
		arch-wiki-docs arch-wiki-lite
		pv
	)
	create_dir "$ROOT"
	print "Installing packages"
	install_packages "$ROOT" "${PACKAGES[@]}"

	initialise_chroot "$ROOT"

	print "Configuring system"
	write "$ROOT/etc/locale.gen" "$INSTALL_LANG $INSTALL_LANG_ENC"
	write "$ROOT/etc/locale.conf" "LANG=$INSTALL_LANG"
	write "$ROOT/etc/vconsole.conf" "KEYMAP=$INSTALL_KEYMAP" "FONT=$INSTALL_FONT"
	create_symlink /usr/share/zoneinfo/"$INSTALL_TZ" "$ROOT/etc/localtime"
	write "$ROOT/etc/hostname" "$INSTALL_HOSTNAME"
	write "$ROOT/etc/hosts" '127.0.0.1 localhost' '0::1 localhost'
	write "$ROOT/etc/sudoers.d/wheel" "%wheel ALL=(ALL) ALL"
	write "$ROOT/etc/gdm/custom.conf" '[daemon]' 'AutomaticLoginEnable=True' "AutomaticLogin=$INSTALL_USER" 'WaylandEnable=true'
	copy_file inputrc "$ROOT/etc/inputrc"
	# shellcheck disable=SC2016
	chr "$ROOT" bash -c '(id -u -- "$1" || useradd --base-dir /home --create-home -- "$1")' bash "$INSTALL_USER"
	chr "$ROOT" usermod --append --groups wheel,lp --shell /bin/bash -- "$INSTALL_USER"
	chr "$ROOT" locale-gen
	chr "$ROOT" passwd -d root
	chr "$ROOT" passwd -d "$INSTALL_USER"
	register_service "$ROOT" NetworkManager bluetoothd seatd dbus gdm qemu-ga spice-vdagentd cupsd ufw cronie
	create_symlink nvim "$ROOT"/usr/bin/vi
	create_symlink nvim "$ROOT"/usr/bin/vim
	create_dir "$ROOT/home/$INSTALL_USER/.config" "$ROOT/home/$INSTALL_USER/.config/dconf"
	copy_dir firefox "$ROOT/home/$INSTALL_USER/.mozilla"
	chr "$ROOT" chown -Rv -- "$INSTALL_USER" "/home/$INSTALL_USER/.mozilla"
	chr "$ROOT" chown -Rv -- "$INSTALL_USER" "/home/$INSTALL_USER/.config"
	chr "$ROOT" sudo -u "$INSTALL_USER" /bin/dbus-run-session /bin/dconf load / <settings.dconf

	print "Creating initramfs image with dracut"
	write "$ROOT/etc/dracut.conf.d/02-squashfs.conf" 'add_dracutmodules+=" dmsquash-live dm udev-rules kernel-modules kernel-modules-extra "' 'install_items+=" /usr/bin/lsblk /usr/bin/lsmod /usr/bin/lspci /usr/bin/lsusb "'
	chr "$ROOT" bash - <update-dracut

	finalise_chroot "$ROOT"
	delete_file "$ROOT"/etc/machine-id "$ROOT"/var/lib/dbus/machine-id

	print "Creating ext4 filesystem image, this may take a while"
	delete_file "$ROOTFS"
	# creates a ext4 image of the root directory, excluding the boot, log, and cache directories
	tar --directory "$ROOT" --anchored --verbose \
		--exclude='./boot/*' --exclude='./var/cache/*' --exclude='./root/.cache/*' --exclude='./home/'"$INSTALL_USER"'/.cache/*' \
		--exclude='./var/log/audit' --exclude='./var/log/httpd' --exclude='./var/log/clamav' --exclude='./var/log/old' \
		--exclude='./var/log/pacman.log' --exclude='./var/log/dinit/*' \
		--one-file-system --create \
		--xattrs --acls --numeric-owner --preserve-permissions --xattrs --xattrs-include=* --acls --atime-preserve . 2>/tmp/tar.log |
		mkfs.ext4 -d - -- "$ROOTFS" 10G
	resize2fs -M -- "$ROOTFS" # shrink the filesystem to the minimum size

	print "Creating squashfs filesystem image, this may take a while"
	SQUASHFS_DIR="$BOOT/LiveOS"
	SQUASHFS_PATH="$SQUASHFS_DIR/squashfs.img"
	create_dir "$SQUASHFS_DIR"
	mksquashfs - "$SQUASHFS_PATH" -p '/ d 0755 0 0' -p '/LiveOS/ d 0755 0 0' -p '/LiveOS/rootfs.img f 0644 0 0 cat' -noappend -no-recovery <"$ROOTFS"

	print "Generating GRUB configuration"
	create_dir "$BOOT"/boot
	create_dir "$BOOT"/boot/grub
	awk -v VOLID="$VOLID" '{gsub(/\$VOLID/, VOLID);}1' grub.cfg >"$BOOT"/boot/grub/grub.cfg

	print "Creating ISO"
	grub-mkrescue -o "$ISO" "$BOOT" -- -volid "$VOLID"

	print "Output ISO: $ISO"
}
function main() {
	if [[ "$#" == "0" ]]; then
		print "Available commands:"
		print "clean: Remove generated files"
		print "make: Create the recovery ISO"
		print "chroot: Mount the chroot"
		print "unchroot: Unmount the chroot if already mounted"
		print "install: Install required packages with pacman"
		return 0
	fi
	if [[ "$#" != "1" ]]; then
		error "Invalid usage" >&2
		return 1
	fi
	case "$1" in
	clean)
		check_installed || return "$?"
		clean
		return "$?"
		;;
	make)
		check_installed || return "$?"
		make
		return "$?"
		;;
	chroot)
		check_installed || return "$?"
		initialise_chroot "$ROOT"
		return "$?"
		;;
	unchroot)
		check_installed || return "$?"
		finalise_chroot "$ROOT"
		return "$?"
		;;
	install)
		pacman_install
		return "$?"
		;;
	*)
		error "Invalid command" >&2
		return 1
		;;
	esac
}
main "$@"