#!/bin/sh -
#
# Shuffle nameservers listed based on the query times, if enabled in
# resolv.conf. Bare in mind the special syntax.
#
# Rick van der Zwet <rick@wirelessleiden.nl>
#

# ``Random'' sleep to avoid ``slamming'' on the DNS door
verbose=true
if [ "$1" = "cron" ]; then
  verbose=false
  sleep `expr $$ % 30`
fi

TAG='^# START DYNAMIC LIST - updated by /tools/nameserver-shuffle$'
$verbose && echo "# Searching in /etc/resolv.conf for tag '$TAG'"

TDIR=`mktemp -d -t $(basename $0)`
# Cleanup before going home
trap "rm -Rf $TDIR; exit 1" 1 2 3 15
trap "rm -Rf $TDIR; exit 0" 0

DYNLIST=$TDIR/dynlist
RESULTLIST=$TDIR/resultlist
NEWRESOLV=$TDIR/resolv.conf.new

# Get enabled DYNAMIC LIST nameservers
sed "1,\+$TAG+d" /etc/resolv.conf > $DYNLIST || exit 1
NAMESERVERS=`awk '/^nameserver/ {print $2}' $DYNLIST | sort -u` || exit 1

# Only do something if we have dynamic nameservers
$verbose && echo "# Processing" `echo $NAMESERVERS | wc -w` nameservers
if [ -n "$NAMESERVERS" ]; then
  # Find query times
  for NAMESERVER in $NAMESERVERS; do 
    $verbose && printf "## Testing nameserver %-16s query time: " $NAMESERVER
    # Strict checking to avoid buggy links to return decent results.
    QUERY_TIME=`timeout 0.5 drill SOA wleiden.net @$NAMESERVER 2>/dev/null | awk '/Query time:/ {print $4}'`
    # Failed to complete succesfully
    [ -z "$QUERY_TIME" ] && QUERY_TIME="9999"
    $verbose && echo "$QUERY_TIME"
    echo "$QUERY_TIME $NAMESERVER" >> $RESULTLIST
  done
  
  # Get the header part 
  sed -n "1,\+$TAG+p" /etc/resolv.conf > $NEWRESOLV || exit 1
  
  # Output sorted list
  NAMESERVERS=`sort -n $RESULTLIST | uniq -f 1 | awk '{print $2}'`
  for NAMESERVER in $NAMESERVERS; do
    QUERY_TIME=`grep -m 1 "${NAMESERVER}$" $RESULTLIST | cut -d' ' -f1`
    [ "$QUERY_TIME" = "9999" ] && STATUS="Query time: down" || STATUS="Query time: $QUERY_TIME"
    # awk magic to get maximum length of comment field (for display purposes).
    ML=`awk '/^nameserver/ {l=length($4);if (l>ml){ml=l}}END{print ml}' $DYNLIST`
    # awk magic allows adding or updating of status of nameserver
    awk '/^nameserver[[:blank:]]+'"$NAMESERVER"'[[:blank:]]+/ {printf "nameserver %-16s # %-'$ML's (%s)\n", $2, $4,"'"$STATUS"'"}' $DYNLIST >> $NEWRESOLV || exit 1
  done
  $verbose && echo "################################"
  $verbose && echo "## BEGIN new /etc/resolv.conf ##"
  $verbose && echo "################################"
  $verbose && cat $NEWRESOLV
  $verbose && echo "################################"
  $verbose && echo "## END new /etc/resolv.conf   ##"
  $verbose && echo "################################"
  cat $NEWRESOLV > /etc/resolv.conf || exit 1

  # Update unbound forwarders list
  unbound-control -q forward $(awk '/^nameserver 172/ {print $2}' /etc/resolv.conf | uniq | head -3) 1>&2
fi