#!/bin/bash
set -e

# Variables
IPT=$(which iptables)
IPCMD=$(which ip)
MAINIP=$("${IPCMD}" -4 -o a | grep inet | grep -v "lo " | awk '{print $4}' | cut -d "/" -f 1)
DPORTSIN="20:22,80,443"
MOSHIN="60000:61000"

# Trap errors
check_error()
{
   if [ "$?" = 0 ]; then
       echo "===> OK";
   else
       echo "===> ERROR!";
       exit
   fi
}

trap check_error EXIT

# Flush rules and set default rules
${IPT} -F
${IPT} -P INPUT DROP
${IPT} -P FORWARD DROP
${IPT} -P OUTPUT ACCEPT

# Set host sepcific rules
### Accept all data from "lo" interface
${IPT} -A INPUT -i lo -j ACCEPT

### Accept all connections after the three-way handshake
${IPT} -A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT

### Block country zones
if [ "${IPSET}" == "" ]; then
        echo "No ipset command found. Skipping"
else
    for chain in $(${IPSET} list -name);
    do if [ "${chain}" == "" ]; then
	   echo "No ipset chain found. Skipping"
       else
	   ${IPT} -A INPUT -p tcp -m set --match-set ${chain} src -j DROP;
       fi
    done
fi

### Services rules
##### Allow connections to the server main ip for SSH, MOSH and psyBNC
##### Use multiports extension to set more ports in a oneline rule
${IPT} -A INPUT -p tcp -d ${MAINIP} -m multiport --dports ${DPORTSIN} -j ACCEPT
${IPT} -A INPUT -p udp -d ${MAINIP} -m multiport --dports ${MOSHIN} -j ACCEPT