#!/bin/bash
#
# Iptables Setup and Initialization Script
# Version 1.0
#
# Written by Michael Crute (mcrute@gmail.com)
# for SoftGroup Interactive on July 10, 2006
# Last updated: same
#

if [ $UID != 0 ]; then
    echo "You are not root. Goodbye."
    exit 1
fi

IPT='/sbin/iptables'

GREEN='\033[0;32m'
WHITE='\033[0m'
RED='\033[0;31m'

ALLOW_PING=0
FUNNY_BUSINESS=0

function do_log 
{
    timestamp=`date +"[%Y-%m-%d %H:%M:%S]"`
    if [ $2 == 0 ]; then
        echo -e "$timestamp $1: [$GREEN OK $WHITE]"
    else
        echo -e "$timestamp $1: [$RED FAILED $WHITE]"
    fi
}

function flush_all 
{
    # Flush all existing chains and erase personal chains
    CHAINS=`cat /proc/net/ip_tables_names 2>/dev/null`
    for i in $CHAINS; do
        $IPT -t $i -F
        $IPT -t $i -X
        do_log "Flushing chain $i" $? 
    done

    # Flush the conntrack database if we can
    if [ `which conntrack > /dev/null 2>&1` ]; then
        conntrack -F
        do_log "Flushing conntrack" $?
    fi

    $IPT -P INPUT ACCEPT
    do_log "All chains flushed" $?
}

function block_shenanigans
{
    iptables -A INPUT -f -j DROP
    do_log "Blocking packet fragments" $?
    iptables -A INPUT -p tcp --tcp-flags ALL NONE -j DROP
    do_log "Blocking null packets" $?
    iptables -A INPUT -p tcp --tcp-flags ALL ALL -j DROP
    do_log "Blocking christmas-tree packets" $?
}

function set_ping
{
    policy="ACCEPT"
    [[ $1 == 0 ]] && policy="DROP"

    $IPT -A INPUT -p icmp --icmp-type 8 -j $policy
    do_log "$policy pings" $?
}

function block_port
{
    $IPT -A INPUT -p $1 -m state --state NEW -m $1 --dport $2 -j DROP
    $IPT -A OUTPUT -p $1 -m state --state NEW -m $1 --dport $2 -j DROP
    do_log "Blocking $1 port $2" $?
}

function load_policy
{
    policy=`cat /etc/firewall/profile 2>/dev/null || echo 'default'`
    [[ ! -z "$1" ]] && policy=$1
    LOADFILE="/etc/firewall/iptables-$policy"

    if [ -r $LOADFILE ]; then
        source $LOADFILE
        flush_all
        do_log "Loading ruleset $LOADFILE" $?
    else
        do_log "Could not load a ruleset" 1
        exit 1
    fi
}

function filter_host
{
    $IPT -I INPUT -s $2 -j $1
    do_log "$1 host $2" $?
}

function filter_port 
{
    $IPT -A INPUT -p $2 -m state --state NEW -m $2 --dport $1 -j $3
    do_log "$3 $2 port $1" $?
}

function add_exception
{
    ip=`echo $1 | cut -d' ' -f1`
    prot=`echo $1 | cut -d' ' -f2`
    port=`echo $1 | cut -d' ' -f3`

    $IPT -A INPUT -p $prot -m state --state NEW -s $ip -m $prot --dport $port -j ACCEPT

    do_log "Allowing $ip for $prot $port" $?
}

function set_defaults
{
    $IPT -A INPUT -s 127.0.0.1 -j ACCEPT
    $IPT -A INPUT -m state --state RELATED,ESTABLISHED -j ACCEPT
    $IPT -P INPUT DROP
    $IPT -P FORWARD DROP
    $IPT -P OUTPUT ACCEPT 
}

# Lockdown is the most strict policy that only passes
# ssh traffic. This is good for new servers and emergencies
function lockdown
{
    flush_all

    $IPT -P INPUT DROP
    $IPT -P FORWARD DROP
    $IPT -P OUTPUT DROP
    do_log "All chains set to drop" $?

    $IPT -A INPUT -m state --state RELATED,ESTABLISHED -j ACCEPT
    $IPT -A OUTPUT -m state --state RELATED,ESTABLISHED -j ACCEPT
    $IPT -A INPUT -p tcp -m state --state NEW -m tcp --dport 22 -j ACCEPT
    $IPT -A OUTPUT -p tcp -m state --state NEW -m tcp --dport 22 -j ACCEPT
    do_log "Allowing ssh only" $?

    do_log "Firewall locked down" $?
}

function build_firewall
{
    load_policy $1
    [[ $FUNNY_BUSINESS == 0 ]] && block_shenanigans
    set_defaults
    set_ping $ALLOW_PING
    
    # Set exceptions
    for idx in $(seq ${#RULES[@]}); do
        item=${RULES[$idx]}
        add_exception $item
    done;

    # Allow or deny hosts
    for host in $HOSTS_DENY; do
        filter_host "DROP" $host
    done
    for host in $HOSTS_ALLOW; do
        filter_host "ALLOW" $host
    done
    
    # Allow or deny tcp ports
    for port in $TCP_DENY; do
        block_port "tcp" $port
    done
    for port in $TCP_ALLOW; do
        filter_port $port "tcp" "ACCEPT"
    done

    # Allow or deny udp ports
    for port in $UDP_DENY; do
        block_port "udp" $port
    done
    for i in $UDP_ALLOW; do
        filter_port $port "udp" "ACCEPT"
    done
}

if [[ $1 == "--flush" ]]; then
    flush_all
    exit 0
elif [[ $1 == "--lockdown" ]]; then
    lockdown
    exit 0
else
    build_firewall $1
    do_log "Firewall started" $?
fi