#!/bin/bash
###############################################################################
#                                                                             #
# IPFire.org - A linux based firewall                                         #
# Copyright (C) 2012  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/>.       #
#                                                                             #
###############################################################################

function iptables() {
	local arg
	local args
	local table="filter"
	local src dst

	# Default is both protocols.
	local proto="6 4"

	# Check if the directory where we put our rules in is set and
	# exists.
	assert isset IPTABLES_TMPDIR
	assert [ -d "${IPTABLES_TMPDIR}" ]

	# Parsing arguments
	while [ $# -gt 0 ]; do
		case "${1}" in
			# Select IPv4 protocol.
			-4)
				proto="4"
				shift
				;;
			# Select IPv6 protocol.
			-6)
				proto="6"
				shift
				;;
			-t)
				table=${2}
				shift 2
				;;
			-A)
				args="${args} -A ${2^^}"
				shift 2
				;;
			*)
				args="${args} ${1}"

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

	# Check that the nat table is not used for IPv6.
	if isoneof 6 ${proto}; then
		assert [ "${table}" != "nat" ]
	fi

	# Detect the version of the IP protocol.
	local src_proto
	isset src && src_proto=$(ip_detect_protocol ${src})

	local dst_proto
	isset dst && dst_proto=$(ip_detect_protocol ${dst})

	# Check that the source and destinations are not
	# using different versions of the IP protocol.
	if isset src_proto && isset dst_proto; then
		assert [ "${src_proto}" = "${dst_proto}" ]
	fi

	local rulesfile
	local p
	for p in ${proto}; do
		case "${p}" in
			6)
				listmatch ipv4 ${src_proto} ${dst_proto} \
					&& continue
				;;
			4)
				listmatch ipv6 ${src_proto} ${dst_proto} \
					&& continue
				;;
		esac

		rulesfile=$(iptables_rulesfile ipv${p} ${table})
		assert isset rulesfile

		print "${args:1:${#args}}" >> ${rulesfile}
	done
}

# Calls the binary iptables command.
function _iptables() {
	local iptables_cmd=$(which iptables)
	assert isset iptables_cmd

	cmd ${iptables_cmd} $@
}

function iptables_status() {
	_iptables -L -n -v
}

# Returns which tables exist for the given protocol.
function iptables_tables() {
	local proto=${1}
	assert isset proto

	case "${proto}" in
		ipv6)
			print "filter mangle"
			;;
		ipv4)
			print "filter mangle nat"
			;;
		*)
			return ${EXIT_ERROR}
			;;
	esac

	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 policy=${1}
	assert isset policy

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

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

	# Add NAT table for IPv4.
	iptables -4 -t nat "* nat"
	iptables_chain_create -4 -t nat PREROUTING  --policy=ACCEPT
	iptables_chain_create -4 -t nat OUTPUT      --policy=ACCEPT
	iptables_chain_create -4 -t nat POSTROUTING --policy=ACCEPT
}

# Load the created ruleset into the kernel.
function iptables_load() {
	# If first argument is present and true, we
	# run in test mode.
	local test="${1}"

	local rulesfile

	# Concat the table rulesets into one big file.
	local proto
	for proto in 6 4; do
		rulesfile=$(iptables_rulesfile ipv${proto})
		assert isset rulesfile

		local table
		local tablefile
		for table in $(iptables_tables ipv${proto}); do
			tablefile=$(iptables_rulesfile ipv${proto} ${table})

			fread ${tablefile}

			# Add the COMMIT statement for every table.
			if [ -s "${tablefile}" ]; then
				print "COMMIT"
			fi
		done > ${rulesfile}

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

	local error="false"
	local ret

	# First check if everything is correctly formatted.
	for proto in 6 4; do
		rulesfile=$(iptables_rulesfile ipv${proto})
		assert isset rulesfile

		_iptables_load ipv${proto} ${rulesfile} true
		if [ $? -ne ${EXIT_OK} ]; then
			log CRITICAL "Ruleset load check failed for IPv${proto}"
			error="true"
		fi
	done

	# Check if there has been an error in the load check.
	if enabled error; then
		iptables_dump CRITICAL

		log CRITICAL "New firewall rules could not be loaded."
		return ${EXIT_ERROR}
	fi

	# Dump the data, we are going to load.
	iptables_dump

	# If we are running in test mode, we are done here.
	enabled test && return ${EXIT_OK}

	# If we got until here, everything is fine to load the ruleset.
	for proto in 6 4; do
		rulesfile=$(iptables_rulesfile ipv${proto})

		_iptables_load ipv${proto} ${rulesfile}
	done
	return ${EXIT_OK}
}

function _iptables_load() {
	local proto=${1}
	local file=${2}
	local testmode=${3}

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

	if enabled testmode; then
		command="${command} --test"
	fi

	local time_started=$(date -u "+%s")

	cmd ${command} < ${file}
	local ret=$?

	case "${ret}" in
		${EXIT_OK})
			local time_finished=$(date -u "+%s")
			time_finished=$(( ${time_finished} - ${time_started} ))
			log INFO \
				"Successfully loaded new firewall ruleset for IPv${proto/ipv/} in ${time_finished}s!"
			;;
		*)
			if ! enabled testmode; then
				log CRITICAL "Error loading firewall ruleset for IPv${proto/ipv/}!"
			fi
			;;
	esac

	return ${ret}
}

function iptables_dump() {
	local log_facility=${1-DEBUG}

	# Here is nothing to do, if we are not running in
	# debug mode.
	enabled DEBUG || return ${EXIT_OK}

	local rulesfile
	local counter
	local line

	local proto
	for proto in 6 4; do
		rulesfile=$(iptables_rulesfile ipv${proto})
		[ -e "${rulesfile}" ] || continue

		log ${log_facility} "Firewall ruleset for IPv${proto}:"

		counter=1
		while read -r line; do
			printf -v line "%4d | %s" "${counter}" "${line}"
			log ${log_facility} "${line}"

			counter=$(( $counter + 1 ))
		done < ${rulesfile}
	done
}

function iptables_chain_create() {
	local chain
	local table="filter"
	local policy="-"
	local proto
	local args
	while [ $# -gt 0 ]; do
		case "${1}" in
			-6|-4)
				proto=${1}
				;;
			-t)
				table=${2}
				shift
				;;
			--policy=*)
				policy=$(cli_get_val ${1})
				;;
			*)
				chain=${1}
				;;
		esac
		shift
	done

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

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

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

	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
}