#!/bin/bash

# BASICS -----------------------------------------------------------------------
# shell options: extended globbing
shopt -qs extglob

# global arrays: script metadata, user config, user commands
declare -Ax script config calltable
script[version]='3.5'
script[name]="${0##*/}"
script[pid]=$$
script[cmd_prefix]='cmd:'
script[verb_levels]='debug info status warning error critical fatal'
script[verb_override]=3
script[external_programs]='which iptables iptables-save iptables-restore tc sysctl readlink sed sort'

# USER CONFIG -----------------------------------------------------------------
# make sure this contains the path for "which", otherwise the script will abort
# in doubt type "which which"
config[which]='/usr/bin/which'

# the rest of the values are case insensitive

# default policy when not specified
# values: accept|drop
config[policy]='accept'

# verbosity level: increment to see more output
# 1:warnings, 2:status, 3:information, 4=debug messages
# error, critical and fatal messages ignore this value, and stop execution
config[verbosity]=2

# EXPOSED FUNCS ----------------------------------------------------------------
cmd1:help() {
    IFS="\n" _msg<<HELP
${script[name]} version ${script[version]}

Iptables helper

Usage: ${script[name]} ${script[valid_commands]// /|}

    help
        Shows a help screen.

    start [profiles]
        Initializes the firewall according to configuration, optionally loading profiles.

    stop
        Removes any iptables state and configuration, and disables forwarding.

    switch <profile>
        Switches iptables configuration to the specified profile.

    load <profiles>
        Load the specified profile(s) into the iptables configuration preserving existing rules.

    save [-f] <profile>
        Stores current iptables configuration and state into a profile. The -f switch causes to
        overwrite an existing profile, otherwise an overwrite attempt triggers an error.

    flush [tables] [chains]
        Flushes (removes all rules) in the specified tables and/or chains.

    zero [tables] [chains]
        Zeroes packet counters in the specified tables and/or chains.

    user [tables]
        Flushes (removes) all user-defined chains in the specified tables.

    clear [tables]
        Equivalent to calling flush+zero+user in sequence.

    policy [tables] [chains] [policy]
        Sets default policy for the specified tables and/or chains.

    Notes on [tables], [chains] and [policy] arguments:
        *) All these arguments are optional, case-insensitive and can appear in any order.
        *) Policy can be one of DROP or ACCEPT, defaulting to user configuration.
        *) Whenever a [tables] or [chains] argument is completely missing, the default is to INCLUDE all tables or chains.
        *) Tables and chains are comma-delimited lists, without spaces.
        *) All the commands that operate with these arguments currently IGNORE bogus strings.

    Notes on profiles:
        *) When multiple profiles are expected they can be separated by comma or space.
        *) By default, profiles are loaded or stored from/to the same directory of this script.
           To specify another location, prepend a relative or absolute path to the profile name.
HELP
}

cmd2:start() {
    _status "${script[name]} ${script[version]} starting..."

    # clear iptables configuration and state, set everything to default policy
    _call clear
    _call policy

    # forward any arguments to "load"
    (( $# )) && {
        _info "a request to load profiles has been received"
        _call load $@
    }

    _status "entering user section..."
    # --------------------------------------------------------------------------
    # ADD YOUR START UP COMMANDS HERE ------------------------------------------
    # --------------------------------------------------------------------------

    # enable kernel forwarding
    $_sysctl -w net.ipv4.ip_forward=1

    # simple qdisc to optimize ethernet
    $_tc qdisc add dev eth0 root tbf rate 256kbit latency 50ms burst 1540

    # here I force my output filter to accept
    _call policy output filter accept

    # --------------------------------------------------------------------------
    # ADD YOUR START UP COMMANDS HERE ------------------------------------------
    # --------------------------------------------------------------------------
    _status 'finished starting process.'
}

cmd3:stop() {
    (( $# )) && _error "extra arguments"

    _status "${script[name]} ${script[version]} stopping..."

    # clear iptables configuration and state, set everything to default policy
    _call clear
    _call policy

    _status "entering user section..."
    # --------------------------------------------------------------------------
    # ADD YOUR SHUTDOWN COMMANDS HERE ------------------------------------------
    # --------------------------------------------------------------------------

    # turn off kernel forwarding
    $_sysctl -w net.ipv4.ip_forward=0

    # remove ethernet qdisc
    $_tc qdisc del dev eth0 root
    
    # --------------------------------------------------------------------------
    # ADD YOUR SHUTDOWN COMMANDS HERE ------------------------------------------
    # --------------------------------------------------------------------------
    _status 'finished stopping process.'
}

cmd4:switch() { 
    (( ! $# )) && _error "<profile> expected"
    (( $# > 1 )) && _error "extra arguments"

    local file=$(_getprofile "$1")
    [[ ! -f "$file" ]] && _error "file not found: $file"

    _status "Switching to profile $file ..."
    $_iptables_restore < "$file"
    _status 'done switch.'
}

cmd5:load() {
    (( ! $# )) && _error "at least one <profile> expected"

    _status "loading profiles..."
    local file
    for file in ${@//,/ }; do
        file=$(_getprofile "$file")
        [[ ! -f "$file" ]] && _error "file not found: $file"

        _info "profile $file"
        $_iptables_restore -n < "$file"
    done
    _status 'done loading.'
}

cmd6:save() {
    (( ! $# )) && _error "<profile> expected"
    (( $# > 1 )) && _error "extra arguments"

    local file overwrite arg
    for arg in $@; do
        [[ "$arg" == '-f' ]] && overwrite='yes' || file=$(_getprofile "$arg")
    done

    [[ -f "$file" ]] && {
        [[ ! $overwrite ]] && _error "file \"$file\" already exists.\\\nIf you want to overwrite it try again with -f" 
        _warning "file \"$file\" already exists, but will overwrite as requested." 
    }

    _status "saving current state to $file ..."
    $_iptables_save > "$file"
    _status 'done saving.'
}

cmd7:flush() {
    _status 'flushing builtin chains...'
    _tbchain flush $@
    _status 'done flushing builtin chains.'
}

cmd8:zero() {
    _status 'zeroing counters...'
    _tbchain zero $@
    _status 'done zeroing counters.'
}

cmd9:user() {
    _status 'flushing user chains...'
    _tbchain user $@
    _status 'done flushing user chains.'
}

cmd10:clear() {
    _call flush $@
    _call zero $@
    _call user $@
}

cmd11:policy() {
    _status 'setting policies...'

    # default policy: user-defined
    local policy="${config[policy]^^}" arg

    # get argument policy
    for arg in $@; do
        [[ "${arg^^}" =~ (ACCEPT|DROP) ]] && policy="${arg^^}" && continue
    done
 
    # set policies
    _tbchain policy $policy $@
    
    _status 'done setting policies.'
}

# INTERNAL FUNCS ---------------------------------------------------------------

# returns a full path to a profile, starting from an ambiguous path/profile
# argument
_getprofile() {
    [[ "$1" =~ ^(/|./|../).* ]] && $_readlink -f "$1" || $_readlink -f "${script[directory]}/$1"
}

# process a command in some table(s) and chain(s)
_tbchain() {
    local -A tables
    local action="$1" arg table chain chainarg
    shift

    # get arguments
    for arg in $@; do
        [[ "${arg^^}" =~ (FILTER|NAT|MANGLE|RAW) ]] && {
            while IFS='' read -rd',' table; do
                tables[$table]=''
            done <<< "${arg,,},"
            continue
        }
        [[ "${arg^^}" =~ (INPUT|OUTPUT|FORWARD|PREROUTING|POSTROUTING) ]] && chainarg="${arg^^}" && continue
    done

    # default tables: all tables
    (( ! "${#tables[@]}" )) && tables=([filter]='' [nat]='' [mangle]='' [raw]='')
    
    # default chains: all chains
    [[ ! "$chainarg" ]] && chainarg="INPUT OUTPUT FORWARD PREROUTING POSTROUTING"

    # ensure that chains are defined only at tables where they exist
    for table in ${!tables[@]}; do
        while read chain; do
            [[ "$chainarg" =~ ($chain) ]] && tables[$table]="$chain ${tables[$table]}"
        done< <($_iptables -L -t $table | $_sed -n 's/^Chain \(.*\) (.*$/\1/p')
    done
    
    # apply actions on selected tables and chains
    for table in ${!tables[@]}; do
        for chain in ${tables[$table]}; do
            case $action in
                policy)
                    # note that $policy is local to policy() but we inherit as
                    # in this case we're its child! that puzzled me for a while
                    # morale: beware with locals inheritance!!!
                    _info "$table $chain $policy"
                    $_iptables -t $table -P $chain $policy 
                    ;;
                flush)
                    _info "$table $chain"
                    $_iptables -t $table -F $chain
                    ;;
                zero)
                    _info "$table $chain"
                    $_iptables -t $table -Z $chain
                    ;;
                user)
                    _info "$table"
                    $_iptables -t $table -X
                    continue 2
                    ;;
            esac
        done
    done
}

_msg() {
    # for 3.6 this function will get a complete rewrite
    # stdin probably will be the only input for message text, and $@
    # will be just for parameters (as god intended ??) hehe
    # proper management/options for stream redir, colors, escapes, etc
    # will be stuff put in place here

    # use params for messages with header
    (( $# )) && {
        echo -e "${script[name]} (${script[pid]}): $@"
        return
    }

    # read messages from stdin
    local stuff
    while read stuff; do
        echo -e "$stuff"
    done
}

# calls exposed functions internally via lookup table
_call() {
    local command="$1"
    shift
    ${calltable[$command]} $@
}

_init() {
    # used in eval expansions
    local macro

    # shortcut functions for messages:
    # visibility is determined @rt by config[verbosity] + script[verb_override]
    # non overridable levels include an exit with corresponding errlevel
    local verbosity level=0
    for verbosity in ${script[verb_levels]}; do
        IFS="" read -d '' macro <<MACRO
        _$verbosity() {
            (( \${config[verbosity]} + $level > ${script[verb_override]} || $level > ${script[verb_override]} )) &&  _msg <<< "${verbosity}: \"\$@\"";
            $( (( $level > ${script[verb_override]} )) && echo "exit $(( level - ${script[verb_override]} ));" )
        }
MACRO
        eval "$macro"
        (( level++ ))
    done

    # globals with validated full path to external programs:
    # we include "which" too for the case the user lefts it misconfigured
    local program path
    macro=''
    for program in ${script[external_programs]}; do
        path="$( ${config[which]} "$program" )"
        [[ ! -x "$path" ]] && _fatal "\"$program\" could not be found or is not executable. Aborting."

        # sanitize varname and spaces in path, push expr into to the macro
        # for varname rules see bash manual, DEFINITIONS section, "name" paragraph
        macro="_${program//@([![:alnum:]])/_}=\"${path// /\ }\"; $macro"
    done
    eval "$macro"

    # find functions named after pattern ${script[cmd_prefix]}. An optional
    # index can be used before the last character of the pattern to aid the
    # help system to sort the commands for presentation.

    # would liked a single datastruct (ie ordered hash) to contain sorted funcs
    # along with the mappings of command<->realfunc, but couldn't do it, so the
    # options were RT overhead at _call() and/or string manips (even) weirder
    # than these below

    # alternatively I could decouple user-cmd/help hinting from funcnames and
    # just use pure hashes. Less nice, less "automatic", IMO...
    # trades: (space|static|verbose config) vs (cpu|dynamic|verbose code)
    # hope for suggestions
    local index command function
    while IFS=" " read index command function; do
        # lookup table mapping cmd names to functions
        calltable[$command]="$function"

        # sorted list of functions exposed as user commands
        # the conditional handles intermixed indexed/non-indexed commands
        (( 10#$index )) && script[valid_commands]="$command${script[valid_commands]:+ ${script[valid_commands]}}" \
                        || script[valid_commands]="${script[valid_commands]:+${script[valid_commands]} }$command"

    done< <(declare -F | $_sed -nr "s/.*(${script[cmd_prefix]%%?}([0-9]*)${script[cmd_prefix]: -1}(.*))/0\2 \3 \1/p" | $_sort -rg)

    # get the script directory
    script[directory]="$($_readlink -f "${0%/*}")"
}

# START ------------------------------------------------------------------------

# initial stuff
_init

# command loop
while IFS="\n" read command; do
    case "$command" in
        "") _msg <<< 'try with ?' ;;
        \?) _msg <<< "valid commands: ${script[valid_commands]}" ;;
        @(${script[valid_commands]// /|})?(+([[:space:]])*)) _call $command ;;
        *) _error "invalid command: $command\\\nto see available commands: ${script[name]} ?" ;;
    esac
done <<< "$@"