#!/bin/bash

# Recursive/Caching Nameserver configuration
#
# Focused on supported datacenters / ip space
# ny components removed 2026 (and earlier)
#
# FIXME how many ns to keep?
# FIXME consider shellcheck improvements
# FIXME focused on ipv4 only right now, extend to ipv6?
#

# some default values
REMOVEBAD="N"
CHECK="Y"
ADDLOCAL="N"
DEBUG="N"
MAXGOODDNS="3"

# Dallas, Standard Network
nsc_dl1="72.249.191.254"
# Dallas, Premium Network
nsc_dl2="206.123.113.254"
# Reading
nsc_rdg1="185.17.252.125"
# Auckland
nsc_akl1="103.16.180.254"
# Sydney
nsc_syd1="43.239.97.254"


# blacklist here from resolver configs
# - level3/lumen name servers (4.2.2.x) occasionally redirect bad requests to
#   placeholder instead of returning a fail. Dont use on servers
# - rimuhosting ns1/2/3... are not intended to be resolvers
# - includes old/deprecated rimuhosting nsc- resolvers
#
# historically 149.112.112.112 used as a default "nameserver" which can
# be set up e.g. in errata but will be replaced when this script is run.
#
bad_nameservers=(
4.2.2.1
4.2.2.2
27.50.65.23
60.234.1.1
60.234.2.2
66.199.228.130
66.199.228.250
66.199.228.253
66.199.228.254
66.199.235.50
72.249.0.34
72.29.96.250
72.9.108.146
92.48.122.126
94.76.200.49
118.127.6.6
118.127.6.7
122.100.15.254
149.112.112.112
202.125.32.4
202.125.32.5
202.60.64.6
202.60.64.7
203.25.185.119
206.123.113.132
206.123.64.245
206.123.69.254
206.123.69.4
207.210.212.202
207.99.0.1
207.99.0.2
207.99.0.41
207.99.0.42
210.56.80.56
)

# selection of ideal ns servers
ln_nameservers="$nsc_rdg1 $nsc_dl1 $nsc_dl2"
fr_nameservers="$nsc_rdg1 $nsc_dl1 $nsc_dl2"
dl_nameservers="$nsc_dl1 $nsc_dl2 $nsc_rdg1"
au_nameservers="$nsc_syd1 $nsc_akl1 $nsc_dl1"
ak_nameservers="$nsc_akl1 $nsc_syd1 $nsc_dl1"

# well known public resolvers, which we can default to
gg_v4_ns="8.8.8.8 8.8.4.4"
gg_v6_ns="2001:4860:4860::8888 2001:4860:4860::8844"
cf_v4_ns="1.1.1.1 1.0.0.1"
cf_v6_ns="2606:4700:4700::1111 2606:4700:4700::1001"

# ip matches that are help us guess if one of the above ns sets above can be
# explicitly selected via cli
ln_ranges="185.17. 213.229. 217.112. 31.193. 85.234. 92.48. 94.76."
fr_ranges="185.17. 84.200."
dl_ranges="206.123. 207.210.2 65.99. 72.249. 72.29.10 74.50. 199.231. 169.6"
au_ranges="223.252. 43.239.9 103.52.11 117.120.6 117.20.6"
ak_ranges="103.248.17 202.37.12 49.50.24 103.6.21 103.16.18"


####################################################################
# Functions to modify nameserver information in each config system #
#                                                                  #
# These functions abstract accessors that were previously inline.  #
# Functions generally print results to stdout                      #
####################################################################

function msg_con () {
  [ "${DEBUG}" = 'N' ] && return
  echo "${1:-'Unknown Error'}" 1>&2
  return 0
}

# Get a list of nameservers, one per line
#
# Netplan.io and resolvconf have the concept of nameservers for an
# interface (or reachable through an interface).
function get_nameservers() {
  if [ "$RESOLVSYS" = "traditional" ] ; then
    awk '/^nameserver/ {print $2}' /etc/resolv.conf
  elif [ "$RESOLVSYS" = "netplan" ] ; then
    python3 -c 'import yaml;print(" ".join(yaml.load(open("/etc/netplan/rimu_nameservers.yaml").read())["network"]["ethernets"]["'${NIC_PRIMARY}'"]["nameservers"]["addresses"]))' 2> /dev/null || echo
  elif [ "$RESOLVSYS" = "systemd" ] ; then
    grep 'DNS=' /etc/systemd/resolved.conf | cut -d= -f2 | tr ' ' '\n'
  else
    echo "${FUNCNAME[0]}: Resolver configuration scheme unknown or not implimented"
    exit 1
  fi
  return 0
}

# Get the number of nameservers currently configured
function get_num_nameservers() {
  get_nameservers | wc --words
}

function is_locahost_nameservers_only() {
  local nameservers="$(get_nameservers)"
  [ "$nameservers" == "127.0.0.1" ] && return 0
  return 1 
}

# Create a config if nothing exists
function create_if_empty() {
  if [ "$RESOLVSYS" = "traditional" ] ; then
    if [ ! -e /etc/resolv.conf  ] ; then
      touch /etc/resolv.conf
    fi
  elif [ "$RESOLVSYS" = "netplan" ] ; then
    # use a placeholder address from the bad list at the top, this will get
    # overwritten later
    if [ ! -e /etc/netplan/rimu_nameservers.yaml  ] ; then
      cat << EOF > /etc/netplan/rimu_nameservers.yaml
network:
  version: 2
  renderer: networkd
  ethernets:
    $NIC_PRIMARY:
      nameservers:
        addresses:
          - 149.112.112.112
EOF
    fi
  elif [ "$RESOLVSYS" = "systemd" ] ; then
    touch /etc/systemd/resolved.conf
  else
    echo "${FUNCNAME[0]}: Resolver configuration scheme unknown or not implimented"
    exit 1
  fi
  return 0
}

# Backup the old config. uses shell expansion to copy 'a' to 'a$currentpid'
function backup_old() {
  if [ "$RESOLVSYS" = "traditional" ] ; then
    old="/etc/resolv.conf-$$"
    cp /etc/resolv.conf $old
  elif [ "$RESOLVSYS" = "netplan" ] ; then
    old="/etc/netplan/rimu_nameservers.yaml-$$"
    cp /etc/netplan/rimu_nameservers.yaml $old
  elif [ "$RESOLVSYS" = "systemd" ] ; then
    old="/etc/systemd/resolved.conf-$$"
    cp /etc/systemd/resolved.conf $old
  else
    echo "${FUNCNAME[0]}: Resolver configuration scheme unknown or not implimented"
    exit 1
  fi
  return 0
}

# remove a nameserver from the config, returning success if that was removed,
# or failure if that server wasn't there.
# Actually just comments out the line where possible
function remove_nameserver() {
  if [ "$RESOLVSYS" = "traditional" ] ; then
    if [ $(grep -c "^nameserver $1" /etc/resolv.conf) -gt 0 ]; then
      if ! sed s/"^nameserver $1"/"#nameserver $1"/g --in-place /etc/resolv.conf; then
        echo "! ${FUNCNAME[0]}: Unable to remove nameserver : $1"
        exit 1
      fi
    else
      return 1
    fi
  elif [ "$RESOLVSYS" = "netplan" ] ; then
    if grep -e "[[:space:]]${1}$" /etc/netplan/rimu_nameservers.yaml; then
      if ! sed -ir "/[[:space:]]${1}$/d" /etc/netplan/rimu_nameservers.yaml; then
        echo "! ${FUNCNAME[0]}: Unable to remove nameserver : $1"
        exit 1
      fi
    else
      return 1
    fi
  elif [ "$RESOLVSYS" = "systemd" ] ; then
    if grep -e "^DNS=.*${1}.*" /etc/systemd/resolved.conf; then
      if ! sed -ir "/[[:space:]]?${1}/d" /etc/systemd/resolved.conf; then
        echo "! ${FUNCNAME[0]}: Unable to remove nameserver : $1"
        exit 1
      fi
    else
      return 1
    fi
  else
    echo "! ${FUNCNAME[0]}: Resolver configuration scheme unknown or not implimented"
    exit 1
  fi
  echo "  Removed nameserver : $1"
  return 0
}

# add a nameserver, returning success if added, failure if it was already
# there. nameserver is added at the *end* of the current list
function add_nameserver() {
  if [ "$RESOLVSYS" = "traditional" ] ; then
    # name server already listed?
    if [ $(grep -c "^nameserver $1" /etc/resolv.conf) -gt 0 ]; then
      return 1
    fi
    # add name server
    echo "nameserver $1" >> /etc/resolv.conf
  elif [ "$RESOLVSYS" = "netplan" ] ; then
    # name server already listed?
    if grep -e "[[:space:]]${1}$" /etc/netplan/rimu_nameservers.yaml; then
      return 1
    fi
    # add name server
    echo "          - $1" >> /etc/netplan/rimu_nameservers.yaml
  elif [ "$RESOLVSYS" = "systemd" ] ; then
    # name server already listed?
    if grep -e "^DNS=.*${1}.*" /etc/systemd/resolved.conf; then
      return 1
    fi
    # add name server. for matching line, substitute end of line with the new value
    if [ $(grep '^DNS=' /etc/systemd/resolved.conf | wc --words) -gt 1 ]; then
      sed -i "/^DNS=.*$/ s/$/ ${1}/" /etc/systemd/resolved.conf
    else
      sed -i "/^DNS=.*$/ s/$/${1}/" /etc/systemd/resolved.conf
    fi
  else
    echo "${FUNCNAME[0]}: Resolver configuration scheme unknown or not implimented"
    exit 1
  fi
  return 0
}

# Show changes we made, apply the changes
function show_changes() {
  if [ "$RESOLVSYS" = "traditional" ] ; then
    diff -u $old /etc/resolv.conf || rm -f $old
  elif [ "$RESOLVSYS" = "netplan" ] ; then
    if diff -u $old /etc/netplan/rimu_nameservers.yaml; then
      rm -f $old
    else
      netplan apply
    fi
  elif [ "$RESOLVSYS" = "systemd" ] ; then
    if diff -u $old /etc/systemd/resolved.conf; then
      rm -f $old
    else
      systemctl restart systemd-resolved
    fi
  else
    echo "${FUNCNAME[0]}: Resolver configuration scheme unknown or not implimented"
    exit 1
  fi
  return 0
}


function usage() {
  echo "
$0 --help | [options]

  Update nameservers/resolvers used by server to support DNS lookups.

  - Must be run as root to allow access to system networking files
  - Currently only knows about ipv4 nameservers
  - prefers any already present working nameservers

  Options:
  --dc ln|fr|dl|au|ak  - optional, specify prefered location to select resolvers
                         by default tries to auto detect or fallback to using
                         public dns
  --nameservers <list> - optional, space seperated list of prefered resolvers
  --addlocal           - optional, add 127.0.0.1 as first nameserver, eg in
                         case there is a local caching resolver
  --[no]removebad      - optional, comment out 'bad' nameservers, prefix
                         with no to explicitly invert, default is to not remove
  --[no]check          - optional, check current name servers, prefix with no
                         to explicitly disable, default is to check
  --max-ns <num>       - optional, upper limit of name servers to set (default 3)
  --help               - returns this usage message, and exits

  Optional overrides (in case detection fails. CAUTION)
  --interface <name>   - override the primary interface name
  --ip <ip>            - specify the IP the primary interface has
  "
}
declare -xrf usage


#####################################################
# End of abstracted functions, start of main script #
#####################################################

if [ ${EUID} -ne "0" ]; then
  echo "! Must be root to allow access to networking system files"
  usage
  exit 1
else
  msg_con "* Running: ${0} $*"
fi

if ! command -v "dig" &> /dev/null; then
  echo "! 'dig' command required but not installed. Install and try again?" >&2
  exit 1
fi


# We know about a few systems for specifying nameservers, namely:
#
# * Traditional, nameservers stored in static /etc/resolv.conf
# * resolveconf package in debian/ubuntu/other distros: nameservers stored
#      in /etc/network/interfaces
# * systemd (eg /etc/systemd/resolved.conf
# * netplan.io, potentiallty in ubuntu 18.04 and later: nameservers stored in
#      /etc/netplan/rimunameservers.yaml
# * perhaps other systems such as low-level systemd.networkd
#
# We set RESOLVESYS to match the system in use
RESOLVSYS=traditional
if [ -e /etc/resolveconf ]; then
  RESOLVSYS=resolveconf
elif [ $(ls /etc/netplan/ 2>/dev/null|wc --words) -gt 0 ] ; then
  RESOLVSYS=netplan
elif [ -e /etc/systemd/resolved.conf ]; then
  # systemd is present, but it will honour a static resolv.conf, before using its own values
  if [ ! -f /etc/resolv.conf ]; then
    RESOLVSYS=systemd
  fi
fi

# autodetection, parsing things like...
# 2: enp2s0f0: <BROADCAST,MULTICAST,SLAVE,UP,LOWER_UP> mtu 1500 qdisc mq master bond0 state UP group default qlen 1000
# 3: enp2s0f1: <BROADCAST,MULTICAST,SLAVE,UP,LOWER_UP> mtu 1500 qdisc mq master bond0 state UP group default qlen 1000
# 4: bond0: <BROADCAST,MULTICAST,MASTER,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default qlen 1000
# 5: br2: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default qlen 1000
# 2: eth0: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc fq_codel state UP group default qlen 1000
NIC_PRIMARY=$(ip addr show | egrep -v '^[[:space:]]|lo|vif|master' | grep qlen | grep -m1 'state UP' | cut -d: -f2 | sed 's/ //g')
# FIXME add older method eg ifconfig to detect interface?

IP_PRIMARY="$(ip -4 addr show "${NIC_PRIMARY}" | grep 'inet' | head -n1 | awk '{print $2}' | cut -f1 -d'/')"
if [ -z "${IP_PRIMARY}" ]; then # try older command
  IP_PRIMARY="$(/sbin/ifconfig "${NIC_PRIMARY}" | grep 'inet ' | sed 's/inet addr:/inet /' | awk '{print $2}')"
fi
if [ -z "${IP_PRIMARY}" ]; then # best effort
  IP_PRIMARY="$(hostname -I | awk '{print $1}')"
fi

while [ -n "$1" ]; do
  case "$1" in
  --dc)
    [ $# -lt 1 ] && echo "Option --dc requires a value" >&2 && usage && exit 1
    shift
    eval good_nameservers=\$\{${1}_nameservers\}
    if [ -z "${good_nameservers}" ]; then # we didnt get a match
      echo "! Unable to determine nameserver set for --dc ${1}"
      usage
      exit 1
    fi
    msg_con "* Using name servers for location : ${1}"
    ;;
  --nameservers|--name-servers)
    [ $# -lt 1 ] && echo "! Option --nameserver requires a quoted space seperated list of ips" >&2 && usage && exit 1
    shift
    good_nameservers="${1}"
    ;;
  --addlocal|--add-local)
    ADDLOCAL="Y"
    ;;
  --removebad|--remove-bad)
    REMOVEBAD="Y"
    ;;
  --noremovebad|--no-removebad|--no-remove-bad|--noremove-bad)
    REMOVEBAD="N"
    ;;
  --nocheck|--nochecks|--no-check|--no-checks)
    CHECK="N"
    ;;
  --interface)
    [ $# -lt 1 ] && echo "Option --interface requires a value" >&2 && usage && exit 1
    shift
    NIC_PRIMARY="${1}"
    echo "! CAUTION: forcing primary interface to be known as ${NIC_PRIMARY}"
    ;;
  --ip)
    [ $# -lt 1 ] && echo "Option --ip requires a value" >&2 && usage && exit 1
    shift
    IP_PRIMARY="${1}"
    echo "! CAUTION: forcing primary IP to be known as ${IP_PRIMARY}"
    ;;
  --check|--checks)
    CHECK="Y"
    ;;
  --debug)
    DEBUG="Y"
    ;;
  --max-ns)
    [ $# -lt 1 ] && echo "Option --max-ns requires a value" >&2 && usage && exit 1
    shift
    MAXGOODDNS=$1
    ;;
  --help|--usage)
    usage
    exit 0
    ;;
  *)
    echo "! Unexpected argument $1" >&2
    usage
    exit 1
    ;;
  esac
  shift
done

if [ -z "${NIC_PRIMARY}" ]; then
  echo "! Unable to detect primary interface. Check network configuration. Use --interface to override manually"
  usage
  exit 1
fi
if [ -z "${IP_PRIMARY}" ]; then
  echo "! Unable to detect primary network address. Check network configuration. Use --ip to override manually"
  usage
  exit 1
fi
if ! ping -q -c 1 -W 1 8.8.8.8 >/dev/null; then
  echo "! Deeming the internet to not be accessible, please fix before re-running." >&2
  usage
  exit 1
fi

is_locahost_nameservers_only && echo "Only localhost nameservers, leaving the setup as is." && exit 0

# if --dc wasnt explicitly given, run through locations to check if
# primary ip matches a well known range and add that set
for dc in ln fr dl au ak; do
  [ -n "${good_nameservers}" ] && break
  eval ranges=\$\{${dc}_ranges\}
  for range in ${ranges}; do
    if [ $(echo "${IP_PRIMARY}" | grep -c ^${range}) -gt 0 ]; then
      eval good_nameservers=\$\{${dc}_nameservers\}
      msg_con "* Selecting nameservers from ip   : ${IP_PRIMARY}"
      msg_con "* Using nameservers for location  : ${dc}"
      break
    fi
  done
done

# safe defaults
if [ -z "${good_nameservers}" ]; then
  good_nameservers="${cf_v4_ns} ${gg_v4_ns}"
  msg_con "
  Defaulting to well known public name servers. Use --help to find how to
  set or select prefered nameservers instead.
  "
fi

if [ "${ADDLOCAL}" = "Y" ]; then
  good_nameservers="127.0.0.1 ${good_nameservers}"
  msg_con "* Adding localhost as first nameserver"
fi

echo "* Detected resolver configuration : ${RESOLVSYS}"
echo "* Recomended name servers are : ${good_nameservers}"

if [ "${RESOLVSYS}" = "resolveconf" ]; then
  msg_con "
  ! server appears to have the resolvconf package installed, that may
  override your changes. Set the recomended name servers manually or consider
  uninstalling resolveconf, then reruning this script.
  " >&2
  exit 1
fi

create_if_empty
# this file should at least be present, as a static, link, or generated file
if [ ! -e /etc/resolv.conf ]; then
  echo "! unable to find /etc/resolv.conf"
  exit 1
fi
# workaround so services can use resolve.conf (eg postfix)
chmod o+r /etc/resolv.conf

backup_old

nameservers=$(get_nameservers)

if [ "${REMOVEBAD}" == "Y" ]; then
  msg_con "  check all currently configured nameservers, add to the bad server list if any"
  for server in ${nameservers}; do
    if [ $(dig @${server} +short google.com | grep -v '^;' | wc -l) -lt 1 ]; then
      bad_nameservers+=($server)
      msg_con "! Currently configured nameserver is not responding, will disable/remove : ${server}"
    fi
  done

  msg_con "  actually remove (comment out) any bad nameservers"
  for ((i=0;i<${#bad_nameservers[@]};i++)); do
    ns=${bad_nameservers[$i]}
    remove_nameserver "${ns}"
  done
fi

for ns in $good_nameservers; do
  if [ $(get_num_nameservers ) -eq ${MAXGOODDNS} ]; then
    msg_con "! Done adding more name servers, we have enough"
    break
  fi

  # skip adding a 'good' nameserver if its not actually working
  [ $(dig @${ns} +short google.com | grep -v '^;' | wc -l) -lt 1 ] && continue;

  msg_con "  Add nameserver: ${ns}"
  add_nameserver "${ns}"
  # ret 1 if already present
done

# add some public dns name servers if there are not enough (as a fallback if
# they are all bad, may be dead code now)
for ns in ${cf_v4_ns}; do
  if [ $(get_num_nameservers ) -lt ${MAXGOODDNS} ]; then
    msg_con "  Not enough good name servers added, adding public DNS"
    add_nameserver "${ns}"
  fi
done

# show any changes. if none remove the old file
show_changes

if [ "N" = "$CHECK" ] ; then
  echo "* Skipping further checks as reuqested"
  exit 0
fi

msg_con "* Checking configured name servers:"
retcode=0
gooddns=0
nameservers=$(get_nameservers)
for server in ${nameservers}; do
  if [ $(dig @${server} +short google.com | grep -v '^;' | wc -l) -lt 1 ]; then
    echo "! Server '${server}' in your active configuration is not working.  Rerun this script with --removebad to remove that."
    ((retcode++))
    continue
  fi
  ((gooddns++))
  echo "- ${server} is working"
done

if [ $gooddns -lt 3 ] ; then
  echo "! Despite all efforts, only ${gooddns} nameservers are working in the current configuration. Please check manually (firewall?)."
fi
exit $retcode

#example from lenny, dont use 'host', is no longer consistent output across distros. 'dig' is better
#dcs:~# host google.com
#google.com              A       74.125.225.98
#google.com              A       74.125.225.97
#google.com              A       74.125.225.101
#google.com              A       74.125.225.110
#google.com              A       74.125.225.100
#google.com              A       74.125.225.105
#google.com              A       74.125.225.102
#google.com              A       74.125.225.103
#google.com              A       74.125.225.96
#google.com              A       74.125.225.99
#google.com              A       74.125.225.104

