#!/bin/bash
# -*- mode: sh; sh-basic-offset: 2; indent-tabs-mode: nil; -*-
# vim:set ft=sh et sw=2 ts=2:
#
# dig-mock - dig command replacement for testing
# Author: Scott Shambarger <devel@shambarger.net>
#
# Copyright (C) 2021 Scott Shambarger
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
VERSION=0.3.1

VERBOSE='' RREC='' HOSTS=() BRIEF='' SERVER='' HEADER=''

# default name state file
DIG_MOCK_NAMES=${DIG_MOCK_NAMES:-conf/dig-mock-names}

# used to trigger failure
DIG_MOCK_FAIL=${DIG_MOCK_FAIL-}

err() {
  local IFS; unset IFS
  printf >&2 "%s\n" "$*"
}

fail() {
  local -i rc=0
  [[ $1 != 0 ]] && { printf 2>/dev/null -v rc "%d" "$1" || rc=1; }
  shift
  [[ $1 ]] && { if [[ $rc == 0 ]]; then echo "$@"; else err "$@"; fi; }
  exit "$rc"
}

verbose() { [[ $VERBOSE ]] && err "$@"; }

usage() {
  echo "Usage: ${0##*/} [@server] [q-type] {q-opt} host"
  echo "Where: q-type  is one of (a,aaaa,ns,...)"
  echo "       q-opt is one of:"
  echo "             -q name (specify query name)"
  echo "             -t type (specify query type)"
  echo "             -<any>  (Ignored)"
  echo "             +short  (Display only short form of answer)"
  echo "             +<any>  (Ignored)"
  echo "       @server       (only match entries for \$SERVER)"
  echo "       -d            (show debug messages)"
  echo "       -h            (print help and exit)"
  echo "       -v            (print version and exit)"
  echo "       +fail[=rc]    (mock failure)"
  echo "All matches done against \"${DIG_MOCK_NAMES}\""
}

# inspired by https://stackoverflow.com/a/51573758/14179001
caseadj() { # <retvar> <string> <from> <to>
  local _r="${2-}" _i
  for ((_i=0; _i<${#3}; _i++)); do _r=${_r//${3:$_i:1}/${4:$_i:1}}; done
  printf -v "$1" "%s" "$_r"
}
lowercase() { # <retvar> <string>
  caseadj "$1" "${2-}" "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz"
}
uppercase() { # <retvar> <string>
  caseadj "$1" "${2-}" "abcdefghijklmnopqrstuvwxyz" "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
}

display_item() { # <fqdn> <ttl> <rrec> <value>
  local IFS; unset IFS
  printf "%-23s %-7d IN      %-7s %s\n" "${1%.}." "${@:2:3}"
}

display_header() { # answer
  [[ ${HEADER} || ${BRIEF} ]] && return 0
  HEADER=${1:-ANSWER}
  echo "; <<>> DiG-mock ${VERSION} <<>> ${ARGS}"
  echo ";; global options: +cmd"

  echo ";; QUESTION SECTION:"
  display_item ";${HOST}" "" "${RREC}"
  echo
  echo ";; ${HEADER} SECTION:"
}

display_result() { # <fqdn> <ttl> <rrec> <value>
  [[ ${BRIEF} ]] && { echo "$4"; return 0; }
  display_header
  display_item "$@"
}

display_nxdomain() {
  local host domain
  [[ ${BRIEF} ]] && return 0
  display_header "AUTHORITY"
  if [[ ${HOST} =~ [^.]+\..+$ ]]; then
    host=${HOST%.*.*}
  else
    host=${HOST%.*}
  fi
  domain=${HOST#${host}.}
  display_item "${domain}" 10800 SOA "a.root-servers.net."
}

match_query() { # <server> <fqdn> <ttl> <rrec> <value>
  local server=$1 fqdn=$2 rrec=$4

  verbose "matching" "$@"
  [[ ${SERVER} && ${server} ]] && {
    [[ ${SERVER} == "${server}" ]] || return 1; }
  shift
  [[ ${RREC} == "${rrec}" ]] || return 1
  [[ ${HOST} == "${fqdn}" ]] || return 1
  display_result "$@"
}

# Name state file format:
#   $SERVER <server>
#   $ORIGIN <domain>
#   $TTL <secs>
#   [ <name> ] <rrec> <value>

do_query() { # <filename>
  local file=$1 IFS='' line found='' ttl=3600 lineno=0
  local name='' rrec='' domain='' server='' value=''

  [[ -f ${file} ]] || fail 1 "State file $file not found"

  while read -r line || [[ ${line} ]]; do
    ((lineno++)) || :
    [[ ${line} =~ ^# ]] && continue
    unset IFS
    # shellcheck disable=SC2086
    set -- $line
    IFS=''
    if [[ ${3-} ]] && ! [[ ${line} =~ ^[[:space:]] ]]; then
      # <name> <rrec> <value>
      lowercase name "$1"
      uppercase rrec "$2"
      unset IFS; value=${*:3}; IFS=''
    elif [[ ${2-} ]]; then
      case $1 in
        \$ORIGIN) lowercase domain "${2%.}" ;;
        \$SERVER) lowercase server "@$2" ;;
        \$TTL) ttl=$2 ;;
        *)
          [[ ${name} ]] || fail 1 "State file ${file}#${lineno}: missing <name>"
          uppercase rrec "$1"
          unset IFS; value=${*:2}; IFS=''
          ;;
      esac
    elif [[ ${1-} ]]; then
      fail 1 "State file ${file}#${lineno}: invalid format"
    fi
    [[ ${name} && ${rrec} && ${value} ]] && {
      # always return lowercase AAAA
      [[ ${rrec} == "AAAA" ]] && lowercase value "${value}"
      match_query "${server}" "${name}${domain:+.${domain}}" \
                  "${ttl}" "${rrec}" "${value}" && found=1
      value=''
    }
  done < "${file}"
  [[ ${found} ]] || display_nxdomain
  return 0
}

is_rrec() { # <arg>
  local rrec
  uppercase rrec "$1"
  case ${rrec} in
    A|AAAA|NS|SOA|CNAME|MX|HINFO|MINFO|TXT|PTR) return 0 ;;
  esac
  return 1
}

parse_args() {
  local arg nextarg='' IFS; unset IFS
  ARGS="$*" # keep for header
  for arg in "$@"; do
    [[ ${arg} ]] || "dig: '' is not a legal name (unexpected end of input)"
    [[ ${nextarg} ]] && {
      case ${nextarg} in
        RREC) if is_rrec "${arg}"; then
                uppercase RREC "${arg}"
              else
                err ";; Warning, ignoring invalid type $arg"
              fi ;;
        HOST)
          lowercase arg "${arg%.}."
          HOSTS+=("${arg}") ;;
      esac
      nextarg=''
      continue
    }
    case ${arg} in
      -d) VERBOSE=1 ;;
      -h) usage; exit 0 ;;
      -q) nextarg=HOST ;;
      -t) nextarg=RREC ;;
      -v) err "${0##*/} ${VERSION}"; exit 0 ;;
      +short) BRIEF=1 ;;
      +fail=*) DIG_MOCK_FAIL=${arg#+fail=} ;;
      +fail) DIG_MOCK_FAIL=1 ;;
      +*|-*) : ;;
      @*) lowercase SERVER "${arg}" ;;
      *)
        if is_rrec "${arg}"; then
          uppercase RREC "${arg}"
        else
          lowercase arg "${arg%.}"
          HOSTS+=("${arg}")
        fi ;;
    esac
  done
  [[ ${#HOSTS[*]} != 0 || ${RREC} ]] || RREC=NS
  RREC=${RREC:-A}
  [[ ${#HOSTS[*]} != 0 ]] || HOSTS=(".")
  return 0
}

[[ ${MOCK_DEBUG} ]] && { err "${0##*/}" "$@"; VERBOSE=1; }

parse_args "$@"

[[ ${DIG_MOCK_FAIL-} || ${DIG_MOCK_OUTPUT-} ]] && {
  args='' arg=''; for arg in "$@"; do args+="${args:+ }'${arg}'"; done
  fail "${DIG_MOCK_FAIL:-0}" \
       "${DIG_MOCK_OUTPUT+${DIG_MOCK_OUTPUT/@ARGS@/${args}}}"
}

verbose "Using state file: ${file}"

for HOST in "${HOSTS[@]}"; do
  HEADER=''
  verbose "looking up ${RREC} record ${HOST}"
  do_query "${DIG_MOCK_NAMES}"
done
