#!/usr/bin/env bash 

iptables=iptables

PROGRAM="${0}"
cmd="${1}"
wg_iface="${2}"
uplink_iface="${3}"
wg_port="${4}"

cmd() {
	echo "[#] $*" >&2
	"$@"
}

print_usage() {
	cat >&2 <<-_EOF
    ${PROGRAM} [ set | unset ] <wg_interface> <uplink_interface> <wg_port> 
	_EOF
}

HAVE_SET_NAT=0
add_nat() {
    cmd ${iptables} -t nat -A POSTROUTING -o ${uplink_iface} -j MASQUERADE
    HAVE_SET_NAT=1
}

HAVE_SET_FORWARD=0
add_forward() {
    cmd ${iptables} -I FORWARD 1 -i ${uplink_iface} -o ${wg_iface} -j ACCEPT
    cmd ${iptables} -I FORWARD 1 -i ${wg_iface} -o ${uplink_iface} -j ACCEPT
    HAVE_SET_FORWARD=1
}

HAVE_SET_ISOLATION=0
add_isolation() {
    cmd ${iptables} -I FORWARD -i ${wg_iface} -o ${wg_iface} -j DROP
    HAVE_SET_ISOLATION=1
}

HAVE_OPEN_PORT=0
add_port() {
    cmd ${iptables} -I INPUT 1 -i ${uplink_iface} -p udp --dport ${wg_port} -j ACCEPT
    HAVE_OPEN_PORT=1
}

add_rules() {
    trap 'rm_rules; exit' INT TERM EXIT
    add_forward || exit 1
    add_nat || exit 1
    add_isolation || exit 1
    add_port || exit 1
    trap - INT TERM EXIT
}

rm_forward() {
    cmd ${iptables} -D FORWARD -i ${uplink_iface} -o ${wg_iface} -j ACCEPT
    cmd ${iptables} -D FORWARD -i ${wg_iface} -o ${uplink_iface} -j ACCEPT
}

rm_nat() {
    cmd ${iptables} -t nat -D POSTROUTING -o ${uplink_iface} -j MASQUERADE
}

rm_isolate() {
    cmd ${iptables} -D FORWARD -i ${wg_iface} -o ${wg_iface} -j DROP
}

rm_port() {
    cmd ${iptables} -D INPUT -i ${uplink_iface} -p udp --dport ${wg_port} -j ACCEPT
}

rm_rules() {
    [[ $HAVE_SET_FORWARD -eq 0 ]] || rm_forward
    [[ $HAVE_SET_NAT -eq 0 ]] || rm_nat
    [[ $HAVE_SET_ISOLATION -eq 0 ]] || rm_isolate
    [[ $HAVE_OPEN_PORT -eq 0 ]] || rm_port
}

if [ ! $# -eq 4 ]
then
    print_usage
    exit 1
fi

if [ "${cmd}" == "set" ]
then
    add_rules

elif [ "${cmd}" == "unset" ]; 
then
    HAVE_OPEN_PORT=1
    HAVE_SET_ISOLATION=1
    HAVE_SET_NAT=1
    HAVE_SET_FORWARD=1
    rm_rules
else
   # cat << "Invalid command. Use set or unset" >&2
   echo "Invalid command. Use set or unset"
fi