#!/bin/sh
#
# ipset      Start and stop ipset firewall sets
#
# config: /etc/sysconfig/ipset-config

IPSET_BIN=/usr/sbin/ipset
IPSET_CONFIG=/etc/sysconfig/ipset-config
IPSET_DATA_COMPAT=/etc/sysconfig/ipset
IPSET_DATA_COMPAT_BACKUP=${IPSET_DATA_COMPAT}.save
IPSET_DATA_DIR=/etc/sysconfig/ipset.d
IPSET_DATA_DIR_BACKUP=${IPSET_DATA_DIR}.save
IPSET_DATA_SAVED_FLAG=${IPSET_DATA_DIR}/.saved
IPSET_LOCK=/run/ipset.lock
IPSET_RUN=/run/ipset.run
CLEAN_FILES=""

trap "rm -rf \${CLEAN_FILES}" EXIT

info() {
	echo "ipset: ${*}" >&2
}

warn() {
	echo "<4>ipset: ${*}" >&2
}

err() {
	echo "<3>ipset: ${*}" >&2
}

[ -x ${IPSET_BIN} ] || { err "Cannot execute ${IPSET_BIN}"; exit 1; }

# Source ipset configuration
# shellcheck source=ipset-config
[ -f ${IPSET_CONFIG} ] && . ${IPSET_CONFIG}

set -f

lock() {
	CLEAN_FILES="${CLEAN_FILES} ${IPSET_LOCK}"
	until mkdir ${IPSET_LOCK} 2>/dev/null; do :; done
}

save() {
	fail=0

	# Make backups of existing configuration first, if any
	[ -d ${IPSET_DATA_DIR} ] && mv -Tf ${IPSET_DATA_DIR} ${IPSET_DATA_DIR_BACKUP}
	[ -f ${IPSET_DATA_COMPAT} ] && mv -Tf ${IPSET_DATA_COMPAT} ${IPSET_DATA_COMPAT_BACKUP}

	rm -f ${IPSET_DATA_SAVED_FLAG}

	# Save each set in a separate file
	mkdir -p ${IPSET_DATA_DIR}
	chmod 0700 ${IPSET_DATA_DIR}
	IFS="
"
	for set in $(${IPSET_BIN} list -n -t); do
		# Empty name allowed, use ".set" as suffix. 'ipset save' doesn't
		# quote set names with spaces: if we have a space in the name,
		# work around this by quoting it ourselves in the output.
		# shellcheck disable=SC2003 # No POSIX equivalent to expr index
		if expr index "${set}" " " >/dev/null; then
			:> "${IPSET_DATA_DIR}/${set}.set"
			for line in $(${IPSET_BIN} save "${set}"); do
				create=0
				echo "${line}" | grep -q "^create " && create=1
				if [ $create -eq 1 ]; then
					line=${line#create *}
				else
					line=${line#add *}
				fi
				line=${line#${set} *}
				set="$(echo "${set}" | sed 's/"/\\"/g')"
				if [ $create -eq 1 ]; then
					echo "create \"${set}\" ${line}" >> "${IPSET_DATA_DIR}/${set}.set"
				else
					echo "add \"${set}\" ${line}" >> "${IPSET_DATA_DIR}/${set}.set"
				fi
			done
		else
			${IPSET_BIN} save "${set}" > "${IPSET_DATA_DIR}/${set}.set" || fail=1
		fi
		[ -f "${IPSET_DATA_DIR}/${set}.set" ] && chmod 600 "${IPSET_DATA_DIR}/${set}.set"
		[ $fail -eq 1 ] && err "Cannot save set ${set}" && unset IFS && return 1
	done
	touch ${IPSET_DATA_SAVED_FLAG} || { unset IFS; return 1; }
	unset IFS

	# Done: remove backups
	rm -rf ${IPSET_DATA_DIR_BACKUP}
	rm -rf ${IPSET_DATA_COMPAT_BACKUP}

	return 0
}

# Generate a grep regexp matching abbreviated command forms. E.g., for create:
# \(c\|cr\|cre\|crea\|creat\|create\)
cmd_short_expr() {
	out=
	cmd_len=1
	while [ "${cmd_len}" -le "${#1}" ]; do
		[ -z "${out}" ] && out='\(' || out="${out}"'\|'
		# shellcheck disable=SC2003 # No POSIX equivalent to expr substr
		out="${out}$(expr substr "${1}" 1 "${cmd_len}")"
		cmd_len=$((cmd_len + 1))
	done
	echo "${out}"'\)'
}

ipset_restore() {
	file="${1}"

	retfile="$(mktemp -q /tmp/ipset.XXXXXX)"
	CLEAN_FILES="${CLEAN_FILES} ${retfile}"

	# If restore fails due to invalid lines, remove them and retry
	while ! restore_err="$( (${IPSET_BIN} -f "${file}" -! restore 2>&1; echo $? >"${retfile}") | head -n1; exit "$(cat "${retfile}")" )"; do
		warn "${restore_err}"
		case ${restore_err#*: } in
		"No command specified"*)
			line="$(grep -m1 -n "^${restore_err##* }" "${file}")"
			line="${line%:*}"
			;;
		"Missing second mandatory argument to command "*)
			cmd="${restore_err##* }"
			cmd_expr="$(cmd_short_expr "${cmd}")"
			line="$(grep -n '^'"${cmd_expr}" "${file}" | grep -m1 -v '^[0-9]\+\:'"${cmd_expr}"'[[:blank:]]\+[^[:blank:]]\+[[:blank:]]\+[^[:blank:]]\+')"
			line="${line%:*}"
			;;
		"Missing mandatory argument to command "*)
			cmd="${restore_err##* }"
			cmd_expr="$(cmd_short_expr "${cmd}")"
			line="$(grep -n '^'"${cmd_expr}" "${file}" | grep -m1 -v '^[0-9]\+\:'"${cmd_expr}"'[[:blank:]]\+[^[:blank:]]\+')"
			line="${line%:*}"
			;;
		"Command "*"is invalid in restore mode"*)
			restore_err_cmd="${restore_err#*: }"
			restore_err_cmd="${restore_err_cmd#*\`}"
			restore_err_cmd="${restore_err_cmd%%\'*}"
			cmd="${restore_err_cmd##* }"
			cmd_expr="$(cmd_short_expr "${cmd}")"
			line="$(grep -m1 -ne '^'"${cmd_expr}"'[[:blank:]]\+' -e '^'"${restore_err_cmd}"'$' "${file}")"
			line="${line%:*}"
			;;
		"Error in line "*)
			line="${restore_err%: *}"
			line="${line##* }"
			;;
		*)
			rm "${retfile}"
			CLEAN_FILES="${CLEAN_FILES%* ${retfile}}"
			return 1
			;;
		esac

		[ -z "${line}" ] && return 1

		warn "Skipped invalid entry: $(sed "${line}q;d" "${file}")"
		sed -i -e "${line}d" "${file}"

		[ -s "${file}" ] || return 1
	done

	rm "${retfile}"
	CLEAN_FILES="${CLEAN_FILES%* ${retfile}}"
}

load() {
	if [ -f ${IPSET_DATA_SAVED_FLAG} ]; then
		# If we have a cleanly saved directory with all sets, we can
		# delete any left-overs and use it
		rm -rf ${IPSET_DATA_DIR_BACKUP}
		rm -f ${IPSET_DATA_COMPAT_BACKUP}
	else
		# If sets weren't cleanly saved, restore from backups
		[ -d ${IPSET_DATA_DIR_BACKUP} ] && rm -rf ${IPSET_DATA_DIR} && mv -Tf ${IPSET_DATA_DIR_BACKUP} ${IPSET_DATA_DIR}
		[ -f ${IPSET_DATA_COMPAT_BACKUP} ] && rm -f ${IPSET_DATA_COMPAT} && mv -Tf ${IPSET_DATA_COMPAT_BACKUP} ${IPSET_DATA_COMPAT}
	fi

	if [ ! -d ${IPSET_DATA_DIR} ] && [ ! -f ${IPSET_DATA_COMPAT} ]; then
		info "No existing configuration available, none loaded"
		touch ${IPSET_RUN}
		return
	fi

	# Merge all sets into temporary file
	merged="$(mktemp -q /tmp/ipset.XXXXXX)"
	CLEAN_FILES="${CLEAN_FILES} ${merged}"
	chmod 600 "${merged}"
	set +f
	if [ -d ${IPSET_DATA_DIR} ]; then
		# Copy create commands from each saved set first, then the rest:
		# list:set entries depend on other sets, so make sure they all
		# get created first
		for f in "${IPSET_DATA_DIR}"/*; do
			[ "${f}" = "${IPSET_DATA_DIR}/*" ] && break
			[ -f "${f}" ] || continue
			grep '^c' "${f}" >> "${merged}"
		done
		for f in "${IPSET_DATA_DIR}"/*; do
			[ "${f}" = "${IPSET_DATA_DIR}/*" ] && break
			[ -f "${f}" ] || continue
			grep -v '^c' "${f}" >> "${merged}"
		done
	fi
	set -f
	[ -f ${IPSET_DATA_COMPAT} ] && cat ${IPSET_DATA_COMPAT} >> "${merged}"

	# Drop sets that aren't in saved data, mark conflicts with existing sets
	conflicts=""
	IFS="
"
	for set in $(${IPSET_BIN} list -n -t); do
		grep -q "^create ${set} " "${merged}" && conflicts="${conflicts}|${set}" && continue

		# We can't destroy the set if it's in use, flush it instead
		if ! ${IPSET_BIN} destroy "${set}" 2>/dev/null; then
			${IPSET_BIN} flush "${set}"
		fi
	done
	unset IFS
	conflicts="${conflicts#|*}"

	# Common case: if we have no conflicts, just restore in one shot
	if [ -z "${conflicts}" ]; then
		if ! ipset_restore "${merged}"; then
			err "Failed to restore configured sets"
			exit 1
		fi
		rm "${merged}"
		CLEAN_FILES="${CLEAN_FILES%* ${merged}}"
		touch ${IPSET_RUN}
		return
	fi

	# Find a salt for md5sum that makes names of saved sets unique
	salt=0
	while true; do
		unique=1
		IFS="
"
		for set in $(${IPSET_BIN} list -n -t); do
			if grep -q "^create $(echo "${salt}${set}" | md5sum | head -c31) " "${merged}"; then
				unique=0
				break
			fi
		done
		unset IFS
		[ ${unique} -eq 1 ] && break
		salt=$((salt + 1))
	done

	# Add sets, mangling names for conflicting sets
	mangled="$(mktemp -q /tmp/ipset.XXXXXX)"
	CLEAN_FILES="${CLEAN_FILES} ${mangled}"
	chmod 600 "${mangled}"

	awk '/^(add|create) ('"${conflicts}"')/ { printf "%s ",$1; system("echo '${salt}'" $2 " | md5sum | head -c31"); $1=""; $2=""; print; next} {print}' "${merged}" > "${mangled}"
	if ! ipset_restore "${mangled}"; then
		err "Failed to restore configured sets"
		exit 1
	fi

	rm "${mangled}"
	CLEAN_FILES="${CLEAN_FILES%* ${mangled}}"

	# Swap and delete old sets
	IFS='|'
	for set in ${conflicts}; do
		mangled="$(echo "${salt}${set}" | md5sum | head -c31)"
		if ! ${IPSET_BIN} swap "${set}" "${mangled}" 2>/dev/null; then
			# This fails if set types are different: try to destroy
			# existing set
			if ! ${IPSET_BIN} destroy "${set}" 2>/dev/null; then
				# Conflicting set is in use, we can only warn
				# and flush the existing set
				err "Cannot load set \"${set}\", set with same name and conflicting type in use"
				${IPSET_BIN} flush "${set}"
				${IPSET_BIN} destroy "${mangled}"
			else
				${IPSET_BIN} rename "${mangled}" "${set}"
			fi
		else
			${IPSET_BIN} destroy "${mangled}"
		fi
	done
	unset IFS

	rm "${merged}"
	CLEAN_FILES="${CLEAN_FILES%* ${merged}}"
	touch ${IPSET_RUN}
}

cleanup() {
	${IPSET_BIN} flush || err "Failed to flush sets"

	# Try to destroy all sets at once. This will fail if some are in use,
	# destroy all the other ones in that case
	${IPSET_BIN} destroy 2>/dev/null && return
	IFS="
"
	for set in $(${IPSET_BIN} list -n -t); do
		if ! ${IPSET_BIN} destroy "${set}"; then
			err "Failed to destroy set ${set}"
		fi
	done
	unset IFS
}

stop() {
	[ -f ${IPSET_RUN} ] || { info "Not running"; return 0; }
	[ "${IPSET_SAVE_ON_STOP}" = "yes" ] && { save || err "Failed to save sets"; }

	# Nothing to stop if the ip_set module is not loaded
	lsmod | grep -q "^ip_set " || { info "Not running"; rm ${IPSET_RUN}; return 0; }

	# If the xt_set module is in use, then iptables is using ipset, so
	# refuse to stop the service
	if mod="$(lsmod | grep ^xt_set)"; then
		if [ "$(echo "${mod}" | tr -s ' ' | cut -d' ' -f3)" != "0" ]; then
			err "Current iptables configuration requires ipset" && return 1
		fi
	fi

	cleanup

	rm ${IPSET_RUN}
	return 0
}

lock
case "$1" in
start)
	load
	;;
stop)
	stop
	;;
reload)
	cleanup
	load
	;;
save)
	save
	;;
*)
	info "Usage: $0 {start|stop|reload|save}"
	exit 1
esac

exit $?