#!/usr/bin/env bash
#
#     __               __
#    / /_  ____  _____/ /______
#   / __ \/ __ \/ ___/ __/ ___/
#  / / / / /_/ (__  ) /_(__  )
# /_/ /_/\____/____/\__/____/
#
# A program for managing host file entries.
#
# Based on Bash Boilerplate: https://github.com/alphabetum/bash-boilerplate
#
# Based on prior work by:
#
# - https://github.com/nddrylliog
#   - https://gist.github.com/nddrylliog/1368532
# - https://github.com/dfeyer
#   - https://gist.github.com/dfeyer/1369760
#
# Original idea and interface (since changed) via:
#
# https://github.com/macmade/host-manager
#
# Updates copyright (c) 2015 William Melody • hi@williammelody.com

###############################################################################
# Strict Mode
###############################################################################

set -o nounset
set -o errexit
set -o pipefail
IFS=$'\n\t'

###############################################################################
# Globals
###############################################################################

_VERSION="1.4.1"

# DEFAULT_COMMAND
#
# The command to be run by default, when no command name is specified. If the
# environment has an existing $DEFAULT_COMMAND set, then that value is used.
DEFAULT_COMMAND="${DEFAULT_COMMAND:-list}"

# HOSTS_PATH
#
# The path to the hosts file. This will almost always be /etc/hosts
HOSTS_PATH="${HOSTS_PATH:-/etc/hosts}"

###############################################################################
# Debug
###############################################################################

# _debug()
#
# Usage:
#   _debug printf "Debug info. Variable: %s\n" "$0"
#
# A simple function for executing a specified command if the `$_USE_DEBUG`
# variable has been set. The command is expected to print a message and
# should typically be either `echo`, `printf`, or `cat`.
_debug() {
  if [[ "${_use_debug:-"0"}" -eq 1 ]]
  then
    # Prefix debug message with "bug (U+1F41B)"
    printf "🐛  "
    "$@"
    printf "――――――――――――――――――――――――――――――――――――――――――――――――――――――――――――――――\n"
  fi
}
# debug()
#
# Usage:
#   debug "Debug info. Variable: $0"
#
# Print the specified message if the `$_USE_DEBUG` variable has been set.
#
# This is a shortcut for the _debug() function that simply echos the message.
debug() {
  _debug echo "$@"
}

###############################################################################
# Die
###############################################################################

# _die()
#
# Usage:
#   _die printf "Error message. Variable: %s\n" "$0"
#
# A simple function for exiting with an error after executing the specified
# command. The command is expected to print a message and should typically
# be either `echo`, `printf`, or `cat`.
_die() {
  # Prefix die message with "cross mark (U+274C)", often displayed as a red x.
  printf "❌  "
  "$@" 1>&2
  exit 1
}
# die()
#
# Usage:
#   die "Error message. Variable: $0"
#
# Exit with an error and print the specified message.
#
# This is a shortcut for the _die() function that simply echos the message.
die() {
  _die echo "$@"
}

###############################################################################
# Options
###############################################################################

# Get raw options for any commands that expect them.
raw_options="$*"

# Steps:
#
# 1. set expected short options in `optstring` at beginning of the "Normalize
#    Options" section,
# 2. parse options in while loop in the "Parse Options" section.

# Normalize Options ###########################################################

# Source:
#   https://github.com/e36freak/templates/blob/master/options

# The first loop, even though it uses 'optstring', will NOT check if an
# option that takes a required argument has the argument provided. That must
# be done within the second loop and case statement, yourself. Its purpose
# is solely to determine that -oARG is split into -o ARG, and not -o -A -R -G.

# Set short options -----------------------------------------------------------

# option string, for short options.
#
# Very much like getopts, expected short options should be appended to the
# string here. Any option followed by a ':' takes a required argument.
#
# In this example, `-x` and `-h` are regular short options, while `o` is
# assumed to have an argument and will be split if joined with the string,
# meaning `-oARG` would be split to `-o ARG`.
optstring=h

# Normalize -------------------------------------------------------------------

# iterate over options, breaking -ab into -a -b and --foo=bar into --foo bar
# also turns -- into --endopts to avoid issues with things like '-o-', the '-'
# should not indicate the end of options, but be an invalid option (or the
# argument to the option, such as wget -qO-)
unset options
# while the number of arguments is greater than 0
while (($#))
do
  case $1 in
    # if option is of type -ab
    -[!-]?*)
      # loop over each character starting with the second
      for ((i=1; i<${#1}; i++))
      do
        # extract 1 character from position 'i'
        c=${1:i:1}
        # add current char to options
        options+=("-$c")

        # if option takes a required argument, and it's not the last char
        # make the rest of the string its argument
        if [[ $optstring = *"$c:"* && ${1:i+1} ]]
        then
          options+=("${1:i+1}")
          break
        fi
      done
      ;;
    # if option is of type --foo=bar, split on first '='
    --?*=*)
      options+=("${1%%=*}" "${1#*=}")
      ;;
    # end of options, stop breaking them up
    --)
      options+=(--endopts)
      shift
      options+=("$@")
      break
      ;;
    # otherwise, nothing special
    *)
      options+=("$1")
      ;;
  esac

  shift
done
# set new positional parameters to altered options. Set default to blank.
set -- "${options[@]:-}"
unset options

# Parse Options ###############################################################

# Initialize command_argv array
#
# This array contains all of the arguments that get passed along to each
# command. This is essentially the same as the program arguments, minus those
# that have been filtered out in the program option parsing loop. This array
# is initialized with $0, which is the program's name.
command_argv=("$0")
# Initialize $cmd and $_use_debug, which can continue to be blank depending on
# what the program needs.
cmd=""
_use_debug=0

while [ $# -gt 0 ]
do
  opt="$1"
  shift
  case "$opt" in
    -h|--help)
      cmd="help"
      ;;
    --version)
      cmd="version"
      ;;
    --debug)
      _use_debug=1
      ;;
    *)
      # The first non-option argument is assumed to be the command name.
      # All subsequent arguments are added to $command_arguments.
      if [[ -n $cmd ]]
      then
        command_argv+=("$opt")
      else
        cmd="$opt"
      fi
      ;;
  esac
done

# Set $command_parameters to $command_argv, minus the initial element, $0. This
# provides an array that is equivalent to $* and $@ within each command
# function, though the array is zero-indexed, which could lead to confusion.
command_parameters=("${command_argv[@]:1}")

_debug printf "\$cmd: %s\n" "$cmd"
_debug printf "\$raw_options (one per line):\n%s\n" "$raw_options"
_debug printf "\$command_argv: %s\n" "${command_argv[*]}"
_debug printf "\$command_parameters: %s\n" "${command_parameters[*]:-}"

###############################################################################
# Environment
###############################################################################

# $_ME
#
# Set to the program's basename.
_ME=$(basename "$0")

_debug printf "\$_ME: %s\n" "$_ME"

###############################################################################
# Load Commands
###############################################################################

# Initialize defined_commands array.
defined_commands=()

# _load_commands()
#
# Loads all of the commands sourced in the environment.
#
# Usage:
#   _load_commands
_load_commands() {

  _debug printf "_load_commands(): entering...\n"
  _debug printf "_load_commands() declare -F:\n%s\n" "$(declare -F)"

  # declare is a bash built-in shell function that, when called with the '-F'
  # option, displays all of the functions with the format
  # `declare -f function_name`. These are then assigned as elements in the
  # $function_list array.
  local function_list=($(declare -F))

  for c in "${function_list[@]}"
  do
    # Each element has the format `declare -f function_name`, so set the name
    # to only the 'function_name' part of the string.
    local function_name
    function_name=$(printf "%s" "$c" | awk '{ print $3 }')

    _debug printf "_load_commands() \$function_name: %s\n" "$function_name"

    # Add the function name to the $defined_commands array unless it starts
    # with an underscore or is one of the desc(), debug(), or die() functions,
    # since these are treated as having 'private' visibility.
    if ! ( [[ "$function_name" =~ ^_(.*)  ]] || \
           [[ "$function_name" == "desc"  ]] || \
           [[ "$function_name" == "debug" ]] || \
           [[ "$function_name" == "die"   ]]
    )
    then
      defined_commands+=("$function_name")
    fi
  done

  _debug printf \
    "commands() \$defined_commands:\n%s\n" \
    "${defined_commands[*]:-}"
}

###############################################################################
# Main
###############################################################################

# _main()
#
# Usage:
#   _main
#
# The primary function for starting the program.
#
# NOTE: must be called at end of program after all commands have been defined.
_main() {
  _debug printf "main(): entering...\n"
  _debug printf "main() \$cmd (upon entering): %s\n" "$cmd"

  # If $cmd is blank, then set to help
  if [[ -z $cmd ]]
  then
    cmd="$DEFAULT_COMMAND"
  fi

  # Load all of the commands.
  _load_commands

  # If the command is defined, run it, otherwise return an error.
  if ( _contains "$cmd" "${defined_commands[*]:-}" )
  then
    # Pass all comment arguments to the program except for the first ($0).
    $cmd "${command_parameters[@]:-}"
  else
    _die printf "Unknown command: %s\n" "$cmd"
  fi
}

###############################################################################
# Utility Functions
###############################################################################

# _function_exists()
#
# Takes a potential function name as an argument and returns whether a function
# exists with that name.
_function_exists() {
  [ "$(type -t "$1")" == 'function' ]
}

# _contains()
#
# Takes an item and a list and determines whether the list contains the item.
#
# Usage:
#   _contains "$item" "${list[*]}"
_contains() {
  local test_list=(${*:2})
  for _test_element in "${test_list[@]:-}"
  do
    _debug printf "_contains() \$_test_element: %s\n" "$_test_element"
    if [[ "$_test_element" == "$1" ]]
    then
      _debug printf "_contains() match: %s\n" "$1"
      return 0
    fi
  done
  return 1
}

# _join()
#
# Takes a separator and a list of items, joining that list of items with the
# separator.
#
# Usage:
#   _join "," a b c
#   _join "${an_array[@]}"
_join() {
  local separator="$1"
  local target_array=(${@:2})
  local dirty
  dirty="$(printf "${separator}%s"  "${target_array[@]}")"
  local clean="${dirty:${#separator}}"
  printf "%s" "${clean}"
}

# _command_argv_includes()
#
# Takes a possible command argument and determines whether it is included in
# the command argument list.
#
# This is a shortcut for simple cases where a command wants to check for the
# presence of options quickly without parsing the options again.
#
# Usage:
#   _command_argv_includes "an_argument"
_command_argv_includes() {
  _contains "$1" "${command_argv[*]}"
}

# _verify_write_permissions
#
# Print a helpful error message when the specified operation can't be
# performed due to the lack of write permissions.
_verify_write_permissions() {
  if ! test -w "${HOSTS_PATH}"
  then
    _die printf \
"You don't have permission to perform this operation. Try again with:
sudo !!\n"
  fi
}

###############################################################################
# desc
###############################################################################

# desc()
#
# Usage:
#   desc command "description"
#
# Create a description for a specified command name. The command description
# text can be passed as the second argument or as standard input.
#
# To make the description text available to other functions, desc() assigns the
# text to a variable with the format $_desc_function_name
#
# NOTE:
#
# The `read` form of assignment is used for a balance of ease of
# implementation and simplicity. There is an alternative assignment form
# that could be used here:
#
# var="$(cat <<'EOM'
# some message
# EOM
# )
#
# However, this form appears to require trailing space after backslases to
# preserve newlines, which is unexpected. Using `read` simply requires
# escaping backslashes, which is more common.
desc() {
  set +e
  [[ -z $1 ]] && _die printf "desc: No command name specified.\n"
  if [[ -n ${2:-} ]]
  then
    read -d '' "_desc_$1" <<EOM
$2
EOM
    _debug printf "desc() set with argument: _desc_%s\n" "$1"
  else
    read -d '' "_desc_$1"
    _debug printf "desc() set with pipe: _desc_%s\n" "$1"
  fi
  set -e
}

# _print_desc()
#
# Usage:
#   _print_desc <command>
#
# Prints the description for a given command, provided the description has been
# set using the desc() function.
_print_desc() {
  local var="_desc_$1"
  if [[ -n ${!var:-} ]]
  then
    printf "%s\n" "${!var}"
  else
    printf "No additional information for \`%s\`\n" "$1"
  fi
}

###############################################################################
# Default Commands
###############################################################################

# Version #####################################################################

desc "version" <<EOM
Usage:
  $_ME ( version | --version )

Description:
  Display the current program version.

  To save you the trouble, the current version is $_VERSION
EOM
version() {
  printf "%s\n" "$_VERSION"
}

# Help ########################################################################

desc "help" <<EOM
Usage:
  $_ME help [<command>]

Description:
  Display help information for $_ME or a specified command.
EOM
help() {
  if [[ ${#command_argv[@]} = 1 ]]
  then
    cat <<EOM
    __               __
   / /_  ____  _____/ /______
  / __ \/ __ \/ ___/ __/ ___/
 / / / / /_/ (__  ) /_(__  )
/_/ /_/\____/____/\__/____/

A program for managing host file entries.

Version: $_VERSION

Usage:
  $_ME
  $_ME add <ip> <hostname> [comment]
  $_ME remove ( <ip> | <hostname> | <search string> ) [--force]
  $_ME list [enabled | disabled | <search string>]
  $_ME show ( <ip> | <hostname> | <search string> )
  $_ME disable ( <ip> | <hostname> | <search string> )
  $_ME disabled
  $_ME enable ( <ip> | <hostname> | <search string> )
  $_ME enabled
  $_ME edit
  $_ME file
  $_ME -h | --help
  $_ME --version

Options:
  -h --help  Display this help information.
  --version  Display version information.

Help:
  $_ME help [<command>]

$(commands)
EOM
  else
    _print_desc "$1"
  fi
}

# Command List ################################################################

desc "commands" <<EOM
Usage:
  $_ME commands [--raw]

Options:
  --raw  Display the command list without formatting.

Description:
  Display the list of available commands.
EOM
commands() {
  if _command_argv_includes "--raw"
  then
    printf "%s\n" "${defined_commands[@]}"
  else
    printf "Available commands:\n"
    printf "  %s\n" "${defined_commands[@]}"
  fi
}

###############################################################################
# Commands
# ========.....................................................................
#
# Example command group structure:
#
# desc example ""   - Optional. A short description for the command.
# example() { : }   - The command called by the user.
#
#
# desc example <<EOM
#   Usage:
#     $_ME example
#
#   Description:
#     Print "Hello, World!"
#
#     For usage formatting conventions see:
#     - http://docopt.org/
#     - http://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap12.html
# EOM
# example() {
#   printf "Hello, World!\n"
# }
#
###############################################################################

# ------------------------------------------------------------------------- add

desc "add" <<EOM
Usage:
  $_ME add <ip> <hostname> [comment]

Description:
  Add a given IP address and hostname pair, along with an optional comment.
EOM
add() {
  _verify_write_permissions
  local ip=${1:-}
  local hostname=${2:-}
  local comment=${*:3}
  if [[ -z ${ip} ]]
  then
    $_ME help add
    exit 1
  elif [[ -z ${hostname} ]]
  then
    printf "Please include a hostname\n"
    $_ME help add
    exit 1
  elif grep \
    -e "^${ip}\t${hostname}$" \
    -e "^${ip}\t${hostname}\t.*$" "${HOSTS_PATH}"
  then
    _die printf \
      "Duplicate address/host combination, %s unchanged.\n" \
      "${HOSTS_PATH}"
  else
    if [[ -n ${comment} ]]
    then
      local formatted_comment
      formatted_comment=$(_join " " "${comment[@]}")
      printf "%s\t%s\t# %s\n" \
        "${ip}" \
        "${hostname}" \
        "${formatted_comment}" >> "${HOSTS_PATH}"
    else
      printf "%s\t%s\n" \
        "${ip}" \
        "${hostname}" >> "${HOSTS_PATH}"
    fi
  fi
}

# --------------------------------------------------------------------- disable

desc "disable" <<EOM
Usage:
  $_ME disable ( <ip> | <hostname> | <search string> )

Description:
  Disable one or more records based on a given ip address, hostname, or
  search string.
EOM
disable() {
  _verify_write_permissions
  local search_term=$1
  if [[ -z "${search_term}" ]]
  then
    $_ME help disable
    exit 1
  else
    _debug printf "disable() \$search_term: %s\n" "$search_term"

    targets=$(sed -n "s/^\([^#]*${search_term}.*\)$/\1/p" "${HOSTS_PATH}")
    _debug printf "disable() \$targets: %s\n" "$targets"

    printf "Disabling:\n%s\n" "${targets}"

    # -i '' - in place edit. BSD sed requires extension argument, for GNU
    # it's optional. More info: http://stackoverflow.com/a/16746032
    sed -i '' "s/^\([^#]*${search_term}.*\)$/\#disabled: \1/g" "${HOSTS_PATH}"
  fi
}

# -------------------------------------------------------------------- disabled

desc "disabled" <<EOM
Usage:
  $_ME disabled

Description:
  List all disabled records. This is an alias for \`hosts list disabled\`.
EOM
disabled() {
  $_ME list disabled
}

# ------------------------------------------------------------------------ edit

desc "edit" <<EOM
Usage:
  $_ME edit

Description:
  Open the ${HOSTS_PATH} file in your \$EDITOR.
EOM
edit() {
  _verify_write_permissions
  if [[ -z "$EDITOR" ]]
  then
    _die printf "\$EDITOR not set.\n"
  else
    "$EDITOR" "${HOSTS_PATH}"
  fi
}

# ---------------------------------------------------------------------- enable

desc "enable" <<EOM
Usage:
  $_ME enable ( <ip> | <hostname> | <search string> )

Description:
  Enable one or more disabled records based on a given ip address, hostname,
  or search string.
EOM
enable() {
  _verify_write_permissions
  local search_term=$1
  if [[ -z "${search_term}" ]]
  then
    $_ME help enable
    exit 1
  else
    target_regex="s/^\#disabled: \(.*${search_term}.*\)$/\1/"

    targets=$(sed -n "${target_regex}p" "${HOSTS_PATH}")
    printf "Enabling:\n%s\n" "${targets}"

    # -i '' - in place edit. BSD sed requires extension argument, for GNU
    # it's optional. More info: http://stackoverflow.com/a/16746032
    sed -i '' "${target_regex}g" "${HOSTS_PATH}"
  fi
}

# --------------------------------------------------------------------- enabled

desc "enabled" <<EOM
Usage:
  $_ME enabled

Description:
  List all enabled records. This is an alias for \`hosts list enabled\`.
EOM
enabled() {
  $_ME list enabled
}

# ------------------------------------------------------------------------ file

desc "file" <<EOM
Usage:
  $_ME file

Description:
  Print the entire contents of the ${HOSTS_PATH} file.
EOM
file() {
  cat "${HOSTS_PATH}"
}

# ------------------------------------------------------------------------ list

desc "list" <<EOM
Usage:
  $_ME list [enabled | disabled | <search string>]

Description:
  List the existing IP / hostname pairs, optionally limited to a specified
  state. When provided with a seach string, all matching enabled records will
  be printed.
EOM
list() {
  # Get the disabled records up front for the two cases where they are needed.
  disabled_records=$(sed -n "s/^\#disabled: \(.*\)$/\1/p" "${HOSTS_PATH}")

  if [[ -n "$1" ]]
  then
    if [[ "$1" == disabled ]]
    then
      printf "%s\n" "${disabled_records}"
    elif [[ "$1" == enabled ]]
    then
      grep -v -e '^$' -e '^\s*\#' "${HOSTS_PATH}"
    else
      $_ME show "$1"
    fi
  else
    # NOTE: use separate expressions since using a | for the or results in
    # inconsistent behavior.
    grep -v -e '^$' -e '^\s*\#' "${HOSTS_PATH}"
    if [[ -n "${disabled_records}" ]]
    then
      printf "\nDisabled:\n%s\n" "${disabled_records}"
    fi
  fi
}

# ---------------------------------------------------------------------- remove

desc "remove" <<EOM
Usage:
  $_ME remove ( <ip> | <hostname> | <search string> ) [--force]

Options:
  --force  Skip the confirmation prompt.

Description:
  Remove one or more records based on a given IP address, hostname, or search
  string.
EOM
remove() {
  _verify_write_permissions
  local search_string=${1:-}
  if [[ -z $search_string ]]
  then
    $_ME help remove
    exit 1
  else
    local target_records
    target_records=$(sed -n "s/^\(.*${search_string}.*\)$/\1/p" "${HOSTS_PATH}")
    if [[ -z ${target_records:-} ]]
    then
      printf "No matching records found.\n"
      exit 1
    fi
    if ! _command_argv_includes "--force"
    then
      printf "Removing the following records:\n%s\n" "$target_records"
      while true
      do
        read -p "Are you sure you want to proceed? [y/n] " yn
        case $yn in
          [Yy]* )
            break
            ;;
          [Nn]* )
            printf "Exiting...\n"
            exit 0
            ;;
          * )
            printf "Please answer yes or no.\n"
            ;;
        esac
      done
    fi
    # -i '' - in place edit. BSD sed requires extension argument, for GNU
    # it's optional. More info: http://stackoverflow.com/a/16746032
    sed -i '' "/^.*${search_string}.*$/d" "${HOSTS_PATH}"
    printf "Removed:\n%s\n" "${target_records}"
  fi
}

# ------------------------------------------------------------------------ show

desc "show" <<EOM
Usage:
  $_ME show ( <ip> | <hostname> | <search string> )

Description:
  Print entries matching a given IP address, hostname, or search string.
EOM
show() {
  if [[ -n "$1" ]]
  then
    # Run `sed` before `grep` to avoid conflict that supress `sed` output.
    disabled_records=$(sed -n "s/^\#\(disabled: .*$1.*\)$/\1/p" "${HOSTS_PATH}")
    enabled_records=$(grep "^[^#]*$1" "${HOSTS_PATH}")
    # Output disabled records secondly for better organization.
    [[ -n "$enabled_records"  ]] && printf "%s\n" "$enabled_records"
    [[ -n "$disabled_records" ]] && printf "%s\n" "$disabled_records"
  else
    $_ME help show
    exit 1
  fi
}

###############################################################################
# Run Program
###############################################################################

# Calling the _main function after everything has been defined.
_main