#! /bin/sh

#
# Firewall script
#

#
# Options
#

IPTABLES_BIN_LIST="/sbin/iptables /sbin/ip6tables"
IPTABLES_OPTS="-w"
CHAINS_DEFAULT_LIST="INPUT OUTPUT FORWARD PREROUTING POSTROUTING"
VERBOSE=0
FLUSH_ONLY=0
DUMP_ONLY=0
INDENT=0
USER_RULES_DIR="/etc/firewall.d"
SYSTEM_RULES_DIR="/lib/firewall/rules.d"


#
# Functions
#

# Print usage
# ARGS: none
print_usage () {
  printf "Usage: %s [OPTIONS...]\n" "$0"
  printf "\nOPTIONS:\n"
  printf "  %-20s%s\n" "--dump" "Dump rules on standard output"
  printf "  %-20s%s\n" "--stop" "Stop firewall"
  printf "  %-20s%s\n" "--verbose" "Verbose mode"
}

# Print verbose message
# ARGS: string
print_verbose () {
  [ $VERBOSE -eq 1 ] && for MESSAGE in "$@" ; do
    printf "%${INDENT}s%s\n" "" "$MESSAGE"
  done
}

# Test if $IPTABLES is IPv4
# ARGS: none
ipv4 () {
  if printf "%s" "$IPTABLES" | grep -q 6 ; then
    return 1
  else
    return 0
  fi
}

# Test if $IPTABLES is IPv6
# ARGS: none
ipv6 () {
  if printf "%s" "$IPTABLES" | grep -q 6 ; then
    return 0
  else
    return 1
  fi
}

# Get list of available tables
# ARGS: none
get_tables () {
  local TABLE
  for TABLE in filter nat mangle raw security ; do
    if $IPTABLES -t $TABLE -n -L > /dev/null 2>&1 ; then
      printf "%s\n" "$TABLE"
    fi
  done
}

# Get list of chains in a table
# ARGS: table
# WARNING: case-sensitive, table names must be lowercase
get_chains_from_table () {
  local TABLE CHAINS
  TABLE="$1"
  for CHAIN in $(${IPTABLES#echo } -t $TABLE -n -L | sed -E '/^Chain/!d ; s/^Chain ([^[:blank:]]+).*$/\1/') ; do
    if printf "%s" "$CHAINS_DEFAULT_LIST" | grep -w -q "$CHAIN" ; then
      printf "%s\n" "$CHAIN"
    fi
  done
}

# Add a chain name to the list of defined chains
# ARGS: CHAIN_NAME
add_chain_defined () {
  local CHAIN
  CHAIN="$1"
  if ! printf "%s" "$CHAINS_DEFINED_LIST" | grep -w -q "$CHAIN" ; then
    CHAINS_DEFINED_LIST="$CHAINS_DEFINED_LIST $CHAIN"
  fi
}

# Test if a chain name is in the list of defined chains
# ARGS: CHAIN_NAME
is_chain_defined () {
  local CHAIN
  CHAIN="$1"
  if printf "%s" "$CHAINS_DEFINED_LIST" | grep -w -q "$CHAIN" ; then
    return 0
  else
    return 1
  fi
}

# Build list of available chains
# ARGS: none
build_chains_available_list () {
  local RULE_FILE CHAIN
  CHAINS_AVAILABLE_LIST=""
  CHAINS_AVAILABLE_REGEXP=""
  print_verbose "Building list of user-defined chains"
  for RULE_FILE in $USER_RULES_DIR/*.rule $SYSTEM_RULES_DIR/*.rule ; do
    INDENT=$((INDENT+2))
    for CHAIN in $(sed -E "s/($IGNORE|#).*$// ; /-N[[:blank:]]+\\$/d ; /-N/!d ; s/^.*[[:blank:]]+-N[[:blank:]]+([^[:blank:]]+).*$/\1/" "$RULE_FILE") ; do
      if ! printf "%s" "$CHAINS_AVAILABLE_LIST" | grep -w -q "^$CHAIN=" ; then
        print_verbose "Chain $CHAIN available in $RULE_FILE"
        CHAINS_AVAILABLE_LIST="$(printf "%s\n%s" "$CHAINS_AVAILABLE_LIST" "$CHAIN=$RULE_FILE")"
        CHAINS_AVAILABLE_REGEXP="$CHAINS_AVAILABLE_REGEXP|$CHAIN"
      fi
    done
    INDENT=$((INDENT-2))
  done
  CHAINS_AVAILABLE_REGEXP="(${CHAINS_AVAILABLE_REGEXP#|})"
}

# Get the rule file where an available chain is defined
# ARGS: CHAIN_NAME
get_chain_file () {
  local CHAIN_NAME
  CHAIN_NAME="$1"
  printf "%s" "$CHAINS_AVAILABLE_LIST" | sed "/^$CHAIN_NAME=/!d ; s/^$CHAIN_NAME=//"
}

# Parse a rule file
# ARGS: file_name
parse_rule_file () {
  local PARSED_RULE_FILE CHAINS_LIST_NEEDED CHAIN_NAME CHAIN_FILE
  PARSED_RULE_FILE="$1"

  # Check if file was already parsed
  if printf "$SOURCED_FILES_LIST" | grep -w -q "$PARSED_RULE_FILE" ; then
    print_verbose "Rule file $PARSED_RULE_FILE already sourced"
    return
  fi

  # Parse file
  print_verbose "Parsing rule file $PARSED_RULE_FILE"
  INDENT=$((INDENT+2))

  # Check if this rule file needs any of the available user-defined chains
  if ipv4 ; then IGNORE="ipv6" ; else IGNORE="ipv4" ; fi
  for CHAIN_NAME in $(sed -E "s/($IGNORE|#).*$//" "$PARSED_RULE_FILE" | grep -E -w -o "$CHAINS_AVAILABLE_REGEXP") ; do
    if ! printf "%s" "$CHAINS_DEFINED_LIST" | grep -w -q "$CHAIN_NAME" ; then
      CHAIN_FILE="$(get_chain_file "$CHAIN_NAME")"
      if [ "$CHAIN_FILE" != "$PARSED_RULE_FILE" ] ; then
        print_verbose "Chain $CHAIN_NAME available in $CHAIN_FILE"
        parse_rule_file "$CHAIN_FILE"
        add_chain_defined "$CHAIN_NAME"
      fi
    fi
  done

  # Source file
  INDENT=$((INDENT-2))
  print_verbose "Sourcing rule file $PARSED_RULE_FILE"
  SOURCED_FILES_LIST="$SOURCED_FILES_LIST $PARSED_RULE_FILE"
  . "$PARSED_RULE_FILE"
}


#
# Option handling
#

for OPTION in "$@" ; do
  shift
  case "$OPTION" in
    --dump) DUMP_ONLY=1 ;;
    --stop) FLUSH_ONLY=1 ;;
    --verbose) VERBOSE=1 ;;
    *) print_usage ; exit 1 ;;
  esac
done


#
# MAIN
#

for IPTABLES_BIN in $IPTABLES_BIN_LIST ; do

  #
  # Start
  #

  # Skip loop if iptables binary is not installed
  test -x "$IPTABLES_BIN" || continue

  # Join binary and options
  if [ -n "$IPTABLES_OPTS" ] ; then
    IPTABLES="$IPTABLES_BIN $IPTABLES_OPTS"
  else
    IPTABLES="$IPTABLES_BIN"
  fi

  # Only echo commands if DUMP_ONLY is enabled
  [ $DUMP_ONLY -eq 1 ] && IPTABLES="echo $IPTABLES"

  # Report start and set ignore values
  if ipv4 ; then
    print_verbose "Starting firewall - IPv4"
    IGNORE="ipv6"
  else
    print_verbose "" "Starting firewall - IPv6"
    IGNORE="ipv4"
  fi


  #
  # Reset all tables
  #

  print_verbose "Reset all tables"
  INDENT=$((INDENT+2))
  for TABLE in $(get_tables) ; do
    print_verbose "Flushing table $TABLE"
    $IPTABLES -t $TABLE -F
    $IPTABLES -t $TABLE -X
    INDENT=$((INDENT+2))
    for CHAIN in $(get_chains_from_table $TABLE) ; do
      print_verbose "Setting policy for $CHAIN chain to ACCEPT"
      $IPTABLES -t $TABLE -P $CHAIN ACCEPT
    done
    INDENT=$((INDENT-2))
    print_verbose ""
  done
  INDENT=$((INDENT-2))

  # If we want to stop the firewall, we're done
  [ $FLUSH_ONLY -eq 1 ] && continue


  #
  # Basic firewall
  #

  print_verbose "Setting basic firewall"
  INDENT=$((INDENT+2))

  # Drop everything by default
  for CHAIN in $(get_chains_from_table filter) ; do
    print_verbose "Setting policy for $CHAIN chain in filter table to DROP"
    $IPTABLES -t filter -P $CHAIN DROP
  done

  # Accept all from loopback device
  print_verbose "Accepting everything from/to localhost"
  $IPTABLES -t filter -A INPUT -i lo  -j ACCEPT
  $IPTABLES -t filter -A OUTPUT -o lo -j ACCEPT

  # Accept outgoing connections and their answers
  print_verbose "Accepting outgoing connections and their answers"
  ipv4 && $IPTABLES -t filter -A OUTPUT -m conntrack --ctstate NEW,RELATED,ESTABLISHED -j ACCEPT
  ipv6 && $IPTABLES -t filter -A OUTPUT -m conntrack --ctstate NEW,RELATED,ESTABLISHED,UNTRACKED -j ACCEPT
  $IPTABLES -t filter -A INPUT -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
  INDENT=$((INDENT-2))
  print_verbose ""


  #
  # Advanced firewall
  #

  # Rules are stored in $USER_RULES_DIR
  if [ -d $USER_RULES_DIR ] ; then
    build_chains_available_list
    SOURCED_FILES_LIST="" ; CHAINS_DEFINED_LIST=""
    for RULE_FILE in $USER_RULES_DIR/*.rule ; do
      if [ -r "$RULE_FILE" ] ; then
        print_verbose ""
        parse_rule_file "$RULE_FILE"
      fi
    done
  fi

done
