#! /bin/sh

# firewall v0.38
# Very simple firewall
# Copyright (c) 2023 Raphaël Halimi <raphael.halimi@gmail.com>

#
# Variables
#

# Options defaults
EXEC=0
FLUSH=0
CHECK=0
REMOVE=0
UPDATE=0
VERBOSE=0
DEBUG=0

# Internal variables
IPTABLES_BIN_LIST="iptables ip6tables"
IPTABLES_OPTS="-w"
TABLES_DEFAULT="filter nat mangle raw security"
CHAINS_DEFAULT=" INPUT OUTPUT FORWARD PREROUTING POSTROUTING "
TARGETS_REGEXP="(ACCEPT|AUDIT|CHECKSUM|CLASSIFY|CLUSTERIP|CONNMARK|CONNSECMARK|CT|DNAT|DNPT|DROP|DSCP|ECN|HL|HMARK|IDLETIMER|LED|LOG|MARK|MASQUERADE|NETMAP|NFLOG|NFQUEUE|NOTRACK|RATEEST|REDIRECT|REJECT|SECMARK|SET|SNAT|SNPT|SYNPROXY|TCPMSS|TCPOPTSTRIP|TEE|TOS|TPROXY|TRACE|TTL|ULOG)"
LIB_DIR="/usr/lib/firewall"
ETC_DIR="/etc/firewall"
LIB_CHAINS="$LIB_DIR/chains.d"
ETC_CHAINS="$ETC_DIR/chains.d"
LIB_FUNCTIONS="$LIB_DIR/functions.d"
ETC_FUNCTIONS="$ETC_DIR/functions.d"
LIB_RULES="$LIB_DIR/rules.d"
ETC_RULES="$ETC_DIR/rules.d"
CACHE_DIR="/var/cache/firewall"
TIME="$(date +%s%3N)"
TIME_START="$TIME"


#
# Functions
#

# Exit with an error if the script is not run as root
root_only () {
  if [ "$(id -u)" != "0" ] ; then
    print_error "This script must be run as root, aborting."
    exit 1
  fi
}

# Print error message(s)
# ARGS: STRING...
print_error () {
  local MESSAGE
  for MESSAGE ; do
    printf "%-8s%s\n" "[ERROR]" "$MESSAGE" >&2
  done
}

# Print verbose message(s)
# ARGS: STRING...
print_verbose () {
  local MESSAGE
  if [ $VERBOSE -eq 1 ] ; then
    for MESSAGE ; do
      printf "%-8s%s\n" "[INFO]" "$MESSAGE"
    done
  fi
}

# Print debug message(s)
# ARGS: STRING...
print_debug () {
  local MESSAGE
  if [ $DEBUG -eq 1 ] ; then
    for MESSAGE ; do
      printf "%-8s%s\n" "[DEBUG]" "$MESSAGE"
    done
  fi
}

# Print time
# ARGS: [MESSAGE]
print_time () {
  TIME_OLD="$TIME"
  TIME=$(date +%s%3N)
  print_debug "${1:-checkpoint}, $((TIME-TIME_OLD)) ms"
}

# Debug variables
# ARGS: STRING...
debug_var () {
  local VAR VAL
  if [ $DEBUG -eq 1 ] ; then
    for VAR ; do
      VAL="$(eval "printf %s \"\$$VAR\"")"
      print_debug "$VAR='$VAL'"
    done
  fi
}

# Test IPv4
# ARGS: none
ipv4 () {
  if [ "$IPTABLES_BIN" = "iptables" ] ; then
    return 0
  else
    return 1
  fi
}

# Test IPv6
# ARGS: none
ipv6 () {
  if [ "$IPTABLES_BIN" = "ip6tables" ] ; then
    return 0
  else
    return 1
  fi
}

# Disable incompatible options
# ARGS: MAIN_OPT OPT...
disable_opt () {
  local MAIN_OPT OPT VAL
  MAIN_OPT="$1"
  if [ -z "$MAIN_OPT" ] ; then
    print_error "disable_opt: no argument"
    return 1
  fi
  shift
  for OPT ; do
    VAL="$(eval "printf %d \"\$$OPT\"")"
    if [ $VAL -eq 1 ] ; then
      print_error "Option $MAIN_OPT incompatible with $OPT, ignoring $OPT"
      eval "$OPT=0"
      debug_var $OPT
    fi
  done
}

# Find available tables
# ARGS: none
find_tables () {
  TABLES_AVAILABLE=""
  for TABLE in $TABLES_DEFAULT ; do
    if $IPTABLES -t $TABLE -n -L > /dev/null 2>&1 ; then
      TABLES_AVAILABLE="$TABLES_AVAILABLE $TABLE"
    fi
  done
  debug_var TABLES_AVAILABLE
}

# Get list of chains in a table
# ARGS: TABLE [BUILTIN|USER]
# If "builtin" or "user" is specified, list only built-in or user chains
get_chains_from_table () {
  local TABLE CHAIN TYPE
  TABLE="$1"
  if [ -z "$TABLE" ] ; then
    print_error "get_chains_from_table: no argument"
    return 1
  fi
  TYPE="$2"
  for CHAIN in $($IPTABLES -t $TABLE -n -L | sed -E '/^Chain/!d ; s/^Chain ([^ ]+).*$/\1/') ; do
    if [ -z "$TYPE" ] ; then
      printf "%s\n" "$CHAIN"
    elif [ "$CHAINS_DEFAULT" != "${CHAINS_DEFAULT#* $CHAIN *}" ] ; then
      if [ "$TYPE" = "builtin" ] ; then printf "%s\n" "$CHAIN" ; fi
    else
      if [ "$TYPE" = "user" ] ; then printf "%s\n" "$CHAIN" ; fi
    fi
  done
}

# Test if a chain exists
# ARGS: CHAIN [TABLE...]
chain_exists () {
  local CHAIN TABLES TABLE LIST
  CHAIN="$1"
  if [ -z "$CHAIN" ] ; then
    print_error "chain_exists: no argument"
    return 1
  fi
  shift
  for TABLE in ${@:-$TABLES_AVAILABLE} ; do
    LIST=" $(get_chains_from_table "$TABLE" | xargs) "
    if [ "$LIST" != "${LIST#* $CHAIN *}" ] ; then
      return 0
    fi
  done
  return 1
}

# List preferred files in directories
# ARGS: EXTENSION DIR...
# Specify DIR... in order of preference
list_preferred_files () {
  local TYPE DIR DIRS FILE
  DIRS=""
  TYPE="$1"
  case $TYPE in
    chain|function|rule) ;;
    *) print_error "list_preferred_files: unknown type $TYPE" ; return 1 ;;
  esac
  shift
  for DIR ; do
    if [ -e "$DIR" ] ; then
      DIRS="$DIRS $DIR"
    fi
  done
  if [ -n "$DIRS" ] ; then
    for FILE in $(basename -a $(find $DIRS ! -type d -name "*.$TYPE") | sort -u) ; do
      for DIR in $DIRS ; do
        if [ -e "$DIR/$FILE" ] ; then
          if [ -s "$DIR/$FILE" ] ; then printf "%s\n" "$DIR/$FILE" ; fi
          break
        fi
      done
    done
  fi
  return 0
}

# Source functions
# ARGS: none
source_functions () {
  local DIR FILE
  RC=0
  for FILE in $(list_preferred_files "function" "$ETC_FUNCTIONS" "$LIB_FUNCTIONS") ; do
    print_debug "source function file '$FILE'"
    if [ $CHECK -eq 1 ] ; then
      if [ -s "$FILE" ] ; then printf "%s\n" "$FILE" ; fi
    else
      . "$FILE" || RC=1
    fi
  done
  return $RC
}

# Find chains in files
# ARGS: none
find_chains () {
  local DIR CHAINS FILES
  FILES="$(list_preferred_files "rule" "$ETC_RULES" "$LIB_RULES" ; list_preferred_files "chain" "$ETC_CHAINS" "$LIB_CHAINS")"
  CHAINS="$(sed -n -E "s/($IGNORE|#).*$// ; s/[[:blank:]]+/ /g ; s/^.* -[jN] ([^\$][^ ]+).*$/\1/p" ${FILES:-/dev/null} | sort -u)"
  if [ -n "$CHAINS" ] ; then
    CHAINS_REGEXP="($(printf %s "$CHAINS" | grep -E -w -v "$TARGETS_REGEXP" | xargs | sed 's/ /|/g'))"
    debug_var CHAINS_REGEXP
    print_time "done chains list"
    return 0
  else
    print_error "No chains found"
    return 1
  fi
}

# Get the file where a chain is defined
# ARGS: CHAIN
get_chain_file () {
  local CHAIN DIR FILES FILE
  CHAIN="$1"
  if [ -z "$CHAIN" ] ; then
    print_error "get_chain_file: no argument"
    return 1
  fi
  FILES="$(list_preferred_files "rule" "$ETC_RULES" "$LIB_RULES" ; list_preferred_files "chain" "$ETC_CHAINS" "$LIB_CHAINS")"
  FILE="$(grep -E -l "^[^#]+[[:blank:]]+-N[[:blank:]]+$CHAIN[[:blank:]]*$" ${FILES:-/dev/null} | head -n 1)"
  if [ -n "$FILE" ] ; then
    printf "%s\n" "$FILE"
    return 0
  fi
  return 1
}

# Parse a rule file
# ARGS: RULE_FILE
parse_file () {
  local FILE CHAIN CHAIN_FILE
  FILE="$1"

  if [ -z "$FILE" ] ; then
    print_error "parse_file: no argument"
    return 1
  fi

  # Check if file was already sourced (this shouldn't happen, but still)
  if [ "$SOURCED_FILES" != "${SOURCED_FILES#* $FILE *}" ] ; then
    print_debug "file '$FILE' already sourced"
    return 0
  fi

  # Parse file
  print_debug "parse '$FILE'"

  # CHECK: report only
  if [ $CHECK -eq 1 ] ; then
    if [ -s "$FILE" ] ; then
      if [ "$FILE" != "${FILE%.chain}" ] ; then
        INDENT=$((INDENT+2))
        printf %*s $INDENT
      fi
      printf "%s\n" "$FILE"
    fi
  fi

  # Check if this rule file uses a known chain
  for CHAIN in $(sed -E "s/($IGNORE|#).*$//" "$FILE" | grep -E -w -o "$CHAINS_REGEXP" | sort -u) ; do
    if [ "$CHAINS_FOUND" = "${CHAINS_FOUND#* $CHAIN *}" ] ; then
      print_debug "search chain '$CHAIN'"
      CHAIN_FILE="$(get_chain_file "$CHAIN")"
      if [ -n "$CHAIN_FILE" ] ; then
        CHAINS_FOUND="$CHAINS_FOUND $CHAIN "
        if [ "$CHAIN_FILE" = "$FILE" ] ; then
          print_debug "found: same file"
        else
          print_debug "found: '$CHAIN_FILE'"
          parse_file "$CHAIN_FILE" || RC=1
        fi
      else
        print_error "Unknown chain '$CHAIN' in file '$FILE'"
        RC=1
      fi
    fi
  done

  INDENT=$((INDENT-2))

  # Source file
  if [ $RC -eq 0 ] ; then
    print_debug "source '$FILE'"
    SOURCED_FILES="$SOURCED_FILES $FILE "
    # We should add '|| RC=1' here, but it makes ipv4/ipv6 tests fail, and
    # ruleset is never cached
    if [ $CHECK -eq 0 ] ; then
      . "$FILE"
    fi
  else
    print_error "NOT sourcing file '$FILE'"
    RC=0
  fi
  return $RC
}

# Check cache file
# ARGS: none
cache_check () {
  local FILE TYPE VER_BIN VER_FILE CTIME_ETC CTIME_CACHED
  FILE="$1"
  if [ -z "$FILE" ] ; then
    print_error "cache_check: no argument"
    return 1
  fi
  if [ $UPDATE -eq 0 ] ; then
    if [ -e "$FILE" ] ; then
      TYPE="${FILE##*/}" ; TYPE="${TYPE%%-*}"
      VER_BIN="$($IPTABLES_BIN -V | sed -E "s/^$IPTABLES_BIN v([0-9.]+).*$/\1/")"
      VER_FILE="$(sed -n -E "1 { s/^# Generated by $IPTABLES_BIN-save v([0-9.]+).*$/\1/ ; p ; q }" "$FILE")"
      if [ "$VER_BIN" != "$VER_FILE" ] ; then
        print_debug "$IPTABLES_BIN version mismatch"
        return 1
      else
        if [ "$TYPE" = "rules" ] ; then
          for DIR in "$LIB_DIR" "$ETC_DIR" ; do
            if [ -d "$DIR" ] ; then
              DIRS="$DIRS $DIR"
            fi
          done
          if [ -n "$DIRS" ] ; then
            CTIME_ETC="$(date +%s%N -r $(ls -1trd $(find $DIRS) | tail -n 1))"
            CTIME_CACHED="$(date +%s%N -r "$FILE")"
            if [ $CTIME_ETC -gt $CTIME_CACHED ] ; then
              print_verbose "Configuration changed"
              print_debug "config newer than cache file"
              return 1
            else
              print_debug "config older than cache file"
              return 0
            fi
          else
            print_debug "no config, assume cache file up to date"
            return 0
          fi
        else
          print_debug "$IPTABLES_BIN version match, assume cache file type '$TYPE' up-to-date"
          return 0
        fi
      fi
    else
      print_debug "no cache file"
      return 1
    fi
  else
    print_debug "cache update requested"
    return 1
  fi
}

# Flush ruleset
# ARGS: none
apply_flush () {
  local TABLE CHAIN
  print_verbose "Flush ruleset"
  print_debug "start flush"
  RC=0
  if cache_check "$CACHE_FLUSH" ; then
    cache_load "$CACHE_FLUSH"
  else
    TIME_FLUSH="$(date +%s%3N)"
    for TABLE in $TABLES_AVAILABLE ; do
      print_debug "flush table '$TABLE'"
      $IPTABLES -t $TABLE -F || RC=1
      $IPTABLES -t $TABLE -X || RC=1
      for CHAIN in $(get_chains_from_table $TABLE) ; do
        print_debug "set policy ACCEPT for chain '$CHAIN'"
        $IPTABLES -t $TABLE -P $CHAIN ACCEPT || RC=1
      done
    done
    print_verbose "Ruleset flushed"
    TIME="$TIME_FLUSH"
    print_time "done flush"
    if [ $RC -eq 0 ] ; then cache_save "$CACHE_FLUSH" ; fi
  fi
  return $RC
}

# Basic firewall
# ARGS: none
apply_basic () {
  local CHAIN
  print_verbose "Applying basic firewall"
  print_debug "apply basic"
  RC=0
  if cache_check "$CACHE_BASIC" ; then
    cache_load "$CACHE_BASIC"
  else
    # Flush everything
    apply_flush || RC=1
    TIME_BASE="$(date +%s%3N)"
    # Set policies for chains 'INPUT' and 'FORWARD' in table 'filter' to DROP
    for CHAIN in INPUT FORWARD ; do
      print_debug "set policy DROP for chain '$CHAIN' in table 'filter'"
      $IPTABLES -t filter -P $CHAIN DROP || RC=1
    done
    # Allow everything from loopback
    print_debug "allow from loopback"
    $IPTABLES -t filter -A INPUT -i lo  -j ACCEPT || RC=1
    # Allow ICMP
    print_debug "allow ICMP"
    if ipv4 ; then $IPTABLES -t filter -A INPUT -p icmp -j ACCEPT || RC=1 ; fi
    if ipv6 ; then $IPTABLES -t filter -A INPUT -p ipv6-icmp -j ACCEPT || RC=1 ; fi
    # Allow replies to outgoing connections
    print_debug "allow replies to outgoing connections"
    $IPTABLES -t filter -A INPUT -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT || RC=1
    print_verbose "Basic firewall applied"
    TIME="$TIME_BASE"
    print_time "done basic"
    if [ $RC -eq 0 ] ; then cache_save "$CACHE_BASIC" ; fi
  fi
  return $RC
}

# Apply rules
# ARGS: none
apply_rules () {
  local FILE
  print_verbose "Applying rules"
  print_debug "apply rules"
  RC=0
  if ( cache_check "$CACHE_RULES" && [ $CHECK -eq 0 ] ) ; then
    cache_load "$CACHE_RULES"
  else
    if [ $CHECK -eq 0 ] ; then apply_basic ; fi
    CHAINS_FOUND=""
    CHAINS_REGEXP=""
    SOURCED_FILES=""
    find_chains || RC=1
    print_verbose "Parsing rules"
    TIME_RULES="$(date +%s%3N)"
    for FILE in $(list_preferred_files "rule" "$ETC_RULES" "$LIB_RULES") ; do
      if [ -s "$FILE" ] ; then
        INDENT=0
        parse_file "$FILE" || RC=1
        print_time "done parse"
      fi
    done
    print_verbose "Rules applied"
    TIME="$TIME_RULES"
    print_time "done rules"
    if ( [ $RC -eq 0 ] && [ $CHECK -eq 0 ] ) ; then cache_save "$CACHE_RULES" ; fi
  fi
  return $RC
}

# Load cache
# ARGS: none
cache_load () {
  local FILE
  FILE="$1"
  TIME_LOAD="$(date +%s%3N)"
  if [ -z "$FILE" ] ; then
    print_error "cache_load: no argument"
    return 1
  fi
  RC=0
  print_verbose "Loading '${FILE##*/}'"
  print_debug "run $IPTABLES_BIN-restore '$FILE'"
  $IPTABLES_BIN-restore "$FILE"
  RC=$?
  TIME="$TIME_LOAD"
  print_time "done cache load"
  return $RC
}

# Save cache
# ARGS: none
cache_save () {
  local FILE
  FILE="$1"
  TIME_SAVE="$(date +%s%3N)"
  if [ -z "$FILE" ] ; then
    print_error "cache_save: no argument"
    return 1
  fi
  RC=0
  if ! [ -d "$CACHE_DIR" ] ; then
    mkdir -p "$CACHE_DIR" || RC=1
  fi
  print_verbose "Saving '${FILE##*/}'"
  print_debug "run $IPTABLES_BIN-save | sed -E 's/\[[0-9]+:[0-9]+\]/[0:0]/' > '$FILE'"
  $IPTABLES_BIN-save | sed -E 's/\[[0-9]+:[0-9]+\]/[0:0]/' > "$FILE"
  RC=$?
  TIME="$TIME_SAVE"
  print_time "done cache save"
  return $RC
}

# Print help option
# ARGS: OPTION DESCRIPTION
print_option() {
  printf "  %-20s%s\n" "$1" "$2"
}

# Print help
# ARGS: none
print_help () {
  printf "Usage: %s [OPTION...]\n" "$(basename -- "$0")"
  printf "\nOPTIONS:\n"
  print_option "-f" "Flush tables"
  print_option "-r" "Remove cache"
  print_option "-u" "Update cache"
  print_option "-c" "Check"
  print_option "-v" "Verbose mode"
  print_option "-d" "Debug mode"
  print_option "-h" "Print this help message"
  print_option "-e COMMAND" "Execute command (must be last option)"
}


#
# Options processing
#

while getopts "cefruvdh" OPTION ; do
  case $OPTION in
    c) CHECK=1 ;;
    e) EXEC=1 ; break ;;
    f) FLUSH=1 ;;
    r) REMOVE=1 ;;
    u) UPDATE=1 ;;
    v) VERBOSE=1 ;;
    d) DEBUG=1 ;;
    h) print_help ; exit 0 ;;
    *) print_help ; exit 1 ;;
  esac
done ; shift $((OPTIND-1))


#
# Checks
#

if [ $CHECK -eq 0 ] ; then root_only ; fi

# Debug variables
debug_var \
  EXEC \
  FLUSH \
  CHECK \
  REMOVE \
  UPDATE \
  VERBOSE \
  DEBUG \
  IPTABLES_BIN_LIST \
  IPTABLES_OPTS \
  TABLES_DEFAULT \
  CHAINS_DEFAULT \
  LIB_DIR \
  ETC_DIR \
  LIB_CHAINS \
  ETC_CHAINS \
  LIB_FUNCTIONS \
  ETC_FUNCTIONS \
  LIB_RULES \
  ETC_RULES

for ARG ; do
  debug_var ARG
done


#
# Source external functions
#

RC=0 ; source_functions || RC=1

print_time "done init"


#
# Main
#

for IPTABLES_BIN in $IPTABLES_BIN_LIST ; do

  IPV=""
  TIME_IPV="$(date +%s%3N)"

  # Necessary for check mode without root
  if ( [ $CHECK -eq 1 ] && [ "$(id -u)" != "0" ] ) ; then
    PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
  fi

  # Warn and skip loop if iptables binary is not installed
  IPTABLES="$(command -v $IPTABLES_BIN)"

  if [ -z "$IPTABLES" ] ; then
    print_error "Cannot find binary '$IPTABLES_BIN'"
    continue
  fi

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

  # Set some variables
  if ipv4 ; then
    IPV="IPv4"
    IGNORE="ipv6"
  else
    IPV="IPv6"
    IGNORE="ipv4"
  fi

  CACHE_FLUSH="$CACHE_DIR/flush-$IPV.ipt"
  CACHE_BASIC="$CACHE_DIR/basic-$IPV.ipt"
  CACHE_RULES="$CACHE_DIR/rules-$IPV.ipt"

  # Debug variables
  debug_var \
    IPTABLES_BIN \
    IPTABLES \
    IPV \
    IGNORE \
    CACHE_FLUSH \
    CACHE_BASIC \
    CACHE_RULES

  print_time "done $IPV checks"

  print_verbose "Starting firewall - $IPV"
  print_debug "start $IPV"

  # CHECK: list sourced rules files
  if [ $CHECK -eq 1 ] ; then
    disable_opt CHECK FLUSH REMOVE EXEC UPDATE
    printf "%s\n" "----- $IPV -----"
    if ( [ -d "$LIB_RULES" ] || [ -d "$ETC_RULES" ] ) ; then
      apply_rules
    else
      print_error "No rules found"
    fi
    continue
  fi

  # Find available tables
  find_tables

  # FLUSH: flush ruleset and exit
  if [ $FLUSH -eq 1 ] ; then
    disable_opt FLUSH EXEC UPDATE
    apply_flush
    RC=$?
    if [ $REMOVE -eq 0 ] ; then continue ; fi
  fi

  # REMOVE: remove cache files and exit
  if [ $REMOVE -eq 1 ] ; then
    disable_opt REMOVE EXEC UPDATE
    RC=0
    for CACHE in "$CACHE_FLUSH" "$CACHE_BASIC" "$CACHE_RULES" ; do
      print_debug "remove cache file '$CACHE'"
      rm -f "$CACHE" || RC=1
    done
    print_time "done remove"
    continue
  fi

  # EXEC: execute single command
  if [ $EXEC -eq 1 ] ; then
    disable_opt EXEC UPDATE
    print_debug "execute command: $*"
    eval "$@"
    RC=$?
    print_time "done exec"
    continue
  fi

  # Apply ruleset
  if ( [ -d "$LIB_RULES" ] || [ -d "$ETC_RULES" ] ) ; then
    apply_rules
    RC=$?
  else
    apply_basic
    RC=$?
  fi

  TIME="$TIME_IPV"
  print_time "done $IPV"

done

TIME="$TIME_START"
print_time "done all"
exit $RC
