#!/bin/bash
###############################################################################
#                                                                             #
# IPFire.org - A linux based firewall                                         #
# Copyright (C) 2012-2013  IPFire Network Development Team                    #
#                                                                             #
# This program is free software: you can redistribute it and/or modify        #
# it under the terms of the GNU General Public License as published by        #
# the Free Software Foundation, either version 3 of the License, or           #
# (at your option) any later version.                                         #
#                                                                             #
# This program is distributed in the hope that it will be useful,             #
# but WITHOUT ANY WARRANTY; without even the implied warranty of              #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
# GNU General Public License for more details.                                #
#                                                                             #
# You should have received a copy of the GNU General Public License           #
# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
#                                                                             #
###############################################################################

IPTABLES_TABLES="filter mangle nat"

function iptables() {
	local protocol="${1}"
	assert isset protocol
	shift

	# Rules go to the filter table by default
	local table="filter"

	# Argument list
	local args

	# Cached arguments
	local src dst

	# Parsing arguments
	while [ $# -gt 0 ]; do
		case "${1}" in
			# Filter to which table this rule should go.
			-t)
				table="${2}"
				shift 2

				assert isoneof table ${IPTABLES_TABLES}
				;;

			# Automatically convert ICMP to ICMPv6 for IPv6
			--protocol|-p)
				local proto="${2}"

				if [ "${protocol}" = "ipv6" -a "${proto}" = "icmp" ]; then
					proto="icmpv6"
				fi

				list_append args "${1} ${proto}"
				shift 2
				;;
			*)
				list_append args "${1}"

				# Save some values for further processing.
				case "${1}" in
					-s)
						src="${2}"
						;;
					-d)
						dst="${2}"
						;;
				esac
				shift
				;;
		esac
	done

	assert isset action

	# Check if given IP addresses or networks match the protocol version.
	local src_proto
	if isset src; then
		src_proto="$(ip_detect_protocol ${src})"

		assert [ "${protocol}" = "${src_proto}" ]
	fi

	local dst_proto
	if isset dst; then
		dst_proto="$(ip_detect_protocol ${dst})"

		assert [ "${protocol}" = "${dst_proto}" ]
	fi

	# Check if the directory where we put our rules in is set and
	# exists.
	assert isset IPTABLES_TMPDIR
	local rulesfile="${IPTABLES_TMPDIR}/${protocol}-${table}"

	print "${args}" >> "${rulesfile}"
	assert_check_retval $?
}

function iptables_chain_create() {
	local protocol="${1}"
	assert isset protocol
	shift

	local chain
	local table="filter"
	local policy="-"

	while [ $# -gt 0 ]; do
		case "${1}" in
			-t)
				table="${2}"
				shift
				;;
			--policy=*)
				policy="$(cli_get_val ${1})"
				;;
			-*)
				log WARNING "Unrecognized argument: ${1}"
				;;
			*)
				chain=${1}
				;;
		esac
		shift
	done

	assert isset chain
	assert isset table
	assert isoneof policy ACCEPT DROP "-"

	iptables "${protocol}" -t "${table}" ":${chain} ${policy} [0:0]"
}

# Calls the binary iptables command.
function _iptables() {
	local protocol="${1}"
	assert isset protocol
	shift

	local cmd
	case "${protocol}" in
		ipv6)
			cmd="ip6tables"
			;;
		ipv4)
			cmd="iptables"
			;;
	esac
	assert isset cmd
	cmd="$(which ${cmd})"

	cmd "${cmd}" "$@"
	return $?
}

function iptables_status() {
	local protocol="${1}"
	assert isset protocol

	local table
	for table in ${IPTABLES_TABLES}; do
		print "${protocol} - ${table}:"
		_iptables "${protocol}" -t "${table}" -L -n -v
		print
	done

	return ${EXIT_OK}
}

function iptables_rulesfile() {
	local proto=${1}
	proto=${proto/ipv/}

	local chain=${2}
	[ -z "${chain}" ] && chain="ruleset"

	print "${IPTABLES_TMPDIR}/${chain}${proto}"
}

function iptables_init() {
	local protocol="${1}"
	assert isset protocol

	local policy="${2}"
	assert isset policy

	# Create filter table and initialize chains.
	iptables "${protocol}" "* filter"
	iptables_chain_create "${protocol}" -t filter INPUT   --policy="${policy}"
	iptables_chain_create "${protocol}" -t filter OUTPUT  --policy="${policy}"
	iptables_chain_create "${protocol}" -t filter FORWARD --policy="${policy}"

	# Create mangle table and initialize chains.
	iptables "${protocol}" -t mangle "* mangle"
	iptables_chain_create "${protocol}" -t mangle PREROUTING  --policy="ACCEPT"
	iptables_chain_create "${protocol}" -t mangle INPUT       --policy="ACCEPT"
	iptables_chain_create "${protocol}" -t mangle OUTPUT      --policy="ACCEPT"
	iptables_chain_create "${protocol}" -t mangle FORWARD     --policy="ACCEPT"
	iptables_chain_create "${protocol}" -t mangle POSTROUTING --policy="ACCEPT"

	# Create NAT table and initialize chains.
	iptables "${protocol}" -t nat "* nat"
	iptables_chain_create "${protocol}" -t nat PREROUTING  --policy="ACCEPT"
	iptables_chain_create "${protocol}" -t nat OUTPUT      --policy="ACCEPT"
	iptables_chain_create "${protocol}" -t nat POSTROUTING --policy="ACCEPT"
}

# Load the created ruleset into the kernel.
function iptables_commit () {
	local protocol="${1}"
	assert isset protocol
	shift

	local testmode="false"

	while [ $# -gt 0 ]; do
		case "${1}" in
			--test)
				testmode="true"
				;;
			*)
				log WARNING "Unrecognized argument: ${1}"
				;;
		esac
		shift
	done

	# Concat all rules into one big file.
	local rulesfile="${IPTABLES_TMPDIR}/ruleset"
	_iptables_commit_cat_rulesfile "${protocol}" "${rulesfile}"

	# Run the following loop twice:
	# 1st: Check if the ruleset can be loaded
	# 2nd: If not in test mode, actually load the ruleset into the kernel
	local load_cmd="--test"
	local ret=0

	local i
	for i in 0 1; do
		_iptables_commit_load_rulesfile "${protocol}" "${rulesfile}" "${load_cmd}"
		ret=$?

		case "${i},${ret}" in
			0,${EXIT_OK})
				iptables_dump "${protocol}" "${rulesfile}" --log-facility="DEBUG"
				log DEBUG "Ruleset load check succeeded (${protocol})"
				;;

			# Loading rules has failed (test)
			0,*)
				iptables_dump "${protocol}" "${rulesfile}" --log-facility="CRITICAL"
				log CRITICAL "Ruleset load check failed (${protocol} - ${ret})"
				return ${ret}
				;;

			1,${EXIT_OK})
				log DEBUG "Ruleset successfully loaded (${protocol})"
				return ${EXIT_OK}
				;;

			1,*)
				log CRITICAL "Ruleset loading failed (${protocol})"
				return ${ret}
				;;
		esac

		# Skip the second loop iteration, if we are running in test mode.
		enabled testmode && break

		load_cmd=""
	done

	return ${EXIT_OK}
}

function _iptables_commit_cat_rulesfile() {
	local protocol="${1}"
	assert isset protocol

	local rulesfile="${2}"
	assert isset rulesfile

	local table
	local file
	for table in ${IPTABLES_TABLES}; do
		file="${IPTABLES_TMPDIR}/${protocol}-${table}"

		fread "${file}"

		# Add the COMMIT statement for every table.
		print "COMMIT"
	done > "${rulesfile}"

	assert [ -s "${rulesfile}" ]
}

function _iptables_commit_load_rulesfile() {
	local protocol="${1}"
	assert isset protocol

	local rulesfile="${2}"
	assert isset rulesfile
	shift 2

	local testmode="false"
	while [ $# -gt 0 ]; do
		case "${1}" in
			--test)
				testmode="true"
				;;
		esac
		shift
	done

	local iptables_cmd
	case "${protocol}" in
		ipv6)
			iptables_cmd="ip6tables-restore"
			;;
		ipv4)
			iptables_cmd="iptables-restore"
			;;
	esac
	assert isset iptables_cmd

	if enabled testmode; then
		list_append iptables_cmd "--test"
	fi

	# Save when importing the rules has started.
	local time_started="$(timestamp)"

	cmd "${iptables_cmd}" < "${rulesfile}"
	local ret=$?

	case "${ret}" in
		${EXIT_OK})
			local time_finished="$(timestamp)"
			time_finished="$(( ${time_finished} - ${time_started} ))"

			enabled testmode && return ${EXIT_OK}

			log INFO "Successfully loaded new firewall ruleset for ${protocol} in ${time_finished}s!"
			;;
		*)
			if ! enabled testmode; then
				log CRITICAL "Error loading firewall ruleset for ${protocol}!"
			fi
			;;
	esac

	return ${ret}
}

function iptables_dump() {
	local protocol="${1}"
	assert isset protocol

	local rulesfile="${2}"
	assert isset rulesfile
	shift 2

	local log_facility="INFO"

	while [ $# -gt 0 ]; do
		case "${1}" in
			--log-facility=*)
				log_facility="$(cli_get_val ${1})"
				;;
			*)
				log WARNING "Unrecognized argument: ${1}"
				;;
		esac
		shift
	done

	# Say what we are going to do:
	log "${log_facility}" "Firewall ruleset for ${protocol}:"

	local counter="0"
	local line
	while read -r line; do
		counter="$(( ${counter} + 1 ))"

		printf -v line "%4d | %s" "${counter}" "${line}"
		log "${log_facility}" "${line}"
	done < "${rulesfile}"
}

function iptables_LOG() {
	local prefix="${1}"
	local ret

	# Automatically append a colon and whitespace.
	case "${prefix}" in
		# Everything is fine.
		"*: ") ;;

		# Ends with colon, add whitespace only.
		"*:")
			prefix="${prefix} "
			;;

		# Append both.
		*)
			prefix="${prefix}: "
			;;
	esac

	case "${FIREWALL_LOG_METHOD}" in
		nflog)
			ret="NFLOG --nflog-threshold ${FIREWALL_NFLOG_THRESHOLD}"
			isset prefix && ret="${ret} --nflog-prefix \"$prefix\""
			;;
		syslog)
			ret="LOG"
			isset prefix && ret="${ret} --log-prefix \"$prefix\""
			;;
	esac

	print "${ret}"
}

function iptables_protocol() {
	local PROTO
	PROTO=$1
	for proto in tcp udp esp ah; do
		if [ "$PROTO" = "$proto" ]; then
			echo "-p $PROTO"
			break
		fi
	done
}

IPTABLES_PORT=0
IPTABLES_MULTIPORT=1
IPTABLES_PORTRANGE=2

function _iptables_port_range() {
	grep -q ":" <<< $@
}

function _iptables_port_multiport() {
	grep -q "," <<< $@
}

function _iptables_port() {
	if _iptables_port_range "$@"; then
		echo $IPTABLES_PORTRANGE
	elif _iptables_port_multiport "$@"; then
		echo $IPTABLES_MULTIPORT
	else
		echo $IPTABLES_PORT
	fi
}

function iptables_source_port() {
	[ -z "$@" ] && return
	local type
	type=$(_iptables_port $@)
	if [ "$type" = "$IPTABLES_MULTIPORT" ]; then
		echo "-m multiport --source-ports $@"
	else
		echo "--sport $@"
	fi
}

function iptables_destination_port() {
	[ -z "$@" ] && return
	local type
	type=$(_iptables_port $@)
	if [ "$type" = "$IPTABLES_MULTIPORT" ]; then
		echo "-m multiport --destination-ports $@"
	else
		echo "--dport $@"
	fi
}