#!/usr/bin/env bash

# get available compression types
declare -A compressTypes
compressTypes=(
	["bz2"]="bzip2 pbzip2 lbzip2"
	["gz"]="gzip pigz"
	["lrz"]="lrzip"
	["lz"]="lzip plzip"
	["lz4"]="lz4"
	["lzo"]="lzop"
	["xz"]="xz pixz"
	["zst"]="zstd"
)
declare -A compressAvail
for ext in "${!compressTypes[@]}"; do
	for exechk in ${compressTypes[${ext}]}; do
		binchk=$(command -v "${exechk}")
		if [[ -n "$binchk" ]]; then
			compressAvail+=(["${ext}"]="$binchk")
		fi
	done
done

# set flag variables to null/default
optExcludeBoot=false
optExcludeConfidential=false
optExcludeLost=true
optQuiet=false
optUserExclude=()
optUserInclude=()
optCompressType="bz2"
optSeperateKernel=false


function showHelp() {
	echo "Usage:"
	echo "$(basename "$0") [-b -c -k -l -q] [-C <type>] [-s || -t <target>] [-e <exclude>...] [-i <include>...] <archive> [-- [tar-opts]]"
	echo "Position Arguments:"
	echo "    <archive>    archive name to create with optional path"
	echo "    [tar-opts]   additional options to pass to tar command"
	echo 
	echo "Options:"
	echo "    -b           excludes boot directory"
	echo "    -c           excludes some confidential files (currently only .bash_history and connman network lists)"
	echo "    -k           separately save current kernel modules and src (creates smaller targetArchives and saves decompression time)"
	echo "    -l           includes lost+found directory"
	echo "    -q           activates quiet mode (no confirmation)"
	echo "    -C <type>    specify tar compression (default: ${optCompressType}, available: ${!compressAvail[*]})"
	echo "    -s           makes archive of current system"
	echo "    -t <path>    makes archive of system located at the <path>"
	echo "    -e <exclude> an additional exclude directory (one dir one -e, do not use it with *)"
	echo "    -i <include> an additional include. This has higher precedence than -e, -t, and -s"
	echo "    -h           display this help message."

	if [[ "$1" -ge 0 ]]; 	then
		exit "$1"
	else
		exit 0
	fi
}

function errorMsg() {
	local rc=0

	if [[ "$1" -gt 0 ]]; then
		rc="$1"
		shift
	fi

	echo "$*" >&2

	if [[ "$rc" -gt 0 ]]; then
		exit "$rc"
	fi
}


# reads options:
tarArgs=()
while [[ $# -gt 0 ]]; do
	while getopts ":t:C:e:i:skqcblh" flag; do
		case "$flag" in
			t)	targetPath="$OPTARG";;
			s)	targetPath="/";;
			C)	optCompressType="$OPTARG";;
			q)	optQuiet=true;;
			k)	optSeperateKernel=true;;
			c)	optExcludeConfidential=true;;
			b)	optExcludeBoot=true;;
			l)	optExcludeLost=false;;
			e)	optUserExclude+=("${OPTARG}");;
			i)	optUserInclude+=("$OPTARG");;
			h)	showHelp 0;;
			\?)	errorMsg 1 "Invalid option: -$OPTARG";;
			:)	errorMsg 1 "Option -$OPTARG requires an argument.";;
		esac
	done || exit 1

	[[ $OPTIND -gt $# ]] && break # reached the end of parameters

	shift $((OPTIND - 1)) # Free processed options so far
	OPTIND=1              # we must reset OPTIND
	if [[ -z "$targetArchive" ]]; then
		targetArchive=$1
	else
		tarArgs[${#tarArgs[*]}]=$1
	fi
	#args[${#args[*]}]=$1  # save first non-option argument (a.k.a. positional argument)
	shift                 # remove saved arg
done


# checks if run as root:
if [[ "$(id -u)" -ne 0 ]]; then
	echo "$(basename "$0"): must run as root"
	exit 250
fi

if [[ -z "$targetPath" ]]; then
	echo "$(basename "$0"): no system path specified"
	exit 1
fi

# make sure targetPath path ends with slash
if [[ "$targetPath" != */ ]]; then
	targetPath="${targetPath}/"
fi

# checks for correct output file specification
if [[ -z "$targetArchive" ]]; then
	echo "$(basename "$0"): no archive file name specified"
	exit 1
fi

# determines if filename was given with relative or absolute path
#if (($(grep -c '^/' <<< "$targetArchive") > 0)); then
if [[ "$targetArchive" =~ ^\/.* ]]; then
	stage4Filename="${targetArchive}"
else
	stage4Filename="$(pwd)/${targetArchive}"
fi

# Check if compression in option and filename
if [[ -z "$optCompressType" ]]; then
	echo "$(basename "$0"): no archive compression type specified"
	exit 1
else
	#stage4Filename="${stage4Filename}.${optCompressType}"
	stage4Ext="tar.${optCompressType}"
fi

# Check if specified type is available
if [[ -z "${compressAvail[$optCompressType]}" ]]; then
	echo "$(basename "$0"): specified archive compression type not supported."
	echo "Supported: ${compressAvail[*]}"
	exit 1
fi

# Check if using seperate kernel archive option
if $optSeperateKernel; then
	optUserExclude+=("${targetPath}usr/src/*")
	optUserExclude+=("${targetPath}lib*/modules/*")
fi

# tarExcludes:
tarExcludes=(
	"dev/*"
	"var/tmp/*"
	"media/*"
	"mnt/*/*"
	"proc/*"
	"run/*"
	"sys/*"
	"tmp/*"
	"var/lock/*"
	"var/log/*"
	"var/run/*"
	"var/lib/docker/*"
	"var/lib/containers/*"
	"var/lib/machines/*"
	"var/lib/libvirt/*"
	"var/lib/lxd/*"
	"home/*/*"
)

tarExcludesPortage=(
	"var/db/repos/*/*"
	"var/cache/distfiles/*"
	"var/cache/binpkgs/*"
	"usr/portage/*"
)

tarIncludes=(
	"dev/console"
	"dev/null"
	"var/db/pkg/*"
)

tarExcludes=("${tarExcludes[@]/#/"$targetPath"}")
tarIncludes=("${tarIncludes[@]/#/"$targetPath"}")

if [[ "$targetPath" == '/' ]]; then
	tarExcludes+=("$(realpath "${stage4Filename}")*")
	#if $optSeperateKernel; then
	#	tarExcludes+=("$(realpath "${stage4Filename}.ksrc.${stage4Ext}")")
	#	tarExcludes+=("$(realpath "${stage4Filename}.kmod.${stage4Ext}")")
	#fi
	if command -v portageq &>/dev/null; then
		portageRepos=$(portageq get_repos /)
		for i in ${portageRepos}; do
			repoPath=$(portageq get_repo_path / "${i}")
			tarExcludes+=("${repoPath}/*")
		done
		tarExcludes+=("$(portageq distdir)/*")
		tarExcludes+=("$(portageq pkgdir)/*")
	else
		tarExcludes+=("${tarExcludesPortage[@]/#/"/"}")
	fi
else
	tarExcludes+=("${tarExcludesPortage[@]/#/"$targetPath"}")
fi

if $optExcludeConfidential; then
	tarExcludes+=("${targetPath}home/*/.bash_history")
	tarExcludes+=("${targetPath}root/.bash_history")
	tarExcludes+=("${targetPath}var/lib/connman/*")
fi

if $optExcludeBoot; then
	tarExcludes+=("${targetPath}boot/*")
fi

if $optExcludeLost; then
	tarExcludes+=("lost+found")
fi

tarExcludes+=("${optUserExclude[@]}")
tarIncludes+=("${optUserInclude[@]}")

# Compression options
compressOptions=("${compressAvail[$optCompressType]}")
if [[ "${compressAvail[$optCompressType]}" == *"/xz" ]]; then
	compressOptions+=("-T0")
fi

# Generic tar options:
tarOptions=(
	"-cpP"
	"--ignore-failed-read"
	"--xattrs-include=*.*"
	"--numeric-owner"
	"--checkpoint=.500"
	"--use-compress-prog=${compressOptions[*]}"
)

tarOptions+=("${tarArgs[@]}")


# if not in optQuiet mode, this message will be displayed:
if ! $optQuiet; then
	echo "Are you sure that you want to make a stage 4 tarball of the system"
	echo "located under the following directory?"
	echo "$targetPath"
	echo
	echo "WARNING: since all data is saved by default the user should exclude all"
	echo "security- or privacy-related files and directories, which are not"
	echo "already excluded by mkstage4 options (such as -c), manually per cmdline."
	echo "example: \$ $(basename "$0") -s /my-backup --exclude=/etc/ssh/ssh_host*"
	echo
	echo "COMMAND LINE PREVIEW:"
	#echo 'tar' "${tarOptions[@]}" "${tarIncludes[@]}" "${tarExcludes[@]}" -f "$stage4Filename" "${targetPath}"
	echo 'tar' "${tarOptions[@]}" "${tarExcludes[@]/#/--exclude=}" -f "${stage4Filename}.${stage4Ext}" "${targetPath}" "${tarIncludes[@]}"
	#${excludes[@]/#/--exclude=
	if $optSeperateKernel; then
		echo
		echo 'tar' "${tarOptions[@]}" -f "${stage4Filename}.ksrc.${stage4Ext}" "${targetPath}usr/src/linux-$(uname -r)"
		echo 'tar' "${tarOptions[@]}" -f "${stage4Filename}.kmod.${stage4Ext}" "${targetPath}lib"*"/modules/$(uname -r)"
	fi
	echo
	echo -n 'Type "yes" to continue or anything else to quit: '
	read -r promptAgree
	if [[ "${promptAgree,,}" == "yes" ]]; then
		optQuiet=true
	fi
fi

# start stage4 creation:
if $optQuiet; then
	#echo "Would've worked"
	tar "${tarOptions[@]}" "${tarIncludes[@]}" "${tarExcludes[@]/#/--exclude=}" -f "${stage4Filename}.${stage4Ext}" "${targetPath}"
	if [[ "$optSeperateKernel" ]]
	then
		tar "${tarOptions[@]}" -f "${stage4Filename}.ksrc.${stage4Ext}" "${targetPath}usr/src/linux-$(uname -r)"
		tar "${tarOptions[@]}" -f "${stage4Filename}.kmod.${stage4Ext}" "${targetPath}lib"*"/modules/$(uname -r)"
	fi
fi