#!/usr/bin/env bash

# !!!using only IPv4 addresses and interfaces!!!

set -o pipefail

[[ "$1" != "nocolor" ]] && source colors

skip_local_dns=0
query_timeout=10

# default DNS over HTTPS resolvers. DNSCrypt is always used if running
#DEF_USE_DOH="cloudflare google quad9 alidns-doh"
DEF_USE_DOH="cloudflare google quad9"

# name servers for queries
#dns_srv=("8.8.8.8" "1.1.1.1" "9.9.9.9" "114.114.114.114")
dns_srv=("8.8.8.8" "1.1.1.1" "9.9.9.9")

# other servers to cache
add_srv=() # (download.hiveos.farm)

# repo list
HIVEREPO_LIST=/etc/apt/sources.list.d/hiverepo.list

# DNSCrypt resolver IP
dc_ip="127.0.2.1"

# get current hosts
oldhosts="$(< /etc/hosts)"$'\n'

newhosts="# DO NOT EDIT. All changes will be overwritten by cache-hive-ip
127.0.0.1 localhost
# The following lines are desirable for IPv6 capable hosts
::1 localhost ip6-localhost ip6-loopback
ff02::1 ip6-allnodes
ff02::2 ip6-allrouters
"


# globals
declare -A HOSTS_CACHE
query_servers=
local_gws=
doh_srv=()


function dns_servers {
	if [[ $skip_local_dns -ne 1 ]]; then
		# get system dns servers
		local dns
		readarray -t dns < <(grep -oP "^nameserver\s+\K[^\s]+" /run/systemd/resolve/resolv.conf)
		# merge all servers
		dns+=("${dns_srv[@]}")
	else
		local dns=("${dns_srv[@]}")
	fi
	query_servers=
	local dns_servers=()
	local addr
	# remove duplicates, empty and localhost
	for addr in "${dns[@]}"
	do
		[[ $addr =~ ^127\.0\.[0-9]{1,3}\.[0-9]{1,3}$ || ! $addr =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ || " ${dns_servers[@]} " =~ " $addr " ]] &&
			continue
		dns_servers+=("$addr")
		query_servers+=" @$addr"
	done
	echo "${CYAN}> Using DNS servers: ${YELLOW}${dns_servers[@]}${NOCOLOR}"
}


function check_resolved { # @array_name
	local -n res="$1"
	local arr=()
	local addr
	for addr in "${res[@]}"; do
		[[ " ${arr[@]} " =~ " $addr " ]] && continue

		[[ ! "$addr" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]] &&
			echo "${RED}> Incorrect address format:${NOCOLOR} $addr" &&
			continue

		[[ "$addr" =~ ^127\. || "$addr" =~ ^192\.168\. || "$addr" =~ ^10\. || "$addr" =~ ^169\.254\. ]] &&
			echo "${RED}> Local/reserved address range:${NOCOLOR} $addr" &&
			continue

		[[ -z "$local_gws" ]] && readarray -t local_gws < <(ip route | grep "via" | awk '{print $3}')
		[[ " ${local_gws[@]} " =~ " $addr " ]] &&
			echo "${RED}> Gateway address:${NOCOLOR} $addr" &&
			continue

		arr+=("$addr")
	done

	[[ ${#arr[@]} -eq 0 ]] && return 1

	# replace input array
	res=("${arr[@]}")
	return 0
}


function resolve_DoH { # @output_array_name, @url, @mode [ cf | gg | dc ]
	local name="$1"
	local host="$2"
	local mode="$3"
	local -n arr="$name"
	local query=

	echo -n "${CYAN}> Resolving: ${WHITE}$host${NOCOLOR}"

	if [[ "$mode" == "dnscrypt" ]]; then
		echo -n " ${YELLOW}[DNSCrypt DoH]${NOCOLOR}"
		readarray -t arr < <(dig -4 +noall +answer +retry=0 +time=$query_timeout "$host" "@$dc_ip" | grep -oP ".*\sA\s\K[^\s]+$")
		if [[ ! -z "$arr" ]]; then
			echo ""
			check_resolved "$name" && return 0
		else
			echo "${RED} - FAILED${NOCOLOR}"
		fi
		return 1
	elif [[ "$mode" == "cloudflare" || "$mode" == "cf" ]]; then
		echo -n " ${YELLOW}[CloudFlare DoH]${NOCOLOR}"
		# https://developers.cloudflare.com/1.1.1.1/encryption/dns-over-https/make-api-requests/
                query=`curl -L --connect-timeout $query_timeout --max-time $query_timeout --silent --resolve "cloudflare-dns.com:443:1.1.1.1" -H "accept: application/dns-json" "https://1.1.1.1/dns-query?name=$host" 2>/dev/null`
	elif [[ "$mode" == "quad9" || "$mode" == "q9" ]]; then
		echo -n " ${YELLOW}[Quad9 DoH]${NOCOLOR}"
		query=`curl -L --connect-timeout $query_timeout --max-time $query_timeout --silent "https://9.9.9.9:5053/dns-query?type=A&name=$host" 2>/dev/null`
	elif [[ "$mode" == "google" || $mode == "gg" ]]; then
		echo -n " ${YELLOW}[Google DoH]${NOCOLOR}"
		query=`curl -L --connect-timeout $query_timeout --max-time $query_timeout --silent --resolve "dns.google:443:8.8.8.8" "https://dns.google/resolve?type=A&name=$host" 2>/dev/null`
	elif [[ "$mode" == "alidns-doh" || $mode == "ali" ]]; then
		echo -n " ${YELLOW}[AliDNS DoH]${NOCOLOR}"
		query=`curl -L --connect-timeout $query_timeout --max-time $query_timeout --silent --resolve "dns.alidns.com:443:223.5.5.5" "https://dns.alidns.com/resolve?type=A&name=$host" 2>/dev/null`
	else
		echo " ${YELLOW}[unknown]${RED} - FAILED${NOCOLOR}"
		return 2
	fi
	if [[ $? -eq 0 ]]; then
		#echo "$query" | jq  .
		jqs='if .Status != 0 then empty else .Answer[] | if .type != 1 then empty else .data end end'
		readarray -t arr < <(echo "$query" | jq -c -r "$jqs" 2>/dev/null)
		if [[ ! -z "${arr[@]}" ]]; then
			echo ""
			check_resolved "$name"
			return
		fi
	fi
	echo "${RED} - FAILED${NOCOLOR}"
	return 1
}


function resolve_host { # @output_array_name, @host
	local name="$1"
	local host="$2"
	local -n arr="$name"
	local doh

	for doh in "${doh_srv[@]}"; do
		resolve_DoH "$name" "$host" "$doh" && return 0
	done
	# prepare dns servers array for query
	[[ -z "$query_servers" ]] && dns_servers
	# using dig to query multiple dns servers at once
	for try in {1..2}; do
		echo -n "${CYAN}> Resolving: ${WHITE}$host${NOCOLOR}"
		readarray -t arr < <(dig -4 +noall +answer +retry=0 +time=$query_timeout $query_servers "$host" | grep -oP ".*\sA\s\K[^\s]+$")
		if [[ ! -z "$arr" ]]; then
			echo ""
			check_resolved "$name" && return 0
		else
			echo "${RED} - FAILED${NOCOLOR}"
			break
		fi
		# trying again without local dns
		[[ $skip_local_dns -eq 1 ]] && break
		skip_local_dns=1
		dns_servers
	done
	return 1
}


function add_host {
	local host="$1"
	local forced_host="$2"
	local resolved=()

	[[ "$host" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]] &&
		echo "${YELLOW}> Skipping host:${NOCOLOR} $host" &&
		return 0

	resolve_host "resolved" "$host" || return

	[[ ! -z "$forced_host" ]] && host="$forced_host"

	for ip in "${resolved[@]}"; do
		[[ " ${HOSTS_CACHE["$host"]} " =~ " $ip " ]] && echo "${YELLOW}> Already in cache:${NOCOLOR} $ip $host" && continue
		HOSTS_CACHE["$host"]+=" $ip"
		echo "${GREEN}> Adding to cache:${NOCOLOR} $ip $host"
		newhosts+="$ip $host"$'\n'
	done

	return 0
}


function update_hosts {
	[[ ! -e $RIG_CONF ]] && echo "> No $RIG_CONF" && return 1
	USE_DOH=
	DOH_ENABLED=
	source $RIG_CONF
	[[ -z "$API_HOST_FILE" ]] && source /etc/environment
	[[ -z "$HIVE_HOST_URL" && -e "$API_HOST_FILE" ]] && source $API_HOST_FILE # fallback api host
	[[ -z "$HIVE_HOST_URL" ]] && echo "> HIVE_HOST_URL is empty" && return 2
	local host=`echo $HIVE_HOST_URL | awk -F'://' '{print $2}'`
	[[ -z "$host" ]] && echo "> Unable to parse host name: $HIVE_HOST_URL" && return 3

	# add rig host
	# not using hostname-check here as it directly writes to /etc/hosts
	[[ -z "$WORKER_NAME" ]] && hostname="$(hostname 2>/dev/null)" || hostname="${WORKER_NAME:-worker}" # fallback hostname
	[[ "$hostname" =~ ^([0-9]+).* ]] && hostname="hive$WORKER_NAME"
	newhosts+=$'\n'"# WORKER NAME"$'\n'"127.0.0.1 $hostname"$'\n'

	newhosts+=$'\n'"# API CACHE BEGIN"$'\n'

	# check dnscrypt
	#[[ -f /etc/resolv.conf && " $( < /etc/resolv.conf ) " =~ nameserver[[:blank:]]*$dc_ip[[:space:]] ]] &&
	[[ `dig 127.0.0.1 +time=1 2>/dev/null | grep -oP "SERVER: \K[^(]+"` =~ "$dc_ip#" ]] &&
		DC_ENABLED=1 ||
		DC_ENABLED=0

	# use defaults
	[[ -z "$USE_DOH" ]] && USE_DOH="$DEF_USE_DOH"

	# do not use DoH if DOH_ENABLED = 0 or empty
	[[ "$DOH_ENABLED" -eq 0 ]] && USE_DOH=

	# disable dnscrypt if DOH_ENABLED != 2 and not empty
	[[ "$DOH_ENABLED" -ne 2 && ! -z "$DOH_ENABLED" && $DC_ENABLED -eq 1 ]] &&
		DC_ENABLED=0 &&
		dnscrypt --quiet --disable

	# install/enable dnscrypt if DOH_ENABLED = 2
	if [[ "$DOH_ENABLED" -eq 2 && $DC_ENABLED -ne 1 ]]; then
		[[ -t 1 ]] && REDIR=/dev/tty || REDIR=/dev/null
		payload="$( dnscrypt --quiet --enable | tee $REDIR )"
		exitcode=$?
		if [[ $exitcode -ge 10 ]]; then
			payload="$( dnscrypt --quiet --install | tee $REDIR )"
			exitcode=$?
		fi
		if [[ $exitcode -ne 0 ]]; then
			echo "$payload" | message error "DNSCrypt error ($exitcode)" payload
			# disable temporarily to prevent further attemps
			echo "DOH_ENABLED=1" >> $RIG_CONF
		else
			DC_ENABLED=1
		fi
		sleep 3 # warming up
	fi

	# always use dnscrypt resolver if it is enabled
	[[ ! " $USE_DOH " =~ ' dnscrypt ' && $DC_ENABLED -eq 1 ]] && USE_DOH="dnscrypt $USE_DOH"

	read -a doh_srv < <( echo "$USE_DOH" )

	# Hive API host
	add_srv+=("$host")
	# Hive shell host
	[[ "$HSSH_SRV" =~ ^([0-9a-z\.-]+)(:[0-9]+)?$ ]] && add_srv+=("${BASH_REMATCH[1]}")
	# Repo hosts
	[[ -f $HIVEREPO_LIST ]] &&
		host=`grep -oP "deb\s*(http://|https://|)\K[0-9a-z\.-]+" $HIVEREPO_LIST` &&
		add_srv+=($host)
	# API URL hosts: server-provided $API_HOST_URLs + baked-in failover pool,
	# so /etc/hosts pre-caching covers fallback hosts too (needed under DNS blocking)
	local API_HOSTS_DEFAULT_FILE=/hive/etc/api_hosts
	for url in $API_HOST_URLs $( [[ -f $API_HOSTS_DEFAULT_FILE ]] && grep -vE '^\s*(#|$)' $API_HOSTS_DEFAULT_FILE ); do
		[[ ! "$url" =~ ^(https?://)?([0-9a-z\.-]+)(:[0-9]+)?$ ]] && continue
		host="${BASH_REMATCH[2]}"
		[[ ! " ${add_srv[@]} " =~ " $host " ]] && add_srv+=("$host")
	done

	local code=0
	# add all hosts
	for host in "${add_srv[@]}"
	do
		add_host "$host" || code=$?
	done

	return $code
}

exitcode=0

if [[ "$1" != "resetonly" ]]; then
	#[[ $(networkctl list | grep -c routable) -eq 0 ]] && echo "> No connection to network" && exit 10
	update_hosts
	exitcode=$?
fi

# update only if needed
if [[ "$oldhosts" != "$newhosts" ]]; then
	echo "$newhosts" > /etc/hosts
	echo "${GREEN}> Hosts update done${NOCOLOR}"
	sync
else
	echo "${GREEN}> Hosts is up to date${NOCOLOR}"
fi

exit $exitcode
