scripts: better error handling and revert for set_rules

Signed-off-by: HeshamTB <hishaminv@gmail.com>

scripts/set_rules.sh

@@ -0,0 +1,101 @@
+#!/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
+