#!/bin/bash
#config
if [ "$EUID" -ne 0 ]
  then echo "Please run as root"
  exit
fi
CFG="/etc/firewall.cfg"

if [[ ! -f "${CFG}" ]]; then
	for i in $(ip -oneline addr show scope global|awk '{print $2}'|uniq); do
		devs+=($i)
	done
	echo "alldevs=(${devs[*]})" > "${CFG}"
	echo "ignore=()" >> "${CFG}"
	echo "home_tcp=(ssh)" >> "${CFG}"
	echo "home_udp=()" >> "${CFG}"
	echo "out_tcp=(ssh)" >> "${CFG}"
	echo "out_udp=()" >> "${CFG}"
	echo "vpn_tcp=(ssh)" >> "${CFG}"
	echo "vpn_udp=()" >> "${CFG}"
	echo "" >> "${CFG}"
	echo "PROFILE=" >> "${CFG}"
	echo "" >> "${CFG}"
	echo "NM=" >> "${CFG}"
	echo "" >> "${CFG}"
	echo "VPN_UNRESTRICTED=1" >> "${CFG}"
	echo "" >> "${CFG}"
	echo "nft="/usr/sbin/nft"" >> "${CFG}"
	echo "" >> "${CFG}"
	echo "declare -A ports=(" >> "${CFG}"
	echo "  [ssh]=22" >> "${CFG}"
	echo "  [smtp]=25" >> "${CFG}"
	echo "  [dns]=53" >> "${CFG}"
	echo "  [netbios-ssn]=139" >> "${CFG}"
	echo "  [dovecot]=143" >> "${CFG}"
	echo "  [microsoft-ds]=445" >> "${CFG}"
	echo "  [cups-ipp]=631" >> "${CFG}"
	echo "  [imaps]=993" >> "${CFG}"
	echo "  [kdeconnect]=1714-1764" >> "${CFG}"
	echo "  [avahi-daemon]=5353" >> "${CFG}"
	echo "  [vncd]=5900" >> "${CFG}"
	echo "  [mpd]=6600" >> "${CFG}"
	echo "  [mpd-audio]=8000" >> "${CFG}"
	echo "  [misc]=8080" >> "${CFG}"
	echo "  [syncthing]=8384" >> "${CFG}"
	echo "  [wireguard]=51820" >> "${CFG}"
	echo "  [mosh]=60000-61000" >> "${CFG}"
	echo ")" >> "${CFG}"
	echo "${CFG} missed. Example created. Exit"
	exit
else
	source "${CFG}"
fi



##script

#source data at the end of file
source <(sed -n '/^#functions/,$p' $(dirname "$0")/$(basename "$0"))

#for index in "${!ports[@]}"; do 
#	echo "$index"; 
#	echo "${ports[$index]}"
#done

while getopts Aia:r:p:fDhnL option
do
	case "${option}" in
		A) AUTOMATIC=1;;
		i) INIT=1;;
		a) ADEVICE+=("${OPTARG}");;
		r) RDEVICE+=("${OPTARG}");;
		p) PROFILE="${OPTARG}";;
		f) FLUSH=1;;
		D) DEBUG=1;;
		h) HELP=1;;
		n) NM=1;;
		L) LIST=1;;
		*) HELP=1;;
	esac
done

help
[[ $LIST == "1" ]] && sudo nft -a list ruleset && exit
flush
check_deviceinput
get_devices
profile
init
add_device
remove_device
set_rules
exit 0




#functions
debug() {
	if [[ $DEBUG == "1" ]]; then
		echo "$@"
	fi
}

die() {
	debug "$@"
	exit 1
}

flush() {
	[[ ! $FLUSH == "1" ]] && return 0
	debug "Flush ruleset"
	sudo "${nft}" flush ruleset
	if [[ $FLUSH == "1" ]]; then
		exit
	fi
}

profile() {
	[[ $INIT == "1" ]] && return 0
	if [[ -z $PROFILE ]]; then
		if [[ -e /run/home ]]; then
			HOME=$(cat /run/home)
			if [[ "${HOME}" == "home" ]]; then
				PROFILE=home
			elif [[ "${HOME}" == "out" ]]; then
				PROFILE=out
			fi
		else
			PROFILE=out
		fi
	fi
	case $PROFILE in
		h|home)
		portstcp=("${home_tcp[@]}")
		portsudp=("${home_udp[@]}");;
		o|out)
		portstcp=("${out_tcp[@]}")
		portsudp=("${out_udp[@]}") ;;
	esac
	debug profile: "${PROFILE[@]}"
	debug tcp ports: "${portstcp[@]}"
	debug udp ports: "${portsudp[@]}"
	if [[ ! "$VPN_UNRESTRICTED" == "1" ]]; then
		debug vpntcp ports: "${vpn_tcp[@]}"
		debug vpnudp ports: "${vpn_udp[@]}"
	elif [[ "$VPN_UNRESTRICTED" == "1" ]]; then
		debug "vpn ports unrestricted"
	fi
}

check_deviceinput() {
	[[ $INIT == "1" ]] && return 0
	MISSING=1
	if [[ -n "$AUTOMATIC" ]]; then
		MISSING=0
	elif [[ -n "${ADEVICE[*]}" ]]; then
		MISSING=0
	elif [[ -n "${RDEVICE[*]}" ]]; then
		MISSING=0
	fi
	if [[ $MISSING == "1" ]]; then
		die "Whether automatic (-A) nor any device to add or remove (-a/-r) given."
	fi

	if [[ -n "$AUTOMATIC" ]] && [[ -n "${ADEVICE[*]}" ]]; then
		MISSING=1
	elif [[ -n "$AUTOMATIC" ]] && [[ -n "${RDEVICE[*]}" ]]; then
		MISSING=1
	fi
	if [[ $MISSING == "1" ]]; then
		die "Automatic (-A) and (-a/-r) can not be used together"
	fi	


	if [[ -n "${ADEVICE[*]}" ]] && [[ -n "${RDEVICE[*]}" ]]; then
		for i in "${ADEVICE[@]}"; do
			for z in "${RDEVICE[@]}"; do
				if [[ "$z" == "$i" ]]; then
					MISSING=1
					REASON+=("$i")
				fi
			done
		done
	fi
	if [[ $MISSING == "1" ]]; then
		die "Device can not be removed and added:" "${REASON[*]}"
	fi

	for i in "${ADEVICE[@]}" "${RDEVICE[@]}"; do
		FOUND=$(grep "$i" <<<"${alldevs[*]}")
		if [ ! "${FOUND}" != "" ]; then
		  MISSING=1
		  REASON+=("$i")
		fi
	done
	if [[ $MISSING == "1" ]]; then
		die "Device not found in cfg:" "${REASON[*]}"
	fi
}

get_devices() {
	[[ $INIT == "1" ]] && return 0
	if [[ -n $AUTOMATIC ]]; then
		if [[ $NM == "1" ]]; then
			for i in $(LANG=en_US.UTF-8 nmcli device status|grep " connected"|awk '{print $1}'); do
				for item2 in "${ignore[@]}"; do
					if [[ $i == "$item2" ]]; then
						continue 2
					fi
				done
				DEVON+=("$i")
			done
		else
			for i in $(ip -oneline addr show scope global|awk '{print $2}'|uniq); do
				for item2 in "${ignore[@]}"; do
					if [[ $i == "$item2" ]]; then
						continue 2
					fi
				done
				DEVON+=("$i")
			done
		fi
		DEVOFF=()
		for item1 in "${alldevs[@]}"; do
		    for item2 in "${DEVON[@]}"; do
		        if [[ $item1 == "$item2" ]]; then
		        	continue 2
		        fi
		    done
		    DEVOFF+=( "$item1" )
		done
	fi

	if [[ -z $AUTOMATIC ]]; then
		for i in "${ADEVICE[@]}"; do
			DEVON+=("$i")
		done
		for i in "${RDEVICE[@]}"; do
			DEVOFF+=("$i")
		done
	fi
	for i in "${ignore[@]}"; do
		DEVOFF+=("$i")
	done
	debug alldevices: "${alldevs[*]}"
	debug devices on: "${DEVON[*]}"
	debug devices off: "${DEVOFF[*]}"
}

function init() {
	if ! $nft -a list ruleset | grep -q "table inet filter"; then
		debug "Initialise ruletable: nft add table inet filter"
		$nft add table inet filter
		$nft add chain inet filter INPUT \{ type filter hook input priority 0 \; policy drop \; \}
		$nft add rule inet filter INPUT ct state invalid drop comment \"early drop of invalid packets\"
		$nft add rule inet filter INPUT ct state \{established, related\} accept comment \"accept all connections related to connections made by us\"
		$nft add rule inet filter INPUT iif lo accept comment \"accept loopback\"
		$nft add rule inet filter INPUT iif != lo ip daddr 127.0.0.1/8 drop comment \"drop connections to loopback not coming from loopback ipv4\"
		$nft add rule inet filter INPUT iif != lo ip6 daddr ::1/128 drop comment \"drop connections to loopback not coming from loopback ipv6\"
		$nft add chain inet filter FORWARD \{ type filter hook forward priority 0\; policy drop \; \}
		$nft add chain inet filter OUTPUT \{ type filter hook output priority 0\; policy accept \; \}
	fi
	[[ $INIT == "1" ]] && exit
}

add_device() {
	if [[ -n "${DEVON[*]}" ]]; then
		for _device in "${DEVON[@]}" "${VPNON[@]}"; do
			if ! $nft -a list table inet filter | grep -q "INPUT_${_device}"; then
				debug "Add device: ${_device}"
				$nft add chain inet filter INPUT_"${_device}"
			fi
			if ! $nft -a list ruleset | grep -q "jump INPUT_${_device}"; then
				$nft add rule inet filter INPUT iif "${_device}" jump INPUT_"${_device}"
			fi
			if ! $nft -a list ruleset | grep -i "${_device}" | grep -q "${_device}: accept icmpv4 types"; then
				$nft add rule inet filter INPUT_"${_device}" ip protocol icmp accept comment \""${_device}": accept icmpv4 types\"
			fi
			if ! $nft -a list ruleset | grep -i "${_device}" | grep -q "${_device}: accept icmpv6 types"; then
				$nft add rule inet filter INPUT_"${_device}" ip6 nexthdr icmpv6 accept comment \""${_device}": accept icmpv6 types\"
			fi
		done
	fi
}

remove_device() {
	if [[ -n "${DEVOFF[*]}" ]]; then
		for _device in "${DEVOFF[@]}"; do
			if $nft -a list table inet filter | grep -q "INPUT_${_device}"; then
				debug "Delete device: ${_device}"
				$nft flush chain inet filter INPUT_"${_device}"
				HANDLE=$($nft -a list table inet filter | grep "jump INPUT_${_device}" | awk '{print $NF}')
				$nft delete rule inet filter INPUT handle "${HANDLE}"
				$nft delete chain inet filter INPUT_"${_device}"
			fi
		done 
	fi
}

getports() {
    dev=$1
    istportstcp=()
    istportsudp=()
    turnontcp=()
    turnonudp=()
    turnofftcp=()
    turnoffudp=()
    unchangedtcp=()
    unchangedudp=()
    for _prot in tcp udp; do
        for _port in $($nft -a list ruleset | grep -e dport | grep "${dev}" | grep "${_prot}" | awk '{print $12}' | tr ' ' '\n'); do
            if [[ $_prot == "tcp" ]]; then
            	istportstcp+=("$_port")
            elif [[ $_prot == "udp" ]]; then
                istportsudp+=("$_port")
            fi
        done
    done

    if [[ ! ${dev} == *"vpn"* ]]; then
    	ptcp=("${portstcp[@]}")
    	pudp=("${portsudp[@]}")
    elif [[ ${dev} == *"vpn"* ]]; then
    	ptcp=("${vpn_tcp[@]}")
    	pudp=("${vpn_udp[@]}")
    fi	
    #TCP
    unchangedtcp=()
    for item1 in "${ptcp[@]}"; do
        for item2 in "${istportstcp[@]}"; do
            if [[ $item1 == "$item2" ]]; then
                unchangedtcp+=("$item1")
                break
            fi
        done
    done
    turnofftcp=()
    for item1 in "${istportstcp[@]}"; do
        for item2 in "${ptcp[@]}"; do
            [[ $item1 == "$item2" ]] && continue 2
        done
        turnofftcp+=("$item1")
    done
    if [[ ! $2 == "off" ]]; then
        turnontcp=()
        for item1 in "${ptcp[@]}"; do
            for item2 in "${istportstcp[@]}"; do
                [[ $item1 == "$item2" ]] && continue 2
            done
            turnontcp+=("$item1")
        done
    fi

	#UDP
	unchangedudp=()
    for item1 in "${pudp[@]}"; do
        for item2 in "${istportsudp[@]}"; do
            if [[ $item1 == "$item2" ]]; then
                unchangedudp+=("$item1")
                break
            fi
        done
    done

    turnoffudp=()
    for item1 in "${istportsudp[@]}"; do
        for item2 in "${pudp[@]}"; do
            [[ $item1 == "$item2" ]] && continue 2
        done

        # If we reached here, nothing matched.
        turnoffudp+=("$item1")
    done
    if [[ ! $2 == "off" ]]; then
    turnonudp=()
        for item1 in "${pudp[@]}"; do
            for item2 in "${istportsudp[@]}"; do
                [[ $item1 == "$item2" ]] && continue 2
            done

            # If we reached here, nothing matched.
            turnonudp+=("$item1")
        done
	    if [[ ${dev} == *"vpn"* ]]; then
		    if [[ ! "$VPN_UNRESTRICTED" == "1" ]]; then
				if $nft -a list ruleset | grep -q "${dev}: accept all"; then
					debug "${dev}: remove rule \"${dev}: accept all\""
					HANDLER=$($nft -a list table inet filter | grep "${dev}: accept all" | awk '{print $NF}')
					$nft delete rule inet filter INPUT_"${dev}" handle "${HANDLER}"
				fi
				MAKEVPNRULES=y
			elif [[ "$VPN_UNRESTRICTED" == "1" ]]; then
				if ! $nft -a list ruleset | grep -q "${dev}: accept all"; then
					debug "${dev}: add rule \"${dev}: accept all\""
					$nft add rule inet filter INPUT_"${dev}" accept comment \""${dev}": accept all\"
				fi
				MAKEVPNRULES=n
				turnofftcp=("${turnontcp[@]}" "${unchangedtcp[@]}")
				turnoffudp=("${turnonudp[@]}" "${unchangedudp[@]}")
			fi
	    fi
	fi
}

function addrule() {
	[[ $MAKEVPNRULES == "n" ]] && return
	portnumber=$1
	portname=$2
	protocol=$3
	device=$4
	if ! $nft -a list table inet filter | grep -q "INPUT: allow ${portname} ${portnumber} ${protocol} ${device}"; then
		debug "$device: add rule for ${portname} (${portnumber}, ${protocol})"
		#$nft add rule inet filter INPUT iif "${device}" "${protocol}" dport "${portnumber}" ct state new log accept comment \"INPUT: allow "${portname}" "${portnumber}" "${protocol}" "${device}"\"
		$nft add rule inet filter INPUT_"${_device}" "${protocol}" dport "${portnumber}" ct state new log accept comment \"INPUT: allow "${portname}" "${portnumber}" "${protocol}" "${device}"\"
	fi
}


function removerule() {
	portnumber=$1
	portname=$2
	protocol=$3
	device=$4
	HANDLER=""
	HANDLER=$($nft -a list table inet filter | grep "INPUT: allow ${portname} ${portnumber} ${protocol} ${device}" | awk '{print $NF}')
	if [[ $HANDLER ]]; then
		debug "$device: remove rule for ${portname} (${portnumber}, ${protocol})" 
		#$nft delete rule inet filter INPUT handle "${HANDLER}"
		$nft delete rule inet filter INPUT_"${device}" handle "${HANDLER}"
	fi
}

set_rules() {
	for _device in "${DEVOFF[@]}"; do
		getports "${_device}" off
		for _pn in "${turnofftcp[@]}"; do
			removerule "${ports[$_pn]}" "${_pn}" tcp "${_device}"
		done
		for _pn in "${turnoffudp[@]}"; do
			removerule "${ports[$_pn]}" "${_pn}" udp "${_device}"
		done
		for _pn in "${unchangedtcp[@]}"; do
			removerule "${ports[$_pn]}" "${_pn}" tcp "${_device}"
		done
		for _pn in "${unchangedudp[@]}"; do
			removerule "${ports[$_pn]}" "${_pn}" udp "${_device}"
		done
	done

	for _device in "${DEVON[@]}"; do
		getports "${_device}"
		for _pn in "${turnofftcp[@]}"; do
			removerule "${ports[$_pn]}" "${_pn}" tcp "${_device}"
		done
		for _pn in "${turnoffudp[@]}"; do
			removerule "${ports[$_pn]}" "${_pn}" udp "${_device}"
		done
		for _pn in "${turnontcp[@]}"; do
			addrule "${ports[$_pn]}" "${_pn}" tcp "${_device}"
		done
		for _pn in "${turnonudp[@]}"; do
			addrule "${ports[$_pn]}" "${_pn}" udp "${_device}"
		done
	done
}

help() {
[[ ! $HELP == "1" ]] && return 0
cat <<EOF
usage: $(basename "$0") [-A] [-i] [-a <device>] [-r <device>] [-l <profile>]
                   [-f] [-D] [-h]

    -A 	Automagic find connected devices
    -a 	device to add to firewall (must be configured in cfg)
    -D 	debug
    -f 	flush firewall
    -h 	is help
    -i 	initialise firewall
    -l 	where are we (home|out)
    -L 	list firewall settings after processing
    -n 	use networkmanager to get online devices
    -r 	device to remove from firewall
EOF
exit
}