#!/bin/sh
#
# Purpose {{{
# This script will set a quota for all users of a given ZFS pool or dataset
#   1. Verify the ZFS pool/dataset is reachable
#   2. Get users list of pool/dataset
# …
#
# 2021-11-22
# }}}
# Vars {{{
PROGNAME=$(basename "${0}"); readonly PROGNAME
PROGDIR=$(readlink -m $(dirname "${0}")); readonly PROGDIR
ARGS="${*}"; readonly ARGS
readonly NBARGS="${#}"
[ -z "${DEBUG}" ] && DEBUG=1
## Export DEBUG for sub-script
export DEBUG

# Default values for some vars
readonly ZFS_DATASET_DEFAULT="datastore"
readonly ZFS_QUOTA_DEFAULT="200G"

# Disable QUOTA_FILE_MODE by default
QUOTA_FILE_MODE=1

## Colors
readonly PURPLE='\033[1;35m'
readonly RED='\033[0;31m'
readonly RESET='\033[0m'
readonly COLOR_DEBUG="${PURPLE}"
# }}}
usage() {                                                       # {{{

	cat <<- EOF
usage: $PROGNAME [-d|-h|-p|-q]

Define a quota to all users of a given ZFS pool/dataset

EXAMPLES :
    - Apply default quota (${ZFS_QUOTA_DEFAULT}) to default pool (${ZFS_DATASET_DEFAULT})
        ${PROGNAME}

    - Apply default quota to a given dataset
        ${PROGNAME} --pool "datastore/backup"

    - Apply a given quota by default (100G) and use a quota file for specific user
        ${PROGNAME} --quota 100G --file /datastore/backup/.duplicati/zfs.quota

OPTIONS :
    -d,--debug
        Enable debug messages.

    -f|--file
        Specify path to a quota file for specific quota per user.

    -h,--help
        Print this help message.

    -p,--pool,--dataset
        Apply quota to the given pool/dataset (default: ${ZFS_DATASET_DEFAULT}).

    -q,--quota,--quotat
        Define the quota to apply to all users of the pool/dataset (default: ${ZFS_QUOTA_DEFAULT}).

	EOF

}
# }}}
debug_message() {                                               # {{{

	local_message="${1}"

	## Print message if DEBUG is enable (=0)
	[ "${DEBUG}" -eq "0" ] && printf '\e[1;35m%-6b\e[m\n' "DEBUG − ${PROGNAME} : ${local_message}"

	return 0
}
# }}}
error_message() {                                               # {{{

	local_error_message="${1}"
	local_error_code="${2}"

	## Print message if DEBUG is enable (=0)
	printf '%b\n' "ERROR − ${PROGNAME} : ${RED}${local_error_message}${RESET}"

	exit "${local_error_code:=66}"
}
# }}}
define_vars() {                                                 # {{{

	## If zfs_dataset wasn't defined (argument) {{{
	if [ -z "${zfs_dataset}" ]; then
		## Use default value
		readonly zfs_dataset="${ZFS_DATASET_DEFAULT}"
	fi
	## }}}
	## If zfs_quota wasn't defined (argument) {{{
	if [ -z "${zfs_quota}" ]; then
		## Use default value
		readonly zfs_quota="${ZFS_QUOTA_DEFAULT}"
	fi
	## }}}
	## If quota_file was defined (argument) {{{
	if [ -n "${quota_file}" ]; then
		## If the file is exists
		## OR exit
		is_file_present "${quota_file}" \
			|| error_message "Given quota file (${quota_file}) isn't readable." 10

		## Enable QUOTA_FILE_MODE
		debug_message  "define_vars − \
${quota_file} will be used to get specific quota for defined users."
		QUOTA_FILE_MODE=0
	fi
	## }}}

	## Temp file vars {{{
	readonly zfs_user_list_path="/tmp/${PROGNAME}.user.list"
	readonly zfs_previous_user_list_path="/tmp/${PROGNAME}.old.user.list"
	## }}}

}
# }}}
is_command_absent() {                                           # {{{

	local_command_absent_cmd="${1}"

	## A command is absent by default
	return_command_absent="0"

	if [ "$(command -v ${local_command_absent_cmd})" ]; then
		debug_message "is_command_absent − \
${RED}${local_command_absent_cmd}${COLOR_DEBUG} seems present on this host."
		return_command_absent="1"
	else
		debug_message "is_command_absent − \
${RED}${local_command_absent_cmd}${COLOR_DEBUG} is not available on this host."
		return_command_absent="0"
	fi

	return "${return_command_absent}"

}
# }}}
is_zfs_dataset_exists() {                                       # {{{

	local_zfs_dataset="${1}"

	## Return False by default
	return_zfs_dataset_exists="1"

	## Use local_zfs_dataset var in zfs command and grep to avoid sub-datasets
	if zfs list -H -- "${local_zfs_dataset}" 2>/dev/null | grep --quiet "${local_zfs_dataset}"; then
		debug_message "is_zfs_dataset_exists − \
${RED}${local_zfs_dataset}${COLOR_DEBUG} ZFS pool/dataset seems present on this host."
		return_zfs_dataset_exists="0"
	else
		debug_message "is_zfs_dataset_exists − \
${RED}${local_zfs_dataset}${COLOR_DEBUG} ZFS pool/dataset is not available on this host."
		return_zfs_dataset_exists="1"
	fi

	return "${return_zfs_dataset_exists}"

}
# }}}
get_dataset_user_list() {                                       # {{{

	local_zfs_dataset="${1}"

	## Return False by default
	return_get_dataset_user_list="1"

	debug_message "get_dataset_user_list − \
Create or empty ${RED}${zfs_user_list_path}${COLOR_DEBUG} file to store user list of ${RED}${local_zfs_dataset}${COLOR_DEBUG} ZFS pool/dataset."
	true > "${zfs_user_list_path}"

	if zfs userspace -Hp -o name "${local_zfs_dataset}" >> "${zfs_user_list_path}" 2>/dev/null; then
		if [ -s "${zfs_user_list_path}" ]; then
			debug_message "get_dataset_user_list − \
${RED}${local_zfs_dataset}${COLOR_DEBUG} users list successfully created (see ${zfs_user_list_path} file)."
			command chmod 0400 -- "${zfs_user_list_path}"
			return_get_dataset_user_list="0"
		else
			debug_message "get_dataset_user_list − \
Error, the user list of ${local_zfs_dataset} is empty (${zfs_user_list_path} file)."
			return_get_dataset_user_list="1"
		fi
	else
		debug_message "get_dataset_user_list − \
Error in ${RED}zfs userspace${COLOR_DEBUG} command for ${local_zfs_dataset} ZFS pool/dataset."
		return_get_dataset_user_list="1"
	fi

	return "${return_get_dataset_user_list}"

}
# }}}
is_file_present() {                                             # {{{

	local_file_present="${1}"

	## File doesn't exist by default
	return_is_file_present="1"

	### Check if the file exists
	# shellcheck disable=SC2086
	if find ${local_file_present} > /dev/null 2>&1; then
		return_is_file_present="0"
		debug_message "is_file_present − \
The file ${RED}${local_file_present}${COLOR_DEBUG} exists."
	else
		return_is_file_present="1"
		debug_message "is_file_present − \
The file ${RED}${local_file_present}${COLOR_DEBUG} doesn't exist."
	fi

	return "${return_is_file_present}"

}
# }}}
is_file_similar() {                                             # {{{

	local_similar_file_one="${1}"
	local_similar_file_two="${2}"

	## Files aren't similar by default doesn't exist by default
	return_is_file_similar="1"

	if diff --brief -- "${local_similar_file_one}" "${local_similar_file_two}" > /dev/null; then
		debug_message "is_file_similar − \
${local_similar_file_one} and ${local_similar_file_two} are ${RED}similar${COLOR_DEBUG}."
		return_is_file_similar="0"
	else
		debug_message "is_file_similar − \
${local_similar_file_one} and ${local_similar_file_two} are ${RED}NOT${COLOR_DEBUG} similar."
		return_is_file_similar="1"
	fi

	return "${return_is_file_similar}"
}
# }}}
apply_zfs_userquota() {                                         # {{{

	local_zfs_user="${1}"
	local_zfs_quota="${2}"

	## Return False by default
	return_apply_zfs_userquota="1"

	## Use local_zfs_dataset var in zfs command and grep to avoid sub-datasets
	if zfs userquota@"${local_zfs_user}"="${local_zfs_quota}" "${local_zfs_dataset}" 2>/dev/null; then
		debug_message "apply_zfs_userquota − \
Set new quota (${RED}${local_zfs_quota}${COLOR_DEBUG}) for ${RED}${local_zfs_user}${COLOR_DEBUG} user."
		return_apply_zfs_userquota="0"
	else
		debug_message "apply_zfs_userquota − \
Can't set a new quota (${RED}${local_zfs_quota}${COLOR_DEBUG}) for ${RED}${local_zfs_user}${COLOR_DEBUG} user \
('zfs userquota' command returns : ${?})"
		return_apply_zfs_userquota="1"
	fi

	return "${return_apply_zfs_userquota}"

}
# }}}

main() {                                                        # {{{

	## This script should run as root {{{
	if ! [ $(id -u) = 0 ]; then
		error_message "Please run this script in root or with sudo." 1
	fi
	## }}}
	## If ZFS command is absent from the system {{{
	### Exit
	is_command_absent "zfs" \
		&& exit 0
	## }}}

	## Define all vars
	define_vars

	## If ZFS pool/dataset is not available {{{
	### Exit
	is_zfs_dataset_exists "${zfs_dataset}" \
		|| error_message "Please verify your ZFS pool (${zfs_dataset} doesn't seems available" 2
	## }}}

	## Try to get the user list of ZFS pool/dataset {{{
	### OR Exit
	get_dataset_user_list "${zfs_dataset}" \
		|| error_message "Can't get the user list of ${zfs_dataset} ZFS pool/dataset. Please use --debug option." 3
	## }}}

	## If a previous list of users exists {{{
	### If the two list are the same
	### Exit
	is_file_present "${zfs_previous_user_list_path}" \
		&& is_file_similar "${zfs_user_list_path}" "${zfs_previous_user_list_path}" \
		&& debug_message "main − No new user from previous run, no more actions required." \
		&& exit 0
	## }}}

	## Information message
	debug_message "Apply quota (default : ${RED}${zfs_quota}${COLOR_DEBUG}) to all users of ZFS pool/dataset : ${RED}${zfs_dataset}${COLOR_DEBUG}"

	## Read user one by one
	while IFS= read -r zfs_username; do
		### Define quota to use {{{
		if [ "${QUOTA_FILE_MODE}" -eq "0" ] \
		&& grep --word-regexp -- "^${zfs_username}" "${quota_file}" | grep --only-matching --perl-regexp --quiet -- '[[:digit:].]*.$'; then
			### From quota file if any information for this user
			user_zfs_quota="$(grep --word-regexp -- "^${zfs_username}" "${quota_file}" |
				grep --only-matching --perl-regexp -- '[[:digit:].]*.$' \
				|| error_message "Can't get quota for ${zfs_username} user from ${quota_file} quota file." 4)"
			debug_message "Get specific quota (${RED}${user_zfs_quota}${COLOR_DEBUG}) for ${RED}${zfs_username}${COLOR_DEBUG} user from quota file (${RED}${quota_file}${COLOR_DEBUG})."
		else
			### Default quota by default
			user_zfs_quota="${zfs_quota}"
		fi
		### }}}
		### Try to apply the quota to the user {{{
		### OR Exit with error
		apply_zfs_userquota "${zfs_username}" "${user_zfs_quota}" \
		|| error_message "Can't define the new quota (${user_zfs_quota}) for ${zfs_username} user in ${zfs_dataset} ZFS pool/dataset. Please use --debug option." 4
		### }}}
	done < "${zfs_user_list_path}"

	## Rename user list for next run of the script
	## AND exit
	mv --force -- "${zfs_user_list_path}" "${zfs_previous_user_list_path}" \
		&& exit 0

}
# }}}

# Manage arguments                                                # {{{
# This code can't be in a function due to argument management

if [ ! "${NBARGS}" -eq "0" ]; then

	manage_arg="0"

	## If the first argument is not an option
	if ! printf -- '%s' "${1}" | grep -q -E -- "^-+";
	then
		## Print help message and exit
		printf '%b\n' "${RED}Invalid option: ${1}${RESET}"
		printf '%b\n' "---"
		usage

		exit 1
	fi

	# Parse all options (start with a "-") one by one
	while printf -- '%s' "${1}" | grep -q -E -- "^-+"; do

	case "${1}" in
		-d|--debug )                          ## debug
			DEBUG=0
			;;
		-f|--file )                           ## Quota file
			## Move to the next argument
			shift
			## Define var
			readonly quota_file="${1}"
			;;
		-h|--help )                           ## help
			usage
			## Exit after help informations
			exit 0
			;;
		-p|--pool|--dataset )                 ## Define zfs_dataset
			## Move to the next argument
			shift
			## Define var
			readonly zfs_dataset="${1}"
			;;
		-q|--quota|--quotat )                 ## Define zfs_quota
			## Move to the next argument
			shift
			## Define var
			readonly zfs_quota="${1}"
			;;
		* )                                   ## unknow option
			printf '%b\n' "${RED}Invalid option: ${1}${RESET}"
			printf '%b\n' "---"
			usage
			exit 1
			;;
	esac

	debug_message "Arguments management − \
${RED}${1}${COLOR_DEBUG} option managed."

	## Move to the next argument
	shift
	manage_arg=$((manage_arg+1))

	done

	debug_message "Arguments management − \
${RED}${manage_arg}${COLOR_DEBUG} argument(s) successfully managed."
else
	debug_message "Arguments management − \
No arguments/options to manage."
fi

# }}}

main

exit 255