#!/bin/bash
set -euo pipefail

## INIT ##
# get script base directory
SCRIPT_PATH=$(readlink -f "${BASH_SOURCE[0]}")
SCRIPT_DIR=$(dirname -- "$SCRIPT_PATH")

# store arguments globally
ARG_LIST=("$@")

# get version number
VERSION=$(cat ${SCRIPT_DIR}/VERSION | xargs)

# load libraries
source "${SCRIPT_DIR}/lib/color"
source "${SCRIPT_DIR}/lib/log"
source "${SCRIPT_DIR}/lib/toolcheck"
source "${SCRIPT_DIR}/lib/sudocheck"

# regex definitions
regex_ip_address="^([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\.([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])){3}$"
regex_ip_or_net="^([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\.([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])){3}(\/([0-9]|[1-2][0-9]|3[0-2]))?$"
regex_local_address="^(10\.|172\.(1[6-9]|2[0-9]|3[0-1])\.|192\.168\.).*$"

# subnet list
subnet_list_special_use="0.0.0.0/8 100.64.0.0/10 127.0.0.0/8 169.254.0.0/16 192.0.0.0/24 192.0.0.0/24 192.88.99.0/24 198.18.0.0/15 198.51.100.0/24 203.0.113.0/24 224.0.0.0/4 233.252.0.0/24 240.0.0.0/4 255.255.255.255/32"
subnet_list_local="10.0.0.0/8 172.16.0.0/12 192.168.0.0/16"



#
# HELPER: Remove a firewall.
#
xynat_fw_remove(){
	## REMOVE CHAINS ##
	# flush chains
	iptables -F "${chain_id}_IN" 2> /dev/null || true
	iptables -F "${chain_id}_OUT" 2> /dev/null || true
	iptables -F "${chain_id}_FWI" 2> /dev/null || true
	iptables -F "${chain_id}_FWO" 2> /dev/null || true
	iptables -t nat -F "${chain_id}_NI" 2> /dev/null || true
	iptables -t nat -F "${chain_id}_NO" 2> /dev/null || true
	
	# remove ipv4 references
	iptables -D INPUT -i "$arg_iface" -j "${chain_id}_IN" 2> /dev/null || true
	iptables -D OUTPUT -o "$arg_iface" -j "${chain_id}_OUT" 2> /dev/null || true
	iptables -D FORWARD -i "$arg_iface" -j "${chain_id}_FWI" 2> /dev/null || true
	iptables -D FORWARD -o "$arg_iface" -j "${chain_id}_FWO" 2> /dev/null || true
	iptables -t nat -D PREROUTING -j "${chain_id}_NI" 2> /dev/null || true
	iptables -t nat -D POSTROUTING -j "${chain_id}_NO" 2> /dev/null || true
	
	# remove ipv6 references
	ip6tables -D INPUT -i "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	ip6tables -D OUTPUT -o "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	ip6tables -D FORWARD -i "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	ip6tables -D FORWARD -o "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	
	# delete chains
	iptables -X "${chain_id}_IN" 2> /dev/null || true
	iptables -X "${chain_id}_OUT" 2> /dev/null || true
	iptables -X "${chain_id}_FWI" 2> /dev/null || true
	iptables -X "${chain_id}_FWO" 2> /dev/null || true
	iptables -t nat -X "${chain_id}_NI" 2> /dev/null || true
	iptables -t nat -X "${chain_id}_NO" 2> /dev/null || true
}



#
# HELPER: Add a firewall.
#
xynat_fw_add(){
	## RESET ##
	# remove references
	iptables -D INPUT -i "$arg_iface" -j "${chain_id}_IN" 2> /dev/null || true
	iptables -D OUTPUT -o "$arg_iface" -j "${chain_id}_OUT" 2> /dev/null || true
	iptables -D FORWARD -i "$arg_iface" -j "${chain_id}_FWI" 2> /dev/null || true
	iptables -D FORWARD -o "$arg_iface" -j "${chain_id}_FWO" 2> /dev/null || true
	iptables -t nat -D PREROUTING -j "${chain_id}_NI" 2> /dev/null || true
	iptables -t nat -D POSTROUTING -j "${chain_id}_NO" 2> /dev/null || true
	
	# remove ipv6 references
	ip6tables -D INPUT -i "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	ip6tables -D OUTPUT -o "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	ip6tables -D FORWARD -i "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	ip6tables -D FORWARD -o "$arg_iface" -j REJECT --reject-with icmp6-no-route 2> /dev/null || true
	
	
	## ADD CHAINS ##
	# create chains
	iptables -N "${chain_id}_IN" 2> /dev/null || true
	iptables -N "${chain_id}_OUT" 2> /dev/null || true
	iptables -N "${chain_id}_FWI" 2> /dev/null || true
	iptables -N "${chain_id}_FWO" 2> /dev/null || true
	iptables -t nat -N "${chain_id}_NI" 2> /dev/null || true
	iptables -t nat -N "${chain_id}_NO" 2> /dev/null || true
	
	# insert references
	iptables -I INPUT -i "$arg_iface" -j "${chain_id}_IN"
	iptables -I OUTPUT -o "$arg_iface" -j "${chain_id}_OUT"
	iptables -I FORWARD -i "$arg_iface" -j "${chain_id}_FWI"
	iptables -I FORWARD -o "$arg_iface" -j "${chain_id}_FWO"
	iptables -t nat -I PREROUTING -j "${chain_id}_NI"
	iptables -t nat -I POSTROUTING -j "${chain_id}_NO"
	
	# block ipv6
	ip6tables -I INPUT -i "$arg_iface" -j REJECT --reject-with icmp6-no-route
	ip6tables -I OUTPUT -o "$arg_iface" -j REJECT --reject-with icmp6-no-route
	ip6tables -I FORWARD -i "$arg_iface" -j REJECT --reject-with icmp6-no-route
	ip6tables -I FORWARD -o "$arg_iface" -j REJECT --reject-with icmp6-no-route
}



#
# HELPER: Update all rulesets.
#
xynat_ruleset_update(){
	xynat_ruleset_update_in
	xynat_ruleset_update_out
	xynat_ruleset_update_fwi
	xynat_ruleset_update_fwo
	xynat_ruleset_update_ni
	xynat_ruleset_update_no
}



#
# HELPER: Update ruleset: `in`.
#
xynat_ruleset_update_in(){
	## FLUSH CHAIN ##
	iptables -F "${chain_id}_IN"
	
	
	## ADD RULES ##
	# enforce correct vm address
	iptables -A "${chain_id}_IN" ! -s "$arg_vm_address" -j REJECT --reject-with icmp-admin-prohibited
	
	# maybe allow host access
	if [[ "${arg_allow_host_mode:-}" == "in_out" ]]; then
		iptables -A "${chain_id}_IN" -j ACCEPT
	elif [[ "${arg_allow_host_mode:-}" == "in" ]]; then
		iptables -A "${chain_id}_IN" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
	fi
	
	# reject all packets
	iptables -A "${chain_id}_IN" -j REJECT --reject-with icmp-host-unreachable
}



#
# HELPER: Update ruleset: `out`.
#
xynat_ruleset_update_out(){
	## FLUSH CHAIN ##
	iptables -F "${chain_id}_OUT"
	
	
	## ADD RULES ##
	# allow related icmp messages
	iptables -A "${chain_id}_OUT" -p icmp -m conntrack --ctstate RELATED -j ACCEPT
	
	# enforce correct vm address
	iptables -A "${chain_id}_OUT" ! -d "$arg_vm_address" -j REJECT --reject-with icmp-admin-prohibited
	
	# maybe allow host access
	if [[ "${arg_allow_host_mode:-}" == "in_out" || "${arg_allow_host_mode:-}" == "in" ]]; then
		iptables -A "${chain_id}_OUT" -j ACCEPT
	fi
	
	# reject all other packets
	iptables -A "${chain_id}_OUT" -j REJECT --reject-with icmp-host-unreachable
}



#
# HELPER: Update ruleset: `fwi`.
#
xynat_ruleset_update_fwi(){
	## FLUSH CHAIN ##
	iptables -F "${chain_id}_FWI"
	
	
	## ADD RULES ##
	# enforce correct vm address
	iptables -A "${chain_id}_FWI" ! -s "$arg_vm_address" -j REJECT --reject-with icmp-admin-prohibited
	
	# reject packets for blocked address(es)
	for a in ${arg_block[*]:-}; do
		iptables -A "${chain_id}_FWI" -d "$a" -j REJECT --reject-with icmp-net-unreachable
	done
	
	# accept allowed local addresses
	for a in ${arg_allow[*]:-}; do
		iptables -A "${chain_id}_FWI" -d "$a" -j ACCEPT
	done
	
	# accept allowed incomming local addresses
	for a in ${arg_allow_in[*]:-}; do
		iptables -A "${chain_id}_FWI" -d "$a" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
	done
	
	# reject filtered packets
	for a in $subnet_list_special_use $subnet_list_local; do
		iptables -A "${chain_id}_FWI" -d "$a" -j REJECT --reject-with icmp-net-unreachable
	done
	
	# default: allow
	iptables -A "${chain_id}_FWI" -j ACCEPT
}



#
# HELPER: Update ruleset: `fwo`.
#
xynat_ruleset_update_fwo(){
	## FLUSH CHAIN ##
	iptables -F "${chain_id}_FWO"
	
	
	## ADD RULES ##
	# enforce correct vm address
	iptables -A "${chain_id}_FWO" ! -d "$arg_vm_address" -j REJECT --reject-with icmp-admin-prohibited
	
	# reject packets for blocked address(es)
	for a in ${arg_block[*]:-}; do
		iptables -A "${chain_id}_FWO" -s "$a" -j REJECT --reject-with icmp-net-unreachable
	done
	
	# accept allowed local addresses
	for a in ${arg_allow[*]:-} ${arg_allow_in[*]:-}; do
		iptables -A "${chain_id}_FWO" -s "$a" -j ACCEPT
	done
	
	# reject filtered packets
	for a in $subnet_list_special_use $subnet_list_local; do
		iptables -A "${chain_id}_FWO" -s "$a" -j REJECT --reject-with icmp-net-unreachable
	done
	
	# default: accept
	iptables -A "${chain_id}_FWO" -j ACCEPT
}



#
# HELPER: Update ruleset: `ni`.
#
xynat_ruleset_update_ni(){
	## FLUSH CHAIN ##
	iptables -t nat -F "${chain_id}_NI"
	
	
	## ADD RULES ##
	# translate destination address (forward to vm)
	if [[ "${arg_public_ip:+x}" ]]; then
		iptables -t nat -A "${chain_id}_NI" -d "$arg_public_ip" -j DNAT --to-destination "$arg_vm_address"
	fi
}



#
# HELPER: Update ruleset: `no`.
#
xynat_ruleset_update_no(){
	## FLUSH CHAIN ##
	iptables -t nat -F "${chain_id}_NO"
	
	
	## ADD RULES ##
	# filter for vm ip address
	iptables -t nat -A "${chain_id}_NO" ! -s "$arg_vm_address" -j RETURN
	
	# translate source address/port (forward to internet)
	if [[ "${arg_public_ip:+x}" ]]; then
		iptables -t nat -A "${chain_id}_NO" -j SNAT --to-source "$arg_public_ip"
	else
		iptables -t nat -A "${chain_id}_NO" -j MASQUERADE
	fi
}



#
# Explain usage of script.
#
xynat_help(){
	echo "XYNat v${VERSION}"
	echo
	echo "Usage: $0 [OPTIONS]"
	echo
	echo "Options:"
	echo "    -a, --add                  - Add new firewall                                        (default, fallback: MODE=start)"
	echo "    -r, --remove               - Remove existing firewall                                (fallback: MODE=stop)"
	echo
	echo "    -i, --iface=iface          - Interface name for virtual switch                       (required, fallback: IFACE)"
	echo "    -s, --vm-address=ip        - IP address of virtual machine                           (required when adding new fw)"
	echo "    -p, --public-ip=ip         - IP address to use for outgoing traffic and DNAT"
	echo
	echo "    -b, --block=ip-or-net      - Block address(es) for all connections                   (multi-use allowed)"
	echo
	echo "    -w, --allow=ip-or-net      - Allow local address(es) for all connections             (multi-use allowed)"
	echo "    -x, --allow-in=ip-or-net   - Allow local address(es) for incomming connections only  (multi-use allowed)"
	echo "    -y, --allow-host           - Allow local host for all connections"
	echo "    -z, --allow-host-in        - Allow local host for incomming connections only"
	echo
	echo "    -h, --help                 - Display this help message and exit"
	echo "    -v, --version              - Display version information and exit"
	echo
	echo "Examples:"
	echo "    Add a new firewall and NAT for a VM which has IP address '192.168.234.2' and is a member of the bridge interface 'br2'"
	echo "    All traffic to and from local addresses will be blocked; Traffic to public addresses will be allowed"
	echo "        $0 --add --iface \"br2\" --vm-address \"192.168.234.2\""
	echo
	echo "    Same as first example, but use '192.168.123.99' as source address for outgoing connections"
	echo "        $0 --add --iface \"br2\" --vm-address \"192.168.234.2\" --public-ip \"192.168.123.99\""
	echo
	echo "    Same as first example, but also block all traffic to and from '233.252.0.0/24'"
	echo "        $0 --add --iface \"br2\" --vm-address \"192.168.234.2\" --block \"233.252.0.0/24\""
	echo
	echo "    Same as second example, but allow incomming and outgoing connections from host device"
	echo "    The VM will also be accessible from devices in '192.168.137.64/30' via its public address"
	echo "        $0 --add --iface \"br2\" --vm-address \"192.168.234.2\" --public-ip \"192.168.123.99\" --allow-host --allow-in \"192.168.137.64/30\""
	echo
	echo "    Remove all firewall and NAT rules for iface 'br2'"
	echo "        $0 --remove --iface \"br2\""
}



#
# Display version information.
#
xynat_version(){
	echo "XYNat v${VERSION}"
	echo "(c) 2024 DrMaxNix"
}



#
# VALIDATOR: `vm-address`.
#
xynat_validate_vm_address(){
	## CHECK FORMAT ##
	if [[ ! "$1" =~ $regex_ip_address ]]; then
		log_error "Malformed vm-address '$1'"
	fi
}



#
# VALIDATOR: `public-ip`.
#
xynat_validate_public_ip(){
	## CHECK FORMAT ##
	if [[ ! "$1" =~ $regex_ip_address ]]; then
		log_error "Malformed vm-address '$1'"
	fi
	
	
	## CHECK IP IFACE ASSIGNMENT ##
	# make sure this is not a primary ip address
	local primary_ip_list="$(xynat_primary_ip_list)"
	for p in $primary_ip_list; do
		if [[ "$p" == "$1" ]]; then
			log_error "Virtual machine's public-ip is in use as primary host address"
		fi
	done
	
	# check whether it exists as a secondary address
	local secondary_ip_list="$(xynat_secondary_ip_list)"
	local found="no"
	for s in $secondary_ip_list; do
		if [[ "$s" == "$1" ]]; then
			found="yes"
		fi
	done
	if [[ "$found" != "yes" ]]; then
		log_warn "Virtual machine's public-ip '$1' is not assigned to any host interface"
	fi
}



#
# VALIDATOR: `block`.
#
xynat_validate_block(){
	for a in $1; do
		## VALIDATE SYNTAX ##
		if [[ ! "$a" =~ $regex_ip_or_net ]]; then
			log_error "Malformed ip address or subnet in blocklist: '$a'"
		fi
	done
}



#
# VALIDATOR: `allow`.
#
xynat_validate_allow(){
	for a in $1; do
		## VALIDATE SYNTAX ##
		if [[ ! "$a" =~ $regex_ip_or_net ]]; then
			log_error "Malformed ip address or subnet in allowlist: '$a'"
		fi
		
		
		## CHECK FOR KNOWN LOCAL PREFIX ##
		if [[ ! "$a" =~ $regex_local_address ]]; then
			log_warn "Allowlist entry does not look like a local address: '$a'"
		fi
	done
}



#
# HELPER: List all primary ip addresses of the host.
#
xynat_primary_ip_list(){
	ip address show | grep -Pe "^\s+inet ([0-9\.]+)\/[0-9]+.*$" | grep -v "\ssecondary\s" | sed -E "s/^\s+inet ([0-9\.]+)\/[0-9]+.*$/\1/g"
}



#
# HELPER: List all secondary ip addresses of the host.
#
xynat_secondary_ip_list(){
	ip address show | grep -Pe "^\s+inet ([0-9\.]+)\/[0-9]+.*$" | grep "\ssecondary\s" | sed -E "s/^\s+inet ([0-9\.]+)\/[0-9]+.*$/\1/g"
}



## TOOLCHECK ##
toolcheck error "grep sed iptables ip"


## PARSE ARGUMENTS ##
while [[ "$#" -gt 0 ]]; do
	case "$1" in
		# help
		-h|--help)
			xynat_help
			exit 0
		;;
		
		# version
		-v|--version)
			xynat_version
			exit 0
		;;
		
		# add
		-a|--add)
			if [[ "${arg_mode:+x}" ]]; then log_error "Cannot set mode twice; already set to '$arg_mode'"; fi
			arg_mode="add"
			shift
		;;
		
		# remove
		-r|--remove)
			if [[ "${arg_mode:+x}" ]]; then log_error "Cannot set mode twice; already set to '$arg_mode'"; fi
			arg_mode="remove"
			shift
		;;
		
		# iface
		-i|--iface|-i=*|--iface=*)
			if [[ "${arg_iface:+x}" ]]; then log_error "Cannot set iface twice; already set to '$arg_iface'"; fi
			if [[ "$1" =~ ^[a-z\-]+=(.+)$ ]]; then
				arg_iface=$(echo $1 | sed -E "s/^[a-z\-]+=(.*)$/\1/g")
			shift; continue; fi
			
			shift; if [[ $# -gt 0 ]]; then
				arg_iface="$1"
			else log_error "Value expected for parameter 'iface'"; fi; shift
		;;
		
		# vm-address
		-s|--vm-address|-s=*|--vm-address=*)
			if [[ "${arg_vm_address:+x}" ]]; then log_error "Cannot set vm-address twice; already set to '$arg_vm_address'"; fi
			if [[ "$1" =~ ^[a-z\-]+=(.+)$ ]]; then
				arg_vm_address=$(echo $1 | sed -E "s/^[a-z\-]+=(.*)$/\1/g")
			shift; continue; fi
			
			shift; if [[ $# -gt 0 ]]; then
				arg_vm_address="$1"
			else log_error "Value expected for parameter 'vm-address'"; fi; shift
		;;
		
		# public-ip
		-p|--public-ip|-p=*|--public-ip=*)
			if [[ "${arg_public_ip:+x}" ]]; then log_error "Cannot set public-ip twice; already set to '$arg_public_ip'"; fi
			if [[ "$1" =~ ^[a-z\-]+=(.+)$ ]]; then
				arg_public_ip=$(echo $1 | sed -E "s/^[a-z\-]+=(.*)$/\1/g")
			shift; continue; fi
			
			shift; if [[ $# -gt 0 ]]; then
				arg_public_ip="$1"
			else log_error "Value expected for parameter 'public-ip'"; fi; shift
		;;
		
		# block
		-b|--block|-b=*|--block=*)
			if [[ "$1" =~ ^[a-z\-]+=(.+)$ ]]; then
				arg_block=(${arg_block[@]:-""} "$(echo $1 | sed -E "s/^[a-z\-]+=(.*)$/\1/g")")
			shift; continue; fi
			
			shift; if [[ $# -gt 0 ]]; then
				arg_block=(${arg_block[@]:-""} "$1")
			else log_error "Value expected for parameter 'block'"; fi; shift
		;;
		
		# allow
		-w|--allow|-w=*|--allow=*)
			if [[ "$1" =~ ^[a-z\-]+=(.+)$ ]]; then
				arg_allow=(${arg_allow[@]:-""} "$(echo $1 | sed -E "s/^[a-z\-]+=(.*)$/\1/g")")
			shift; continue; fi
			
			shift; if [[ $# -gt 0 ]]; then
				arg_allow=(${arg_allow[@]:-""} "$1")
			else log_error "Value expected for parameter 'allow'"; fi; shift
		;;
		
		# allow-in
		-x|--allow-in|-x=*|--allow-in=*)
			if [[ "$1" =~ ^[a-z\-]+=(.+)$ ]]; then
				arg_allow_in=(${arg_allow_in[@]:-""} "$(echo $1 | sed -E "s/^[a-z\-]+=(.*)$/\1/g")")
			shift; continue; fi
			
			shift; if [[ $# -gt 0 ]]; then
				arg_allow_in=(${arg_allow_in[@]:-""} "$1")
			else log_error "Value expected for parameter 'allow-in'"; fi; shift
		;;
		
		# allow-host
		-y|--allow-host|-y=*|--allow-host=*)
			if [[ "${arg_allow_host_mode:+x}" ]]; then log_error "Cannot set allow-host mode twice; already set to '$arg_allow_host_mode'"; fi
			arg_allow_host_mode="in_out"
			shift
		;;
		
		# allow-host-in
		-z|--allow-host-in|-z=*|--allow-host-in=*)
			if [[ "${arg_allow_host_mode:+x}" ]]; then log_error "Cannot set allow-host mode twice; already set to '$arg_allow_host_mode'"; fi
			arg_allow_host_mode="in"
			shift
		;;
		
		# unknown
		*)
			log_error "Unexpected argument '$1'"
		;;
	esac
done


## MAKE SURE WE HAVE SUDO PERMISSION ##
sudocheck


## ARGUMENT FALLBACKS AND VALIDATION ##
# mode
if [[ -z "${arg_mode:+x}" ]]; then
	case "${MODE:-""}" in
		"stop") arg_mode="remove";;
		"start"|*) arg_mode="add";;
	esac
fi

# iface
if [[ -z "${arg_iface:+x}" ]]; then
	if [[ -z "${IFACE:+x}" ]]; then log_error "Missing required argument 'iface'; See '$0 --help' for usage information"; fi
	arg_iface="$IFACE"
fi

# vm-address
if [[ "$arg_mode" != "remove" ]]; then
	if [[ -z "${arg_vm_address:+x}" ]]; then
		log_error "Missing required argument 'vm-address'; See '$0 --help' for usage information"
	fi
	xynat_validate_vm_address "$arg_vm_address"
fi

# public-ip
if [[ "${arg_public_ip:+x}" ]]; then
	xynat_validate_public_ip "$arg_public_ip"
fi

# block
xynat_validate_block "${arg_block[*]:-""}"

# allow
xynat_validate_allow "${arg_allow[*]:-""}"

# allow-in
xynat_validate_allow "${arg_allow_in[*]:-""}"


## BUILD MISSING VALUES ##
# firewall chain id
chain_id="XYNAT_$(echo "$arg_iface" | tr a-z A-Z | tr - _)"


## EXECUTE ACTION ##
case "$arg_mode" in
	# add
	"add")
		xynat_fw_add
		xynat_ruleset_update
	;;
	
	# remove
	"remove")
		xynat_fw_remove
	;;
	
	# should not happen
	*)
		log_error "Invalid mode '$arg_mode'"
	;;
esac