#!/bin/bash
#config
alldevs=(enp0s25 wlp2s0 vpn0 vpn-gateway)

hometcp=(ssh vncd misc)
homeudp=(avahi-daemon)
outtcp=(ssh)
outudp=()

declare -A ports=(
  [ssh]=22
  [smtp]=25
  [dns]=53
  [netbios-ssn]=139
  [dovecot]=143
  [microsoft-ds]=445
  [cups-ipp]=631
  [imaps]=993
  [kdeconnect]=1714-1764
  [avahi-daemon]=5353
  [vncd]=5900
  [mpd]=6600
  [mpd-audio]=8000
  [misc]=8080
  [syncthing]=8384
  [wireguard]=51820
  [mosh]=60000-61000
)

nft="/usr/sbin/nft"

##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:l:fDhnL option
do
	case "${option}" in
		A) AUTOMATIC=1;;
		i) INIT=1;;
		a) ADEVICE+=("${OPTARG}");;
		r) RDEVICE+=("${OPTARG}");;
		l) LOCATION="${OPTARG}";;
		f) FLUSH=1;;
		D) DEBUG=1;;
		h) HELP=1;;
		n) NM=1;;
		L) LIST=1;;
		*) HELP=1;;
	esac
done

help
flush
check_deviceinput
get_devices
location
init
add_device
remove_device
set_rules
[[ $LIST == "1" ]] && echo " " && sudo nft list ruleset -a
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
}

location() {
	[[ $INIT == "1" ]] && return 0
	if [[ -z $LOCATION ]]; then
		if [[ -e /run/home ]]; then
			HOME=$(cat /run/home)
			if [[ "${HOME}" == "home" ]]; then
				LOCATION=home
			elif [[ "${HOME}" == "out" ]]; then
				LOCATION=out
			else
				LOCATION=out
			fi
		fi
	fi
	case $LOCATION in
		h|home)
		portstcp=("${hometcp[@]}")
		portsudp=("${homeudp[@]}");;
		o|out)
		portstcp=("${outtcp[@]}")
		portsudp=("${outudp[@]}") ;;
	esac
	debug location: "${LOCATION[@]}"
	debug tcp ports: "${portstcp[@]}"
	debug udp ports: "${portsudp[@]}"
}

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
				DEVON+=("$i")
			done
		else
			for i in $(ip -oneline addr show scope global|awk '{print $2}'|uniq); do
				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

	debug alldevices: "${alldevs[*]}"
	debug devices on: "${DEVON[*]}"
	debug devices off: "${DEVOFF[*]}"
}

function init() {
	if ! $nft list ruleset -a |grep -q "table inet filter"; then
		debug "Initialise rule: 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 list table inet filter -a|grep -q "INPUT_${_device}"; then
				debug "Add device: ${_device}"
				$nft add chain inet filter INPUT_"${_device}"
			fi
			if ! $nft list ruleset -a |grep -q "jump INPUT_${_device}"; then
				$nft add rule inet filter INPUT iif "${_device}" jump INPUT_"${_device}"
			fi
			if ! $nft list ruleset -a |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 list ruleset -a |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
			if [[ $_device == *"vpn"* ]]; then
				if ! $nft list ruleset -a |grep -q "$_device: accept all"; then
					debug "$_device: add rule \"$_device: accept all\""
					$nft add rule inet filter INPUT_"${_device}" accept comment \""$_device": accept all\"
				fi
			fi
		done
	fi
}

remove_device() {
	if [[ -n "${DEVOFF[*]}" ]]; then
		for _device in "${DEVOFF[@]}"; do
			if $nft list table inet filter -a|grep -q "INPUT_${_device}"; then
				debug "Delete device: ${_device}"
				$nft flush chain inet filter INPUT_"${_device}"
				HANDLE=$($nft list table inet filter -a|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 list ruleset -a|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

    #TCP
    for item1 in "${portstcp[@]}"; 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 "${portstcp[@]}"; do
            [[ $item1 == "$item2" ]] && continue 2
        done
        turnofftcp+=("$item1")
    done
    if [[ ! $2 == "off" ]]; then
        turnontcp=()
        for item1 in "${portstcp[@]}"; do
            for item2 in "${istportstcp[@]}"; do
                [[ $item1 == "$item2" ]] && continue 2
            done
            turnontcp+=("$item1")
        done
    fi


	#UDP
    for item1 in "${portsudp[@]}"; 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 "${portsudp[@]}"; do
            [[ $item1 == "$item2" ]] && continue 2
        done

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

            # If we reached here, nothing matched.
            turnonudp+=("$item1")
        done
    fi
}

function addrule() {
	portnumber=$1
	portname=$2
	protocol=$3
	device=$4
	if ! $nft list table inet filter -a |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 list table inet filter -a |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
		if [[ ! ${_device} == *"vpn"* ]]; then
			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
		fi
	done

	for _device in "${DEVON[@]}"; do
		if [[ ! "${_device}" == *"vpn"* ]]; then
			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
		fi
	done
}

help() {
[[ ! $HELP == "1" ]] && return 0
cat <<EOF
usage: $(basename "$0") [-A] [-i] [-a <device>] [-r <device>] [-l <location>]
                   [-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
}