#!/usr/bin/env bash
#
# buildhosts - Build and maintain /etc/hosts while including lists
#              imported from local and remote hosts list sources
#
# Instructions (as root):
#
# Enable Remote Sources
#   1. The first time you run the 'set' command, /etc/hosts will be
#      moved to /etc/hosts.core, and the file /etc/hosts.sources will
#      be created using the default sources list. The /etc/hosts file
#      will then be generated by combining /etc/hosts.core and the
#      resulting contents of downloaded sources in /etc/hosts.sources.
#   2. You can now open /etc/hosts.sources and remove or comment out
#      (by adding # in front) any sources you don't wish to use, as
#      well as add new ones. Sources can remote as well as local.
#   3. If you want to make any changes to the contents of what used to
#      be /etc/hosts, you should now make them to /etc/hosts.core, and
#      you'll need to rerun the 'set' command to have them applied.
#   4. Once you're done customizing your configuration, run the 'set'
#      command again and the /etc/hosts file will be regenerated using
#      the current list of sources and your up-to-date /etc/hosts.core
#      file.
#
# Disable Remote Sources
#   1. To disable the configured hosts lists, run the 'unset' command
#      and /etc/hosts.core will be moved back to /etc/hosts.
#   *  The /etc/hosts.sources file will continue to exist until it's
#      manually deleted, and it will be used again the next time remote
#      sources are enabled.
#
# Notes on /etc/hosts.sources
#   * On each line of /etc/hosts.sources, after # is commented/skipped.
#   * Colour output can be disabled by setting the variable NOCOL=1.
#   * The $HOSTS_SYSTEM variable can be set to an alternative location
#     to configure a different file to be used for /etc/hosts.
#   * Similarly, the $HOSTS_CORE and $HOSTS_SOURCES variables (which
#     default to $HOSTS_SYSTEM.core and $HOSTS_SYSTEM.sources and will
#     automatically change if $HOSTS_SYSTEM is changed), can be set to
#     alternative locations to configure the use of different files.
#
# by: Kevin MacMartin <prurigro@gmail.com>
# released under the MIT license
#

# User variables
[[ -z "$HOSTS_SYSTEM" ]] && HOSTS_SYSTEM="/etc/hosts"
[[ -z "$HOSTS_CORE" ]] && HOSTS_CORE="${HOSTS_SYSTEM}.core"
[[ -z "$HOSTS_SOURCES" ]] && HOSTS_SOURCES="${HOSTS_SYSTEM}.sources"

# Default list of sources (used to generate $HOSTS_SOURCES when it doesn't exist)
DEFAULT_SOURCES=(
    'http://winhelp2002.mvps.org/hosts.txt'
    'http://adaway.org/hosts.txt'
    'http://hosts-file.net/ad_servers.txt'
    'http://pgl.yoyo.org/adservers/serverlist.php?hostformat=hosts&showintro=0&mimetype=plaintext'
)

# Set the script name
BUILDHOSTS_SCRIPT="$(basename $0)"

# Configure colours
if [[ "$NOCOL" -eq 1 ]]; then
    HEADINGCOL="|"
else
    HEADINGCOL="\e[1;41m"
    SUCCESSCOL="\e[1;40m"
    FAILCOL="\e[1;40m"
    RESETCOL="\e[0m"
fi

# Help function: display help text
function buildhosts_help(){
    echo -e "${HEADINGCOL} ${BUILDHOSTS_SCRIPT}:${RESETCOL} build and maintain /etc/hosts with local and remote hosts list files\n${HEADINGCOL} ${RESETCOL}"
    echo -e "${HEADINGCOL} commands:${RESETCOL}"
    echo -e "${HEADINGCOL} ${RESETCOL}${SUCCESSCOL} ${BUILDHOSTS_SCRIPT} set:${RESETCOL} add hosts lists to /etc/hosts"
    echo -e "${HEADINGCOL} ${RESETCOL}${SUCCESSCOL} ${BUILDHOSTS_SCRIPT} unset:${RESETCOL} remove hosts lists from /etc/hosts"
    echo -e "${HEADINGCOL} ${RESETCOL}${SUCCESSCOL} ${BUILDHOSTS_SCRIPT} help:${RESETCOL} display this help"
    [[ "$1" -eq 1 ]] && exit 0
}

# Error function: display error then exit unsuccessfully
function buildhosts_error() {
    echo -e "${HEADINGCOL} ${RESETCOL}${FAILCOL} ! ERROR:${RESETCOL}${FAILCOL} ${1}${RESETCOL}"
    [[ "$2" -eq 1 ]] && echo -e "${HEADINGCOL} ${RESETCOL}" && buildhosts_help
    exit 1
}

# The set function to generate or regenerate /etc/hosts using /etc/hosts.core and the list of sources
function buildhosts_set() {
    # If $HOSTS_SOURCES doesn't exist, generate one using the default list of sources
    if [[ ! -f "$HOSTS_SOURCES" ]]; then
        echo -e "${HEADINGCOL} Generating Default Sources:${RESETCOL} ${HOSTS_SOURCES}"
        echo '# file:///etc/hosts.d/localhostlist.txt' >> "$HOSTS_SOURCES"
        echo '# http://domain.com/remotehostlist.txt' >> "$HOSTS_SOURCES"
        for source in "${DEFAULT_SOURCES[@]}"; do
            echo "$source" >> "$HOSTS_SOURCES"
        done
    fi

    # Fail if $HOSTS_SOURCES contains no sources after trimming comments
    [[ -n $(sed 's|^\ *#.*$||' "$HOSTS_SOURCES" | tr -d "\n") ]] || buildhosts_error "No sources available in ${HOSTS_SOURCES}"

    # If $HOSTS_CORE doesn't exist and $HOSTS_SYSTEM does, move $HOSTS_SYSTEM to $HOSTS_CORE
    if [ -f "$HOSTS_SYSTEM" -a ! -f "$HOSTS_CORE" ]; then
        echo -e "${HEADINGCOL} Moving:${RESETCOL} ${HOSTS_SYSTEM} to ${HOSTS_CORE}"
        cp "$HOSTS_SYSTEM" "$HOSTS_CORE"
    fi

    # Generate the hosts list using the URLs in the $HOSTS_SOURCES
    unset TMPHOSTS
    while read -r source; do
        if [[ -n "$source" ]]; then
            echo -e "${HEADINGCOL} Downloading:${RESETCOL} ${source}"

            # Download the the current source into $SRCDATA and fail if the result is empty
            SRCDATA=$(curl -C - -s "$source" | tr -d "\r")
            [[ -n "$SRCDATA" ]] || buildhosts_error "Could not download list @ ${source}"

            # If this isn't the first source, add a newline at the top
            [[ -n "$TMPHOSTS" ]] && SRCDATA=$(echo -e "\n${SRCDATA}")

            # Strip comments, then add the source to the end of $TMPHOSTS
            TMPHOSTS=${TMPHOSTS}$(sed 's|\ *#.*$||' <<< "$SRCDATA")
        fi
    done < <(sed 's|^\ *#.*$||' "$HOSTS_SOURCES")

    # Sort and remove duplicates, change 0.0.0.0 to 127.0.0.1, and remove non-127.0.0.1 redirections
    TMPHOSTS=$(sort -u < <(sed 's|\t| |;s|^\ *[0-9]*\.[0-9]*\.[0-9]*\.[0-9]*|127\.0\.0\.1|;
                                s|^127\.0\.0\.1\ \ *localhost\ *$||;
                                s|^\ *[^1].*$||' <<< "$TMPHOSTS"))

    # Add the system hosts file to the hosts list and warn if $HOSTS_CORE is missing
    echo -e "${HEADINGCOL} Writing:${RESETCOL} ${HOSTS_CORE} and $(wc -l <<< "$TMPHOSTS") host source entries to ${HOSTS_SYSTEM}"
    [[ -f "$HOSTS_CORE" ]] \
        && TMPHOSTS=$(cat "$HOSTS_CORE")$(echo -e "\n\n# Generated Host List (${BUILDHOSTS_SCRIPT})")${TMPHOSTS} \
        || echo -e "${HEADINGCOL} ${RESETCOL}${FAILCOL} Warning:${RESETCOL} Core hosts file ${HOSTS_CORE} does not exist"

    # Write $TMPHOSTS to $HOSTS_SYSTEM if it's not empty
    [[ -n "$TMPHOSTS" ]] && echo -e "$TMPHOSTS" > "$HOSTS_SYSTEM"
    if [[ -n $(cat "$HOSTS_SYSTEM") ]]; then
        echo -e "${HEADINGCOL} ${RESETCOL}${SUCCESSCOL} Done!${RESETCOL}"
    else
        echo -e "${HEADINGCOL} ${RESETCOL}${FAILCOL} Failed...${RESETCOL} (unseting to core hosts file)"
        cp "$MAINHOSTS" "$HOSTS_SYSTEM"
        buildhosts_error "The generated hosts file @ ${HOSTS_SYSTEM} could not be created"
        exit 1
    fi
}

function buildhosts_unset() {
    # Fail if $HOSTS_CORE doesn't exist
    [[ -f "$HOSTS_CORE" ]] || buildhosts_error "The core file ${HOSTS_CORE} does not exist, cannot unset"

    # Fail if $HOSTS_CORE is empty when all comments are removed
    [[ -n $(sed 's|^\ *#.*$||' "$HOSTS_CORE" | tr -d "\n") ]] || buildhosts_error "The core file ${HOSTS_CORE} contains no information, cannot unset"

    # Move $HOSTS_CORE to $HOSTS_SYSTEM
    echo -e "${HEADINGCOL} Moving:${RESETCOL} ${HOSTS_CORE} to ${HOSTS_SYSTEM}"
    mv "$HOSTS_CORE" "$HOSTS_SYSTEM" \
        && echo -e "${HEADINGCOL} ${RESETCOL}${SUCCESSCOL} Done!${RESETCOL}" \
        || buildhosts_error "Could not move ${HOSTS_CORE} to ${HOSTS_SYSTEM}"
}

# Fail if the user isn't root
[[ $EUID -eq 0 ]] || buildhosts_error "You must run ${BUILDHOSTS_SCRIPT} as root"

# Fail if run with no commands
[[ -n "$1" ]] || buildhosts_error "You must run ${BUILDHOSTS_SCRIPT} with a valid argument" 1

case "$1" in
    set)
        buildhosts_set
        ;;
    unset)
        buildhosts_unset
        ;;
    help)
        buildhosts_help 1
        ;;
    *)
        buildhosts_error "Invalid command" 1
        ;;
esac