unifi_dns.sh

#!/bin/bash

PACKAGE=$(basename "$0")

#  .----. _ .-.          _  _
#  `--. ::_;: :.-.      :_;:_;
#    ,','.-.: `'.' .--. .-..-.
#  .'.'_ : :: . `.' '_.': :: :
#  :____;:_;:_;:_;`.__.': ::_;
#                     .-. :
#                     `._.'

# CONFIG
SITE_ID="default" # site id if it differs
UNIFI_DATA_DIR="" # the path to your config/data dir (without the trailing slash). I use docker and the script is on the host, so my path is ./config/data
USG_SSH="user@ubnt" # the host/alias you use to connect to your USG over SSH. Use SSH key auth and hopefully have it setup in your SSH config
UNIFI_CONTROLLER_BASEURL="https://unifi:8443" # should be self explanatory
UNIFI_CONTROLLER_CURL="curl --tlsv1 --silent --cookie /tmp/unifi_cookie --cookie-jar /tmp/unifi_cookie --insecure "

# PROVISION CONFIG (CREATE A NEW USER FOR THIS)
# If you plan on using the provision command fill this out. Should be a valid account that can provision on the Unifi controller
UNIFI_USERNAME=""
UNIFI_PASSWORD=""
USG_MAC="" # MAC address of the USG

# FORMATTING
RESET=$(tput sgr0)
BOLD=$(tput bold)

MISSING_DEP="no"
if ! command -v jq &> /dev/null; then
  echo "${BOLD}$PACKAGE:${RESET} jq is not installed."
  MISSING_DEP="yes"
fi

if ! command -v curl &> /dev/null; then
  echo "${BOLD}$PACKAGE:${RESET} curl is not installed."
  MISSING_DEP="yes"
fi

if ! command -v ssh &> /dev/null; then
  echo "${BOLD}$PACKAGE:${RESET} ssh is not installed."
  MISSING_DEP="yes"
fi

if [[ "$MISSING_DEP" == "yes" ]]; then
  exit 1
fi

display_usage() {
  echo "$PACKAGE"
  echo ""
  echo "Description:"
  echo "  Use jq, curl and SSH to view & manage the local DNS mappings."
  echo ""
  echo "Usage:"
  echo "  $PACKAGE show                                  Shows entries in local config."
  echo "  $PACKAGE show-remote                           Shows entries in config on the USG."
  echo "  $PACKAGE leases                                Shows lease hostnames on the USG."
  echo "  $PACKAGE set <host> [target] [--alias=<name>]  Sets or removes your host and target in the config."
  echo "  $PACKAGE provision                             Forces the USG to reprovision so the new hostnames reflect."
  echo "  $PACKAGE -h | --help"
  echo "  $PACKAGE --version"
  echo ""
  echo "Options:"
  echo "  -h --help                Show this screen."
  echo "  --version                Show version."
  echo "  -a<name> --alias=<name>  Alias to set with your record."
}

valid_ip() {
    local ip=$1
    local stat=1

    if [[ $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        OIFS=$IFS
        IFS='.'
        ip=($ip)
        IFS=$OIFS
        [[ ${ip[0]} -le 255 && ${ip[1]} -le 255 \
            && ${ip[2]} -le 255 && ${ip[3]} -le 255 ]]
        stat=$?
    fi
    return $stat
}

TEMP=$(getopt -o 'ha:' --long 'help,version,alias:' -n "${BOLD}$PACKAGE${RESET}" -- "$@")

if [ $? -ne 0 ]; then
  display_usage
  exit 1
fi

eval set -- "$TEMP"
unset TEMP

ALIAS=""
while true; do
  case "$1" in
    '-h'|'--help')
      display_usage
      exit 0
    ;;
    '--version')
      echo "$PACKAGE v1.0.0"
      exit 0
    ;;
    '-a'|'--alias')
      ALIAS="$2"
      shift 2
      continue
    ;;
    '--')
      shift
      break
    ;;
    *)
      echo 'Internal error!' >&2
      exit 1
    ;;
  esac
done

sub__show() {
  jq -r "[\"${BOLD}FQDN${RESET}\", \"${BOLD}Alias${RESET}\", \"${BOLD}IP${RESET}\"], (.system[\"static-host-mapping\"][\"host-name\"] | to_entries[] | [.key + \"$RESET\", (if .value.alias then .value.alias | join(\", \") else \"\" end + \"$RESET\"), (.value.inet | join(\", \") + \"$RESET\")]) | @tsv" "$UNIFI_DATA_DIR/sites/$SITE_ID/config.gateway.json" | column -t
}

sub__show-remote() {
  echo $(ssh -q "$USG_SSH" mca-ctrl -t dump-cfg | jq -r '.system["static-host-mapping"]["host-name"]') | jq -r "[\"${BOLD}FQDN${RESET}\", \"${BOLD}Alias${RESET}\", \"${BOLD}IP${RESET}\"], (. | to_entries[] | [.key + \"$RESET\", (if .value.alias then .value.alias | join(\", \") else \"\" end + \"$RESET\"), (.value.inet | join(\", \") + \"$RESET\")]) | @tsv" | column -t
}

sub__leases() {
  local result=$(ssh -q "$USG_SSH" cat /var/run/dnsmasq-dhcp.leases)
  echo "$result" | awk -v B="$BOLD" -v R="$RESET" 'BEGIN{print B"Hostname" R,B"IP"R,B"MAC"R}; {print $4R,$3R,$2R}' | column -t
}

host_exists() {
  echo $(jq -r --arg host "$1" '.system["static-host-mapping"]["host-name"] | has($host) | if . then 1 else 0 end' "$UNIFI_DATA_DIR/sites/$SITE_ID/config.gateway.json")
}

sub__set() {
  local host="$1"
  local target="$2"

  if [[ "$host" == "" ]]; then
    echo "${BOLD}Error:${RESET} You must provide the hostname."
    exit 1
  fi

  local exists=$(host_exists "$host")

  if [[ "$exists" == "1" ]] && [[ "$target" =~ ^(|null|nothing|empty|unset)$ ]]; then
    echo $(jq --arg host "$host" 'del(.system["static-host-mapping"]["host-name"][$host])' "$UNIFI_DATA_DIR/sites/$SITE_ID/config.gateway.json") > "$UNIFI_DATA_DIR/sites/$SITE_ID/config.gateway.json"
    echo "${BOLD}Success!${RESET} $host was removed from your config."
    exit 0
  fi

  if ! valid_ip "$target"; then
    echo "${BOLD}Error:${RESET} Target must be a valid IP."
    exit 1
  fi

  local json
  if [[ "$ALIAS" != "" ]]; then
    json=$(jq --arg host "$host" --arg alias "$ALIAS" --arg target "$target" '.system["static-host-mapping"]["host-name"][$host] = { alias: [$alias], inet: [$target] }' "$UNIFI_DATA_DIR/sites/$SITE_ID/config.gateway.json")
  else
    json=$(jq --arg host "$host" --arg target "$target" '.system["static-host-mapping"]["host-name"][$host] = { inet: [$target] }' "$UNIFI_DATA_DIR/sites/$SITE_ID/config.gateway.json")
  fi
  echo "$json" > "$UNIFI_DATA_DIR/sites/$SITE_ID/config.gateway.json"
  if [[ "$exists" == "1" ]]; then
    echo "${BOLD}Success!${RESET} $host has been updated in your config."
  else
    echo "${BOLD}Success!${RESET} $host has been added to your config."
  fi
}

usg_state() {
  echo $($UNIFI_CONTROLLER_CURL "$UNIFI_CONTROLLER_BASEURL/api/s/$SITE_ID/stat/device-basic" | jq -r '.data[] | select(.mac == "e0:63:da:cd:33:78") | if .state == 5 then "provisioning" elif .state == 1 then "ready" else "unknown" end')
}

spin()
{
  spinner="/|\\-/|\\-"
  while :
  do
    for i in `seq 0 7`
    do
      echo -n "${spinner:$i:1}"
      echo -en "\010"
      sleep 1
    done
  done
}

sub__provision() {
  # login
  local result=$($UNIFI_CONTROLLER_CURL --data '{"username": "'"$UNIFI_USERNAME"'", "password": "'"$UNIFI_PASSWORD"'"}' "$UNIFI_CONTROLLER_BASEURL/api/login" | jq -r '.meta.rc')
  if [[ "$result" =~ ^(|error)$ ]]; then
    echo "${BOLD}Error:${RESET} Unable to login. Check base URL and credentials and try again."
    exit 1
  fi

  # check if already provisoning
  local state=$(usg_state)
  if [[ "$state" == "provisioning" ]]; then
    echo "${BOLD}Error:${RESET} USG is already provisioning."
    exit 1
  fi

  # send provision command
  local result=$($UNIFI_CONTROLLER_CURL --data '{"mac": "'"$USG_MAC"'", "cmd": "force-provision"}' "$UNIFI_CONTROLLER_BASEURL/api/s/$SITE_ID/cmd/devmgr" | jq -r '.meta.rc')
  if [[ "$result" =~ ^(|error)$ ]]; then
    echo "${BOLD}Error:${RESET} Unable to send provision command."
    exit 1
  fi

  echo "${BOLD}Success!${RESET} Waiting on USG to finish provisioning."

  spin &
  SPIN_PID=$!
  trap "kill -9 $SPIN_PID 2>/dev/null" `seq 0 15`

  sleep 2
  local state=$(usg_state)
  until [[ "$state" =~ ^(ready|unknown)$ ]]; do
    sleep 2
    local state=$(usg_state)
  done

  kill -9 $SPIN_PID
  echo

  # logout
  $UNIFI_CONTROLLER_CURL --data '{}' "$UNIFI_CONTROLLER_BASEURL/api/logout" >/dev/null
  echo "${BOLD}Success!${RESET} USG is provisioned."
}

cmdname=$1; shift
if type "sub__$cmdname" >/dev/null 2>&1; then
  "sub__$cmdname" "$@"
else
  if [[ "$cmdname" == "" ]]; then
    sub__show
  else
    echo "${BOLD}$PACKAGE:${RESET} \"$cmdname\" is not a valid action. Use \"$PACKAGE help\" to see valid actions."
  fi
fi