Domain/hostname DNS resolution for iptables

iptables-resolve.sh

#!/bin/bash

##### Domain/Hostname Resolution for iptables #####
# This script performs a DNS query to get a list of IPs associated with a domain.
# You can then create a cron as root to be run however often your needs require.
#
# Usage: $HOME/.local/bin/iptables-resolve [-a|--add] [-t|--test-run] [-j|--jump=<target>] <hostname>
# Example: $HOME/.local/bin/iptables-resolve -atj DROP www.example.com
#
# If no options are given, the default action is to add all resolved IPs as ACCEPT.
# Unfortunately, a [-d|--delete] option is not yet available to use.
# Another current requirement is that the INPUT rule for the chain must be added manually. The functionality does not yet exist to check for an existing rule.
# Check the TODO list for details.
# 
# TODO: Check if INPUT rule exists for chain before attempting to add it. Otherwise, this rule will keep getting created regardless if it already exists.
# TODO: Capture the chain spec in order to delete it from the INPUT. Otherwise, the chain itself is not able to delete using "iptables -X $CHAIN".

# Get short/long options https://stackoverflow.com/a/28466267/5825414
die() { echo "$*" >&2; exit 2; }  # complain to STDERR and exit with error
needs_arg() { if [ -z "$OPTARG" ]; then die "No arg for --$OPT option"; fi; }

while getopts adtj:-: OPT; do
  # support long options: https://stackoverflow.com/a/28466267/519360
  if [ "$OPT" = "-" ]; then   # long option: reformulate OPT and OPTARG
    OPT="${OPTARG%%=*}"       # extract long option name
    OPTARG="${OPTARG#$OPT}"   # extract long option argument (may be empty)
    OPTARG="${OPTARG#=}"      # if long option argument, remove assigning `=`
  fi
  case "$OPT" in
    a | add )       ACTION='add' ;;
    # d | delete )    ACTION='delete' ;;
    t | test-run )  TEST=1 ;; # dry-run/test-run
    j | jump )      needs_arg; JUMP="${OPTARG^^}" ;; # Uppercase jump/target
    ??* )           die "Illegal option --$OPT" ;;  # bad long option
    ? )             exit 2 ;;  # bad short option (error reported via getopts)
  esac
done
shift $((OPTIND-1)) # remove parsed options and args from $@ list


# Display something if minimum arguments are not met.
if [ 1 -gt $# ]; then
  echo "usage: $0 [-a|--add] [-t|--test-run] [-j|--jump=<target>] <hostname>"
  echo "If no options are given, default action is to add."
  echo "[-d|--delete] is not yet available to use."
  echo "Example: $0 -atj DROP www.example.com"
  exit 1
fi

# Stop script if a command fails.
set -e

# Get hostname/domain from args.
HOSTNAME=$1

# If action is not defined, set it to add
if [ -z "$ACTION" ]; then
  ACTION='add'
fi

# Get the list of IP addresses.
IP_LIST=$(dig $HOSTNAME +short | grep -E "^((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$")

# Set the chain name.
CHAIN="${HOSTNAME//./-}"

# Supported jump targets
TARGET_ARRAY=("ACCEPT" "DROP" "QUEUE" "RETURN")

if [ -z "$JUMP" ]; then
  # Set default if target is empty.
  TARGET="ACCEPT"
elif printf '%s\0' "${TARGET_ARRAY[@]}" | grep -Fxqz -- $JUMP; then
  # If target is in array, this script supports it.
  TARGET=${JUMP}
else
  echo "Jump/Target not supported: $JUMP"
  exit 1
fi

if [ "$IP_LIST" ]; then
  echo "${TARGET} IP addresses:"
  echo "$IP_LIST"
else
  echo "No A records listed for $HOSTNAME"
  exit 1
fi

if [ 'add' = "$ACTION" ]; then
  # Add chain and insert rules. Flush chain if already exists.
  if [ "$TEST" = 1 ]; then
    echo "Chain to be created: $CHAIN"
    echo "Rules to be created:"
    for IP in $IP_LIST; do echo "-A $CHAIN -s $IP -p tcp -m multiport --dports 80,443 -j $TARGET"; done
  else
    sudo iptables -N $CHAIN && echo "Chain Created: $CHAIN" || sudo iptables -F $CHAIN && echo "Existing $CHAIN chain flushed." # Create or flush chain
    for IP in $IP_LIST; do sudo iptables -A $CHAIN -s $IP -p tcp -m multiport --dports 80,443 -j $TARGET; done # Add IPs to chain
    # sudo iptables -A INPUT -j $CHAIN # Add chain to main INPUT
  fi
# elif [ 'delete' = "$ACTION" ]; then
#   # Flush chain, then delete.
#   if [ "$TEST" = 1 ]; then
#     echo "Chain to be removed: $CHAIN"
#   else
#     sudo iptables -F $CHAIN || echo "Chain Name: $CHAIN" && exit 1 # flush
#     CHAIN_SPEC=$(sudo iptables -S INPUT | xargs php -r "preg_match('/(?!-A)INPUT.*$CHAIN.*/', $1, $matches); echo $matches[0];")
#     sudo iptables -D $CHAIN_SPEC
#     sudo iptables -X $CHAIN # delete
#     echo "Chain Removed: $CHAIN"
#   fi
fi

# Save iptables rules.
if [ 1 != "$TEST" ]; then
  sudo service netfilter-persistent save
fi