# ------------------------------------------------------------------------------
#
#                          Functions for OpenSLX-NG
#
# ------------------------------------------------------------------------------
#
# check if we have our environment variables to check
# if we actually got sourced by the main script
if [ -z "${SELF_PID}" -o  -z "${ROOT_DIR}" ]; then
	# not using perror, since we probably don't have it.
	echo "Neither SELF_PID nor ROOT_DIR is set. Was this included by OpenSLX-NG?"
	exit 1
fi

# ------------------------------------------------------------------------------
#
#                          General helper functions
#
# ------------------------------------------------------------------------------
banner () {
	echo -e "\033[38;5;196m\t"'                                    .__          '
	echo -e "\033[38;5;196m\t"'  ____ ______   ____   ____   _____|  | ___  ___ '
	echo -e "\033[38;5;202m\t"' /  _ \\\\____ \\_/ __ \\ /    \\ /  ___/  | \\  \\/  / '
	echo -e "\033[38;5;208m\t"'(  <_> )  |_> >  ___/|   |  \\\\___ \\|  |__>    <  '
	echo -e "\033[38;5;214m\t"' \\____/|   __/ \\___  >___|  /____  >____/__/\\_ \\ '
	echo -e "\033[38;5;220m\t"'       |__|        \\/     \\/     \\/           \\/ '
	echo -e "\033[38;5;220m\t"
	echo -e "\033[38;5;226m\t                     ** OpenSLX Project // 2015 **"
	echo -e "\033[38;5;232m\t                        http://lab.openslx.org/"
	echo -e "\033[00m"
}

pinfo() {
	echo -e "\033[38;5;10m[info]\033[0m $@"
}

pwarning() {
	echo -e "\033[38;5;11m[warning]\033[0m $@"
}

perror() {
	echo -e "\033[38;5;9m[error]\033[0m $@"
	kill "$SELF_PID"
	exit 1
}

# accepts one argument to be printed as a warning before dumping the help
print_usage() {
	[ -n "$1" ] && pwarning "$1"
	pinfo "USAGE:"
	pinfo "$ARG0 <action> <actions_params>"
	pinfo "\t\tActions: '--clone', '--package', '--export', 'update'"
	pinfo ""
	pinfo "CLONING:      rsync remote host to local directory"
	pinfo "$ARG0 --clone --host <host> [--syncdir <path>]"
	pinfo "\t\tIf not specified, --syncdir = './clones/<host>/stage4'"
	pinfo ""
	pinfo "PACKAGING:    pack local rsync directory as qcow2-container"
	pinfo "$ARG0 --package --host <host> [--container <path>] [--syncdir <dir>]"
	pinfo "\t\tIf <container> is not specified, --container = './clones/<host>/stage4.qcow2'"
	pinfo "\t\tIf <syncdir> is not specified, --syncdir = './clones/<host>/stage4'"
	pinfo ""
	pinfo "EXPORTING:    rsync remote host to a new qcow2-container"
	pinfo "$ARG0 --export --host <host> [--container <path>] [--syncdir <dir>]"
	pinfo "\t\tIf <container> is not specified, --container = './clones/<host>/stage4.qcow2'"
	pinfo ""
	pinfo "UPDATING:     rsync remote host to an existing qcow2-container"
	pinfo "$ARG0 --update --host <host> [--container <path>]"
	pinfo "\t\tIf <container> is not specified, --container = './clones/<host>/stage4.qcow2'"
	kill "$SELF_PID"
	exit 1
}

# helper to parse the command line arguments and fill the environment
# with the parameters given. Since the fallbacks for those variables
# are only dependent on the selected action, we will also post-process
# them to make sure they are set when leaving this function!
read_params() {
	# initialize global variables
	declare -g FORCE=0
	unset ACTION REMOTE_HOST CONTAINER_PATH RSYNC_TARGET

	# handle rest of arguments
	while [ "$#" -gt "0" ]; do
		local PARAM="$1"
		shift

		# options to current target
		if [[ "$PARAM" == --* ]]; then
			case "$PARAM" in
				--clone)
					declare -rg ACTION="CLONE"
					;;
				--package)
					declare -rg ACTION="PACKAGE"
					;;
				--export)
					declare -rg ACTION="EXPORT"
					;;
				--update)
					declare -rg ACTION="UPDATE"
					;;
				--host)
					if [ -z "$1" ]; then 
						print_usage "'--host' requires a host as parameter."
					else
						if [[ "$1" == --* ]]; then
							local _msg="A host should not start with '--'."
							_msg="$_msg Parameter for '--host' missing?"
							perror $_msg 
						fi
						declare -rg REMOTE_HOST="$1"
						shift
					fi
					continue
					;;
				--container)
					if [ -z "$1" ]; then 
						print_usage "'--container' requires a path as parameter."
					else
						if [[ "$1" == --* ]]; then
							local _msg="A path should not start with '--'."
							_msg="$_msg Parameter for '--container' missing?"
							perror $_msg 
						fi
						declare -rg CONTAINER_PATH="$1"
						shift
					fi
					continue
					;;
				--syncdir)
					if [ -z "$1" ]; then
						print_usage "'--syncdir' requires a path as parameter."
					else
						if [[ "$1" == --* ]]; then
							local _msg="A path should not start with '--'."
							_msg="$_msg Parameter for '--syncdir' missing?"
							perror $_msg 
						fi
						declare -rg RSYNC_TARGET="$1"
						shift
					fi
					continue
					;;
				--force)
					declare -rg FORCE=1
					;;
				*)
					print_usage "Unknown flag: $PARAM"
					;;
			esac
			continue
		fi
	done
	# done parsing the arguments, exit if no action given
	[ -z "$ACTION" ] && \
		print_usage "No action given"

	# now check for existance of variables
	# and use fallbacks when possible if they were not specified

	# REMOTE_HOST valid? always needed!
	[ -z "$REMOTE_HOST" ] && \
		perror "REMOTE_HOST not set. Use '--host'."
	check_host "$REMOTE_HOST" || \
		perror "'$REMOTE_HOST' is neither an IP nor a known hostname."
	# REMOTE_HOST valid - set build directory for the rest of the operations
	declare -rg BUILD_DIR="${ROOT_DIR}/clones/${REMOTE_HOST}"
	mkdir -p "${BUILD_DIR}" || perror "Could not create '${BUILD_DIR}'."

	# RSYNC_TARGET needs special care
	if [ -z "$RSYNC_TARGET" ]; then
		# none given - use fallbacks
		if [ "x$ACTION" == "xCLONE" ]; then
			# use default path when cloning, no need for CONTAINER_MNT
			pwarning "RSYNC_TARGET not set, using: '${BUILD_DIR}/stage4/'."
			declare -rg RSYNC_TARGET="${BUILD_DIR}/stage4"
		else
			# we always want CONTAINER_MNT here
			declare -rg CONTAINER_MNT="$(mktemp -d)"
			[ -z "${CONTAINER_MNT}" ] && \
				perror "Could not create temporary directory for mounting the container."
			add_cleanup rmdir "${CONTAINER_MNT}"
			# RSYNC_TARGET depends on the action at this point
			if [ "x$ACTION" == "xPACKAGE" ]; then
				# use default path when packaging
				declare -rg RSYNC_TARGET="${BUILD_DIR}/stage4"
			elif [ "x$ACTION" == "xUPDATE" -o "x$ACTION" == "xEXPORT" ]; then
				# for action update/export, we want to sync to the mounted container
				# so create a temporary directory for the mount point that we'll use later
				declare -rg RSYNC_TARGET="${CONTAINER_MNT}"
			fi
		fi
	fi

	# CONTAINER_PATH valid?
	if [ -z "$CONTAINER_PATH" ]; then
		# use default path: ${BUILD_DIR}/stage4.qcow2
		pwarning "CONTAINER_PATH not set. Using '${BUILD_DIR}/stage4.qcow2'."
		declare -rg CONTAINER_PATH="${BUILD_DIR}/stage4.qcow2"
	fi
	# so from now on REMOTE_HOST, RSYNC_TARGET, CONTAINER_PATH are set (and read-only).
}
process_action() {
	if [ "x$ACTION" == "xCLONE" ]; then
		clone_host || perror "Cloning stage4 failed with: $?"
	elif [ "x$ACTION" == "xPACKAGE" ]; then
		pack_clone || perror "Packing as QCoW2 failed with: $?"
	elif [ "x$ACTION" == "xEXPORT" ]; then
		export_host || perror "Exporting failed with: $?"
	elif [ "x$ACTION" == "xUPDATE" ]; then
		update_container || perror "Updating failed with: $?"
	else
		print_usage "No action given."
	fi
	return 0
}

# wrapper to package a cloned stage4 as a qcow2 container
#	- creates empty container at CONTAINER_PATH
#	- mounts it to RSYNC_TARGET
#	- copy RSYNC_TARGET
pack_clone() {
	create_container
	mount_container
	copy_to_container
}
# wrapper to update an existing container
#	- mounts it to RSYNC_TARGET
#	- clone host there
update_container() {
	mount_container
	clone_host
}
# wrapper to export a host directly to a container
#	- create en empty qcow2 container at CONTAINER_PATH
#	- mount it to RSYNC_TARGET
#	- clone host there
export_host() {
	create_container
	mount_container
	clone_host
}
# ------------------------------------------------------------------------------
#
#                          Stage4 related functions
#
# ------------------------------------------------------------------------------
#
# Helper to generate a stage4 export for a remote machine per rsync.
# Usage:
#			clone_host
#
# Note: this functions requires REMOTE_HOST and RSYNC_TARGET to be set.
clone_host() {
	# check if RSYNC_TARGET is valid
	if [ -d "${RSYNC_TARGET}" ]; then
		# does it have the '.stage4' flag? skip this check when exporting directly
		[ "x$ACTION" != "xEXPORT" ] && [ ! -e "${RSYNC_TARGET}/.stage4" ] && \
			perror "'${RSYNC_TARGET}' exists, but no '.stage4' flag found. Refusing to rsync there."
			#touch $RSYNC_TARGET/.stage4 && exit 0
	else
		# not a directory, create it and set the .stage4 flag
		mkdir -p "${RSYNC_TARGET}" || perror "Could not create '${RSYNC_TARGET}'."
	fi

	local EXCLUDE="$BUILD_DIR/exclude-stage4"
	local INCLUDE="$BUILD_DIR/include-stage4"

	pinfo "Building rsync include/exclude files for building stage4...."
	echo "## Exclude file for stage4 of $REMOTE_HOST" > "$EXCLUDE"
	echo "## Include file for stage4 of $REMOTE_HOST" > "$INCLUDE"
	for FILE in $(find "${ROOT_DIR}"/blacklists/*/ -type f); do
		echo "## From $FILE" >> "$EXCLUDE"
		echo "## From $FILE" >> "$INCLUDE"
		grep '^-' "$FILE" >> "$EXCLUDE"
		grep '^+' "$FILE" >> "$INCLUDE"
	done

	# prepare rsync's general options
	local RSYNC_OPTS="${DEFAULT_RSYNC_OPTS}"
	[ -z "${RSYNC_OPTS}" ] && \
		RSYNC_OPTS="
			--acls \
			--hard-links \
			--xattrs \
			--archive \
			--delete \
			--delete-excluded \
			--numeric-ids
		"
	# prepare rsync's remote shell options
	local RSYNC_RSH="$DEFAULT_RSYNC_RSH"
	[ -z "${RSYNC_RSH}" ] && RSYNC_RSH="ssh -c blowfish -oStrictHostKeyChecking=no"
	local RSYNC_SOURCE="root@$REMOTE_HOST:/"

	# if something goes wrong during rsync, we need to recreate the .stage4 flag
	add_cleanup touch ${RSYNC_TARGET}/.stage4
	# run rsync with the exclude/include lists created earlier
	cat "$INCLUDE" "$EXCLUDE" | \
		rsync ${RSYNC_OPTS} \
				--exclude-from=- \
            --rsh "${RSYNC_RSH}" \
				"${RSYNC_SOURCE}" \
				"${RSYNC_TARGET}"

	local -i rsync_ret=$?
	if [ "x$rsync_ret" != "x0" ]; then
		perror "rsync from '${RSYNC_SOURCE}' to '${RSYNC_TARGET}' failed."
	fi
	touch "${RSYNC_TARGET}/.stage4"
	# make sure everything gets flushed
	sync
	pinfo "Cloning '${REMOTE_HOST}' to '${RSYNC_TARGET}' succeeded."
}

# Helper to create the empty container at CONTAINER_PATH
create_container() {
	# CONTAINER_PATH valid?
	[ -d "${CONTAINER_PATH}" ] && perror "Path to container can not be a directory!"
	if [ -f "${CONTAINER_PATH}" ]; then
		if [ $FORCE -eq 0 ]; then
			perror "Container file already exists. Use '--force' to overwrite it."
		else
			# force removal
			rm -f "${CONTAINER_PATH}" || perror "Could not remove '${CONTAINER_PATH}'"
			pinfo "Removed old '${CONTAINER_PATH}'."
		fi
	fi

	# which size for the qcow2 container?
	local CONTAINER_SIZE="${DEFAULT_CONTAINER_SIZE}"
	[ -z "${CONTAINER_SIZE}" ] && CONTAINER_SIZE="20G"

	# so far so good
	pinfo "Creating qcow2-container '${CONTAINER_PATH}'"
	qemu-img create -f qcow2 "${CONTAINER_PATH}" "${CONTAINER_SIZE}" || \
		perror "qemu-img create failed with: $?"
	# now expose it as a loop device
	expose_container
	# filesystem for the qcow2 container?
	local CONTAINER_FILESYSTEM="${DEFAULT_CONTAINER_FILESYSTEM}"
	[ -z "${CONTAINER_FILESYSTEM}" ] && CONTAINER_FILESYSTEM="xfs"
	# check if we have that mkfs helper
	which "mkfs.${CONTAINER_FILESYSTEM}" || \
		perror "Could not find 'mkfs.${CONTAINER_FILESYSTEM}'. Install it and retry." 

	pinfo "Creating '${CONTAINER_FILESYSTEM}' filesystem on '${CONTAINER_PATH}'..."
	mkfs."${CONTAINER_FILESYSTEM}" "${LOOPED_NBD_DEV}" || perror "mkfs failed with: $?"
	return 0
}

# Helper exposing the container as a loop device
expose_container() {
	[ -z "${CONTAINER_PATH}" ] && \
		perror "Internal error - CONTAINER_PATH not set but should be! Check read_params()"
	
	# find usable nbd device
	declare -rg NBD_DEV="$(find_free_nbd)"
	[ -n "${NBD_DEV}" ] || perror "Could not find usable NBD device."
	[ -b "${NBD_DEV}" ] || perror "'${NBD_DEV}' is not a block device!"
	pinfo "Connecting '${CONTAINER_PATH}' to '${NBD_DEV}'"

	qemu-nbd -c "${NBD_DEV}" "${CONTAINER_PATH}" || \
		perror "qemu-nbd failed with: $?"
	add_cleanup disconnect_nbd
	# expose as a loop device
	declare -rg LOOPED_NBD_DEV="$(losetup --find)"
	losetup "${LOOPED_NBD_DEV}" "${NBD_DEV}" || \
		perror "Loop device setup for '${NBD_DEV}' failed with: $?"
	add_cleanup disconnect_loop
	return 0
}

# Helper to mount CONTAINER_PATH to CONTAINER_MNT through expose_container
mount_container() {
	[ -z "${CONTAINER_MNT}" ] && \
		perror "Internal error - CONTAINER_MNT not set but should be! Check read_params()"

	# connect container to a loop device first, if it wasnt already done
	[ -z "${LOOPED_NBD_DEV}" ] && expose_container
	# lets be safe...
	[ -z "${LOOPED_NBD_DEV}" ] && \
		perror "Internal error - LOOPED_NBD_DEV not set but should be! Check expose_container()"
	
	# now we got everything, mount it
	pinfo "Mounting '${LOOPED_NBD_DEV}' to '${CONTAINER_MNT}'..."
	mount "${LOOPED_NBD_DEV}" "${CONTAINER_MNT}" \
		|| perror "Mount failed with: $?"
	add_cleanup umount_container
	return 0
}

# helper to copy the content of RSYNC_TARGET to CONTAINER_MNT
copy_to_container() {
	[ -z "${RSYNC_TARGET}" ] && \
		perror "Internal error - RSYNC_TARGET not set but should be!"
	[ -z "${CONTAINER_MNT}" ] && \
		perror "Internal error - CONTAINER_MNT not set but should be!"
	# sanity checks
	is_dir_empty "$RSYNC_TARGET" && \
		perror "'$RSYNC_TARGET' appears to be empty. Did you clone?"
	# check for '.stage4' flag in the directory, indicating we cloned there
	if [ ! -e "${RSYNC_TARGET}/.stage4" ]; then
		perror "No '.stage4' flag found in '${RSYNC_TARGET}'." \
					"Was this cloned properly?"
	fi
	# copy files from the stage4 directory to the mounted qcow2-container
	pinfo "Copying '${RSYNC_TARGET}' to '${CONTAINER_MNT}'..."
	rsync -avAHX "${RSYNC_TARGET}"/ "${CONTAINER_MNT}"/ \
		|| perror "Rsync failed with: $?"
	#make sure everything is flushed
	sync && return 0
}


###############################################################################
#
#
#
###############################################################################
# helper to find an unused nbd device
# Usage:
#			find_free_nbd
# Echoes the name of the free device to stdout, empty string otherwise.
find_free_nbd() {
	local nbd_size=0
	for nbd_id in {0..15}; do
		[ -b "/dev/nbd${nbd_id}" ] || continue
		[ -r "/sys/block/nbd${nbd_id}/size" ] || continue
		nbd_size=$(cat /sys/block/nbd${nbd_id}/size)
		[ $nbd_size -eq 0 ] && echo "/dev/nbd${nbd_id}" && break
	done
	echo ""
}

# helper to validate an ip
# Usage:
#			valid_ip <ip>
# Returns 0 if valid, 1 otherwise.
valid_ip() {
    local  ip=$1
    local  stat=1

    if [[ $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        OIFS=$IFS
        IFS='.'
        ip=($ip)
        IFS=$OIFS
        [[ ${ip[0]} -le 255 && ${ip[1]} -le 255 \
            && ${ip[2]} -le 255 && ${ip[3]} -le 255 ]]
        stat=$?
    fi
    return $stat
}

# helper to check whether a given host is valid
# Usage:
#			check_host <hostname|ip>
# Returns 0 if valid, 1 otherwise
check_host() {
	local HOST="$1"
	[ -z "$HOST" ] && return 1
	# check if its a valid IP or a valid hostname
	valid_ip "$HOST" && return 0
	host -W 2 "$HOST" && return 0
	# still here? then fail
	return 1
}

# helper to check if a dir is empty or not
# Usage:
#			dir_empty <dir>
# Returns 0 if empty, 1 otherwise
is_dir_empty() {
	[ $# -ne 1 ] && perror "$0 requires directory as paramter, none given."
	local _dir="$1"
	[ -d "$_dir" ] || return 1
	[ -n "$(ls -A $_dir)" ] && return 1 || return 0
}

# helper to ask user for confirmation
# Usage:
#			user_confirm
# Return 0 if confirmed, 1 otherwise
user_confirm() {
	[ $# -ne 1 ] && perror "$0 requires the question as first argument."
	pinfo "$1 [Y/n]"
	local _input
	read _input
	[ "x${_input}" == "x" -o "x${_input}" == "xy" ] && return 0 || return 1
}

###############################################################################
#
# CLEANUP FUNCTIONS
#
###############################################################################
# indexed array for cleanup commands
declare -ag cleanup_commands

# function to add a cleanup command to be executed when the program exits
add_cleanup() {
	# get current command count
	local -i count="${#cleanup_commands[*]}"
	cleanup_commands[$count]="$*"
}

# function trapped to EXIT, SIGINT, SIGTERM
# do the cleanup in FILO style
do_cleanup() {
	trap '' SIGINT SIGTERM EXIT	# from now on, ignore INT and TERM
	for i in $(seq $(( ${#cleanup_commands[*]} - 1 )) -1 0); do
		eval ${cleanup_commands[$i]}
	done
}

# Helper to umount + disconnect the container from all the devices
umount_container() {
	[ -z "${CONTAINER_MNT}" ] && \
		perror "CONTAINER_MNT not set - is it really mounted?"
	# sync?
	umount -l "${CONTAINER_MNT}" || \
		perror "Failed to umount '${CONTAINER_MNT}'."
#	rmdir "${CONTAINER_MNT}" || \
#		pwarning "Could not remove '${CONTAINER_MNT}'."
}

# Helper to disconnect from loop device
disconnect_loop() {
	[ -z "${LOOPED_NBD_DEV}" ] && \
		perror "Container not connected to a loop device?"
	losetup -d "${LOOPED_NBD_DEV}" ||\
		perror "Could not disconnect loop device '${LOOPED_NBD_DEV}'."
}
# Helper to disconnect from nbd device
disconnect_nbd() {
	[ -z "${NBD_DEV}" ] && \
		perror "Container does not seem to be connected to a NBD device?"
	qemu-nbd -d "${NBD_DEV}" || \
		perror "Could not disconnect '${CONTAINER_PATH}' from '${NBD_DEV}'."
}