arkdep/arkdep

#!/usr/bin/env bash

#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.

set -o pipefail

# Print manual if no parameters provided or invalid amount of parameters is provided
if [[ ! -n $1 ]]; then
	cat <<- END
	Usage: arkdep-deploy <action> [target]
	Actions:
	  update		Check for updates, optionally provide a target, if no target provided it defaults to primary
	  deploy		Deploy a new or update an existing deployment
	  init			Initialize arkdep on a new system
	  teardown		Remove all arkdep-deploy related files and folders
	  get-available		List available packages in repo
	  list			List all currently deployed images
	  remove		Remove a specified deployment
	END
	exit 0
fi

## Set common variables
#
declare -r arkdep_dir='/arkdep/'

## Load config file
#
source $(readlink -m $arkdep_dir/config)

## Common functions
#
# Cleanup and quit if error
cleanup_and_quit () {

	# If any paramters are passed we will assume it to be an error message
	[[ -n $1 ]] && printf "\e[1;31m<#>\e[0m $*\e[0m\n" >&2

	# Remove the subvolume we were working on
	# TODO: Make this a generic function and share with the removal of old images?
	if [[ -n ${data[0]} ]]; then
		btrfs property set -f -ts $arkdep_dir/deployments/${data[0]}/rootfs ro false
		rm -rf $arkdep_dir/deployments/${data[0]}
		rm -rf /boot/arkdep/${data[0]}
		rm /boot/loader/entries/${data[0]}.conf
	fi

	# Quit program if argument provided to function
	[[ -n $1 ]] && exit 1

	# Otherwise just quit, there is no error
	exit 0

}

## Error checking
#
# Quit if not root
[[ ! $EUID -eq 0 ]] &&
	printf '\e[1;31m<#>\e[0m\e[1m This program has to be run as root\n\e[0m' &&
	exit 1

# Check if all dependencies are installed, quit if not
for prog in btrfs wget dracut bootctl curl; do
	if ! command -v $prog > /dev/null; then
		printf "\e[1;31m<#>\e[0m\e[1m Failed to locate $prog, ensure it is installed\e[0m\n"
		err=1
	fi

	[[ $err ]] && exit 1
done

## Core functions
#
# Initialize the system for arkdep
init () {

	# Ensure systemd-boot is installed before continuing, for it is the only thing we support
	bootctl is-installed || cleanup_and_quit 'systemd-boot seems to not be installed'

	printf '\e[1;34m-->\e[0m\e[1m Initializing arkdep\e[0m\n'

	[[ -d $arkdep_dir ]] && cleanup_and_quit "$arkdep_dir already exists"

	# Create the /arkdep subvolume
	printf "\e[1;34m-->\e[0m\e[1m Creating $(readlink -m $arkdep_dir) subvolume\e[0m\n"
	btrfs subvolume create $arkdep_dir || cleanup_and_quit "Failed to create btrfs subvolume"

	# Create directory structure
	printf "\e[1;34m-->\e[0m\e[1m Creating directory structure\e[0m\n"
	mkdir -pv $(readlink -m $arkdep_dir/deployments) \
		$(readlink -m $arkdep_dir/deployments) \
		$(readlink -m $arkdep_dir/cache) \
		$(readlink -m $arkdep_dir/templates) \
		$(readlink -m $arkdep_dir/overlay) \
		$(readlink -m $arkdep_dir/shared) ||
		cleanup_and_quit "Failed to create /arkdep and related directories"

	# Create empty database files
	touch $(readlink -m $arkdep_dir/tracker)

	# Add home shared subvolume and make writable
	btrfs subvolume create $(readlink -m $arkdep_dir/shared/home) || cleanup_and_quit "Failed to create home subvolume"
	btrfs subvolume create $(readlink -m $arkdep_dir/shared/root) || cleanup_and_quit "Failed to create root subvolume"
	btrfs property set -f -ts $(readlink -m  $arkdep_dir/shared/home) ro false
	btrfs property set -f -ts $(readlink -m  $arkdep_dir/shared/root) ro false

	# Write default config file
	printf "\e[1;34m-->\e[0m\e[1m Adding default config file\e[0m\n"
	cat <<- END > $arkdep_dir/config
	# Write /arkdep/overlay overlay to root or etc
	enable_overlay=1

	# URL to image repository, do not add trailing slash
	repo_url='https://repo.arkanelinux.org/arkdep'

	# Default image pulled from repo if nothing defined
	repo_default_image='arkanelinux'

	# Keep the latest n+1 deployments, remove anything older
	deploy_keep=2
	END

	# Add default bootloader config file
	cat <<- END > $arkdep_dir/templates/systemd-boot
	title Arkane GNU/Linux - arkdep
	linux /arkdep/%target%/vmlinuz
	initrd /amd-ucode.img
	initrd /intel-ucode.img
	initrd /arkdep/%target%/initramfs-linux.img
	options root="LABEL=arkane_root" rootflags=subvol=/arkdep/deployments/%target%/rootfs rw
	END

	exit 0

}

teardown () {

	cat <<- END
	WARNING: Removing arkdep may leave your system in an unbootable state and you
	may have to manually reconfigure your bootloader etc.. Only proceed if you know
	what you are doing!

	The following changes will be made to your system;
	- All subvolumes under $arkdep_dir will be deleted

	END

	read -p 'Type "I KNOW WHAT I AM GOING" in uppercase to confirm that you know what you are doing: ' input_confirm


	if [[ $input_confirm == 'I KNOW WHAT I AM DOING' ]]; then

		printf '\e[1;34m-->\e[0m\e[1m Tearing down arkdep\e[0m\n'

		# Quit with error if $arkdep_dir does not exist
		if [[ ! -d $arkdep_dir ]]; then
			printf "\e[1;31m<#>\e[0m $(readlink -m $arkdep_dir) does not exist, there is nothing to tear down"
			exit 1
		fi

		# Remove all bootloader entries
		rm -v $(grep -ril arkdep /boot/loader/entries)

		# Ensure all nested volumes in arkdep are writable and remove
		for volume in $(btrfs subvolume list / | grep -oE '[^ ]+$' | grep "^$arkdep_dir" | tac); do
			btrfs property set -f -ts $(readlink -m $volume) ro false
			btrfs subvolume delete $volume
		done

	else
		printf '\e[1;34m-->\e[0m\e[1m Teardown canceled, no changes made to system\e[0m\n'
	fi

	exit 0

}

remove_deployment () {

	# Ensure required vars are set
	[[ -z $2  ]] && \
		printf 'No deployment defined\n' && exit 1

	# Ensure user only provided a single target
	[[ ! -z $3 ]] && \
		printf 'Multiple targets provided, remove only accepts a single target at a time\n' && exit 1

	declare -r hits=($(grep -v $1 $arkdep_dir/tracker))

	if [[ ${#hits[@]} -gt 1 ]]; then
		printf 'Multiple deployments match target, be more specific\n'
		exit 1
	elif [[ ${#hits[@]} -lt 1 ]]; then
		printf 'No deployments match target\n'
		exit 1
	fi

	declare -r target="${hits[0]}"

	# Remove bootloader entry
	rm -rfv /boot/loader/entries/$target.conf
	rm -rfv /boot/arkdep/$target

	# Ensure the deployment and all sub-volumes are writable
	for volume in $(btrfs subvolume list / | grep -oE '[^ ]+$' | grep $target); do
		echo $volume
		btrfs property set -f -ts $volume ro false || printf "failed to make subvol $volume writable\n"
	done

	# Remove the deployment
	rm -rf $(readlink -m $arkdep_dir/deployments/$target)

	# Remove from tracker
	grep -v $2 $arkdep_dir/tracker > $arkdep_dir/tracker_tmp
	mv $arkdep_dir/tracker_tmp $arkdep_dir/tracker

	exit 0

}

# List all available packages defined in the repo's list file
get_available () {
	printf "\e[1;34m-->\e[0m\e[1m Downloading list file from $repo_url\e[0m\n"
	curl -sf "${repo_url}/list" || cleanup_and_quit 'Failed to download repo file'
}

# List all deployments and provide additional information about any untracked deployments
list_deployed () {

	declare tracker=($(cat $arkdep_dir/tracker))
	declare deployed=($(ls $arkdep_dir/deployments/))
	declare untracked=${deployed[*]}

	# Compare items in tracker to actual deployed
	for tracked in ${tracker[@]}; do
		untracked=("${untracked[@]/$tracked}")
	done

	printf "${tracker[*]}\n"

	# Clean whitespaces
	untracked=$(echo $untracked | xargs)

	if [[ ! -z $untracked ]]; then
		printf "\nThe following deployments were found but are not tracked;\n"
		for t in $untracked; do
			printf "$t\n"
		done
	fi

}

# Deploy a new or update an existing deployment
deploy () {

	# target and version are optional, if not defined default to primary as defined in
	# /arkdep/config and latest
	if [[ -n $2 ]]; then
		declare -r deploy_target=$2
	else
		declare -r deploy_target="$repo_default_image"
	fi

	if [[ -n $3 ]]; then
		declare -r deploy_version=$3
	else
		declare -r deploy_version='latest'
	fi

	printf "\e[1;34m-->\e[0m\e[1m Deploying $deploy_target $deploy_version\e[0m\n"

	# If latest is requested grab database and get first line
	printf "\e[1;34m-->\e[0m\e[1m Downloading database from repo\e[0m\n"
	if [[ $deploy_version == 'latest' ]]; then
		declare curl_data=$(curl -sf "${repo_url}/${deploy_target}/database" | head -n 1)
	else
		# Only return first hit
		declare curl_data=$(curl -sf "${repo_url}/${deploy_target}/database" | grep $3 | head -1)
	fi

	# Split latest_version at the delimiter, creating an array with data.0=package ver, data.1=compression method, data.2=sha1 hash
	readarray -d : -t data <<< "$curl_data"
	# A carriage feed is inserted for some reason, lets remove it
	data[2]=${data[2]//[$'\t\r\n']}

	# Lets ensure the requested image is not already deployed
	if [[ -e $arkdep_dir/deployments/${data[0]} ]]; then
		printf "\e[1;33m<!>\e[0m\e[1m ${data[0]} is already deployed, canceling deployment\e[0m\n"
		exit 0
	fi

	# Check if requested version is already downloaded
	if [[ -e $arkdep_dir/cache/${data[0]}.tar.${data[1]} ]]; then
		printf "\e[1;34m-->\e[0m\e[1m ${data[0]} already in cache, skipping download\e[0m\n"
	else

		# Download the tarball if not yet downloaded
		if [[ ! -e $arkdep_dir/cache/${data[0]}.tar.${data[1]} ]]; then
			wget -P $(readlink -m $arkdep_dir/cache/) "$repo_url/$repo_default_image/${data[0]}.tar.${data[1]}" ||
				cleanup_and_quit 'Failed to download tarball'
		fi

	fi

	printf "\e[1;34m-->\e[0m\e[1m Validating integrity\e[0m\n"
	sha1sum "$(readlink -m $arkdep_dir/cache/${data[0]}.tar.${data[1]})" |
		grep "${data[3]}" ||
		cleanup_and_quit "Checksum does not match repo file, got $chksum\e[0m\n"

	# Extract the root image if not yet extracted
	printf "\e[1;34m-->\e[0m\e[1m Writing root\e[0m\n"

	# Create directory using unique deployment name
	mkdir -pv $(readlink -m $arkdep_dir/deployments/${data[0]}) || cleanup_and_quit 'Failed to create deployment directory'

	if [[ ! -e $arkdep_dir/cache/${data[0]}-rootfs.img ]]; then
		tar -xf $(readlink -m $arkdep_dir/cache/${data[0]}.tar.${data[1]}) -C $(readlink -m $arkdep_dir/cache/) "./${data[0]}-rootfs.img" ||
			cleanup_and_quit 'Failed to extract root'
	fi

	# Write the root image
	btrfs receive -f $(readlink -m $arkdep_dir/cache/${data[0]}-rootfs.img) $(readlink -m $arkdep_dir/deployments/${data[0]}) ||
		cleanup_and_quit 'Failed to receive root'

	# Cleanup root image
	rm $(readlink -m $arkdep_dir/cache/${data[0]}-rootfs.img)

	# Extract the etc image if not yet extracted
	printf "\e[1;34m-->\e[0m\e[1m Writing etc\e[0m\n"
	if [[ ! -e $arkdep_dir/cache/${data[0]}-etc.img ]]; then
		tar -xf $(readlink -m $arkdep_dir/cache/${data[0]}.tar.${data[1]}) -C $(readlink -m $arkdep_dir/cache/) "./${data[0]}-etc.img" ||
			cleanup_and_quit 'failed to extract etc'
	fi

	# Write the etc image and create var directory, we have to unlock rootfs temporarily to do this
	btrfs property set -f -ts $(readlink -m $arkdep_dir/deployments/${data[0]}/rootfs) ro false ||
		cleanup_and_quit 'Failed to unlock root to write etc'

	btrfs receive -f $(readlink -m $arkdep_dir/cache/${data[0]}-etc.img) $(readlink -m $arkdep_dir/deployments/${data[0]}/rootfs/) ||
		cleanup_and_quit 'Failed to receive etc'

	printf "\e[1;34m-->\e[0m\e[1m Ensure var, root and arkdep mountpoints exist\e[0m\n"
	mkdir -pv $(readlink -m $arkdep_dir/deployments/${data[0]}/rootfs/var)
	mkdir -pv $(readlink -m $arkdep_dir/deployments/${data[0]}/rootfs/root)
	mkdir -pv $(readlink -m $arkdep_dir/deployments/${data[0]}/rootfs/arkdep)

	# Lock the root volume again
	btrfs property set -f -ts $(readlink -m $arkdep_dir/deployments/${data[0]}/rootfs) ro true ||
		cleanup_and_quit 'Failed to lock root'

	# Unlock the etc deployment
	btrfs property set -f -ts $(readlink -m $arkdep_dir/deployments/${data[0]}/rootfs/etc) ro false ||
		cleanup_and_quit 'Failed to unlock root to write etc'

	# Cleanup etc image
	rm $(readlink -m $arkdep_dir/cache/${data[0]}-etc.img)

	# Write the var image
	if [[ ! -e $arkdep_dir/shared/var ]]; then
		printf "\e[1;34m-->\e[0m\e[1m Writing var\e[0m\n"

		# Extract the var image if not yet extracted
		if [[ ! -e $arkdep_dir/cache/${data[0]}-var.img ]]; then
			tar -xf $(readlink -m $arkdep_dir/cache/${data[0]}.tar.${data[1]}) -C $(readlink -m $arkdep_dir/cache/) "./${data[0]}-var.img" ||
				cleanup_and_quit 'failed to extract var'
		fi

		btrfs receive -f $(readlink -m $arkdep_dir/cache/${data[0]}-var.img) $(readlink -m $arkdep_dir/shared/) ||
			cleanup_and_quit 'Failed to receive var'
	else
		# Notify if var is not deployed
		printf "\e[1;33m<!>\e[0m\e[1m ${data[0]} var is already preset, skipping var deployment\e[0m\n"
	fi

	# Make var writable
	btrfs property set -f -ts $(readlink -m $arkdep_dir/shared/var) ro false ||
		cleanup_and_quit 'Failed to unlock var'

	# Cleanup var image
	rm $(readlink -m $arkdep_dir/cache/${data[0]}-var.img)

	# Add overlay if enabled
	if [[ $enable_overlay -eq 1 ]]; then
		printf "\e[1;34m-->\e[0m\e[1m Copying overlay to deployment\e[0m\n"
		declare -r overlay_files=($(ls $arkdep_dir/overlay/))

		# Check if only /etc is present, if it is we do not have to unlock the root volume
		for file in ${overlay_files[*]}; do
			if [[ ! $file == 'etc' ]]; then
				printf "\e[1;33m<!>\e[0m\e[1m ${data[0]} Non /etc file or directory detected, root will be temporarily unlocked\e[0m\n"
				overlay_unlock_root=1
			fi
		done

		# Unlock root if required
		if [[ $overlay_unlock_root -eq 1 ]]; then
			btrfs property set -f -ts $(readlink -m $arkdep_dir/deployments/${data[0]}) ro false
		fi

		cp -rv $(readlink -m $arkdep_dir/overlay/*) $(readlink -m /$arkdep_dir/deployments/${data[0]}/rootfs/)

		# Lock root again if required
		if [[ $overlay_unlock_root -eq 1 ]]; then
			btrfs property set -f -ts $(readlink -m $arkdep_dir/deployments/${data[0]}) ro true
		fi
	fi

	printf "\e[1;34m-->\e[0m\e[1m Copying kernel image\e[0m\n"
	mkdir -pv $(readlink -m /boot/arkdep/${data[0]})
	cp -v $arkdep_dir/deployments/${data[0]}/rootfs/usr/lib/modules/*/vmlinuz /boot/arkdep/${data[0]}/ ||
		cleanup_and_quit 'Failed to copy kernel image'

	# Install kernel and generate initramfs
	printf "\e[1;34m-->\e[0m\e[1m Generating initramfs\e[0m\n"
	dracut -k $(cd /arkdep/deployments/${data[0]}/rootfs/usr/lib/modules/*; pwd) \
		--kernel-image /boot/arkdep/${data[0]}/vmlinuz \
		--force \
		/boot/arkdep/${data[0]}/initramfs-linux.img || cleanup_and_quit 'Failed to generate initramfs'

	# Add to database
	# TODO: If this step is never reached ensure cleanup, maybe write a "busy file" somewhere
	printf "\e[1;34m-->\e[0m\e[1m Updating database\e[0m\n"
	printf "${data[0]}\n$(cat $(readlink -m $arkdep_dir/tracker | head -$deploy_keep))" > $arkdep_dir/tracker

	# Deploy bootloader configuration
	printf "\e[1;34m-->\e[0m\e[1m Adding bootloader entry\e[0m\n"
	sed "s/%target%/${data[0]}/" $arkdep_dir/templates/systemd-boot > /boot/loader/entries/${data[0]}.conf

	printf "\e[1;34m-->\e[0m\e[1m Setting new bootloader entry as default\e[0m\n"
	# Configuring it with a oneshot for now, for testing
	#bootctl set-default ${data[0]}.conf || cleanup_and_quit "Failed to set default bootloader entry"
	bootctl set-oneshot ${data[0]}.conf || cleanup_and_quit "Failed to set default bootloader entry"

	# Remove entries outside of keep
	declare -r remove_deployments="$(cat $arkdep_dir/tracker | head -$deploy_keep | grep -rvf - $(readlink -m $arkdep_dir/tracker))"

	# Remove old deployments
	for deployment in $remove_deployments; do
		printf "\e[1;34m-->\e[0m\e[1m Removing old deployment $deployment\e[0m\n"
		rm -v /boot/loader/entries/$deployment.conf
		btrfs property set -f -ts $(readlink -m $arkdep_dir/deployments/$deployment) ro false
		rm -rfv $(readlink -m $arkdep_dir/deployments/$deployment)
		rm -rfv $(readlink -m /boot/arkdep/$deployment)
		grep -rv $deployment $(readline -m $arkdep_dir/tracker) > $arkdep_dir/tracker
	done

}

[[ $1 == 'init' ]] && init
[[ $1 == 'teardown' ]] && teardown
[[ $1 == 'update' ]] && check_for_updates
[[ $1 == 'get-available' ]] && get_available
[[ $1 == 'deploy' ]] && deploy $1 $2 $3
[[ $1 == 'list' ]] && list_deployed
[[ $1 == 'remove' ]] && remove_deployment $1 $2 $3