diff --git a/scripts/set_rules.sh b/scripts/set_rules.sh
index 123130a..09d7598 100755
--- a/scripts/set_rules.sh
+++ b/scripts/set_rules.sh
@@ -2,35 +2,97 @@
 
 iptables=iptables
 
+PROGRAM="${0}"
 cmd="${1}"
-wgIface="${2}"
-uplinkIface="${3}"
+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
-    cmd ${iptables} -A FORWARD -i ${wgIface} -j ACCEPT 
-    cmd ${iptables} -A FORWARD -o ${wgIface} -j ACCEPT
-    cmd ${iptables} -t nat -A POSTROUTING -o ${uplinkIface} -j MASQUERADE || exit 1
+    add_forward || exit 1
+    add_nat || exit 1
+    add_isolation || exit 1
+    add_port || exit 1
     trap - INT TERM EXIT
 }
 
-rm_rules() {
-    cmd ${iptables} -D FORWARD -i ${wgIface} -j ACCEPT
-    cmd ${iptables} -D FORWARD -o ${wgIface} -j ACCEPT
-    cmd ${iptables} -t nat -D POSTROUTING -o ${uplinkIface} -j MASQUERADE
+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