#!/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

[ -x ${IPSET_BIN} ] || { echo "ipset: Cannot execute ${IPSET_BIN}" >&2; exit 5; }

# Source ipset configuration
[ -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 -pm 700 ${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.
		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 ] && echo "ipset: Cannot save set ${set}" >&2 && 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
}

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} -a ! -f ${IPSET_DATA_COMPAT} ]; then
		echo "ipset: 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
		for f in ${IPSET_DATA_DIR}/*; do
			[ "${f}" = "${IPSET_DATA_DIR}/*" ] && break
			cat "${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
		${IPSET_BIN} destroy "${set}" 2>/dev/null
		# We can't destroy the set if it's in use, flush it instead
		[ $? -ne 0 ] && ${IPSET_BIN} flush "${set}"
	done
	unset IFS
	conflicts="${conflicts#|*}"

	# Common case: if we have no conflicts, just restore in one shot
	if [ -z "${conflicts}" ]; then
		${IPSET_BIN} restore -! < ${merged}
		[ $? -ne 0 ] && echo "ipset: Failed to restore configured sets" >&2
		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
			grep -q "^create $(echo ${salt}${set} | md5sum | head -c31) " ${merged}
			[ $? -eq 0 ] && unique=0 && break
		done
		unset IFS
		[ ${unique} -eq 1 ] && break
		salt=$((salt + 1))
	done

	# Add sets, mangling names for conflicting sets
	awk '/^(add|create) ('"${conflicts}"')/ { printf "%s ",$1; system("echo '${salt}'" $2 " | md5sum | head -c31"); $1=""; $2=""; print; next} {print}' ${merged} | ${IPSET_BIN} restore -!

	[ $? -ne 0 ] && echo "ipset: Failed to restore configured sets" >&2

	# Swap and delete old sets
	IFS='|'
	for set in ${conflicts}; do
		mangled="$(echo ${salt}${set} | md5sum | head -c31)"
		${IPSET_BIN} swap "${set}" "${mangled}" 2>/dev/null
		if [ $? -ne 0 ]; then
			# This fails if set types are different: try to destroy
			# existing set
			${IPSET_BIN} destroy "${set}" 2>/dev/null
			if [ $? -ne 0 ]; then
				# Conflicting set is in use, we can only warn
				# and flush the existing set
				echo "ipset: Cannot load set \"${set}\", set with same name and conflicting type in use" >&2
				${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 || echo "ipset: Failed to flush sets" >&2

	# 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
		${IPSET_BIN} destroy "${set}" 2>/dev/null
	done
	unset IFS
}

stop() {
	[ -f ${IPSET_RUN} ] || { echo "ipset: not running"; return 0; }
	[ "${IPSET_SAVE_ON_STOP}" = "yes" ] && { save || echo "ipset: Failed to save sets"; }

	# Nothing to stop if the ip_set module is not loaded
	lsmod | grep -q "^ip_set " || { echo "ipset: 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
	mod="$(lsmod | grep ^xt_set)"
	if [ $? -eq 0 ]; then
		if [ "$(echo ${mod} | tr -s ' ' | cut -d' ' -f3)" != "0" ]; then
			echo "ipset: Current iptables configuration requires ipset" >&2 && return 1
		fi
	fi

	cleanup

	rm ${IPSET_RUN}
	return 0
}

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

exit $?