#!/usr/bin/env bash
. colors

[[ -f $RIG_CONF ]] && source $RIG_CONF
[[ $DISABLE_CERT_CHECK -ne 0 ]] && insecure="--insecure" || insecure=


ORIG_HIVE_HOST_URL="$HIVE_HOST_URL"
[[ -z "$API_HOST_FILE" ]] && source /etc/environment
[[ -f "$API_HOST_FILE" ]] && source $API_HOST_FILE # fallback api host

# just to check connection
#OTHER_DNS=(8.8.8.8 1.1.1.1 9.9.9.9 114.114.114.114)
OTHER_DNS=(8.8.8.8 1.1.1.1 9.9.9.9)

# default hosts on ports 80 and 443
API_HOSTS=(api.hiveos.farm helsinki.hiveos.farm msk.hiveos.farm paris.hiveos.farm amster.hiveos.farm naw.hiveos.farm api2msk.hiveos.farm)
# urls for other ports
API_URLS=(https://api.hiveos.farm:8443 http://helsinki.hiveos.farm:8000)
# Repo list
HIVEREPO_LIST=/etc/apt/sources.list.d/hiverepo.list


[[ $1 == "-a" || $1 == "--advanced" ]] &&
	ADV=1 && CYCLES=3

GOODTIME=1000

CASE_WIDTH=48

HEADER_COLOR=$YELLOW
CASE_COLOR=$BCYAN
WARN_COLOR=$LGREEN
ERROR_COLOR=$BRED

OK=$GREEN"[OK]"$NOCOLOR
FAIL=$RED"[FAIL]"$NOCOLOR
WARN=$YELLOW"[WARN]"$NOCOLOR


unset GREP_OPTIONS

# exit handler
trap 'echo -e $NOCOLOR; exit 1' 2

dnscrypt -s >/dev/null && DOH=1 || DOH=0

LOCAL_IP=`hostname -I | awk '{print $1}'`
LOCAL_GW=`ip route | awk '/default/ { print $3 }'`

#LOCAL_DNS=`nslookup server | grep Server | cut -d ":" -f 2 | sed -e 's/^[ \t]*//'`
LOCAL_DNS=(`grep ^nameserver /run/systemd/resolve/resolv.conf | awk '{print $2}'`)
query_servers=

function merge_all_dns {
	ALL_DNS=()
	local DNS=("${LOCAL_DNS[@]}")
	# merge all servers
	DNS+=("${OTHER_DNS[@]}")
	# remove duplicates, empty and localhost
	for ip in "${DNS[@]}"
	do
		[[ $ip =~ ^127\.0\.[0-9]{1,3}\.[0-9]{1,3}$ || ! $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ || " ${ALL_DNS[@]} " =~ " $ip " ]] &&
			continue
		ALL_DNS+=($ip)
		query_servers+=" @$ip"
	done
}


function PrintCase {
	printf "$CASE_COLOR%-${CASE_WIDTH}s" "$1"
}


color_printf() {
	local pad=$1
	local str="$2"
	# colored string length correction for printf formatting
	local wocolors=`echo "$str" | sed 's/\x1b\[[0-9;]*m//g'`
	local len=$(( $pad + ${#str} - ${#wocolors} ))
	printf "%-${len}b" "$str"
}


function PrintCase2 {
	local str="[$1] ${CASE_COLOR}$2"
	color_printf $CASE_WIDTH "$str"
}


function TimeColor {
	[[ $1 -le $GOODTIME ]] && echo "${WHITE}$1 ms${NOCOLOR}" || echo "${RED}$1 ms${NOCOLOR}"
}


# ping wrapper
function Ping {
	local host=$1
	local name=$2
	local errormsg=$3
	[[ -z $errormsg ]] && errormsg="Check firewall rules or connection settings"
	local RESULT=0
	local response=
	[[ -z $name ]] && name=$host

	PrintCase2 "PING" "$name"
	for((i=0; i<${CYCLES:-1}; i++))
	do
		response=`ping -i 0 -q -c 1 -w 4 $host 2>&1`
		RESULT=$?
		if [[ $RESULT -eq 0 ]]; then
			local totaltime=`echo "$response" | awk -F '/' 'END {print $5}'`
			totaltime=`echo "scale=0;($totaltime+0.99)/1" | bc`
			total=`TimeColor $totaltime`
			[[ $i -eq 0 ]] && echo -en "$OK\t$total\t" || echo -en "$total\t"
			[[ $totaltime -ge $((GOODTIME*2)) ]] && break
		else
			if [[ ${CYCLES:-1} -eq 1 ]]; then
				echo -en "$FAIL\t$ERROR_COLOR$errormsg$NOCOLOR"
			elif [[ $i -eq 0 && $CYCLES -gt 1 ]]; then
				echo -en $YELLOW"[???]\t$FAIL\t"
			else
				echo -en "$FAIL\t"
			fi
		fi
	done
	echo
	return $RESULT
}


# curl wrapper
function Curl {
	local host=$1
	local checkApi=$2
	local RESULT=0
	local response=

	if [[ $checkApi != "api" ]]; then
		PrintCase2 "WEB" "$host"
	elif [[ "$HIVE_HOST_URL/" =~ $host/ ]]; then
		PrintCase2 "API" "$host ${WHITE}(CURRENT)"
	elif [[ "$ORIG_HIVE_HOST_URL/" =~ $host/ ]]; then
		PrintCase2 "API" "$host ${BYELLOW}(DEFAULT)"
	else
		PrintCase2 "API" "$host"
	fi

	for((i=0; i<${CYCLES:-1}; i++))
	do
		if [[ $checkApi != "api" ]]; then
			response=`curl $insecure --connect-timeout 7 --max-time 15 --silent $host -w "\n%{time_total}"`
		else
			response=`curl $insecure --connect-timeout 7 --max-time 15 --silent -w "\n%{time_total}"\
				-H "Content-Type: application/json" \
				-X POST -d '{"method":"stats", "params": {"rig_id":"-1", "passwd": "1"}}' \
				${host}/worker/api`
		fi
		if [[ $? -eq 0 ]] && [[ $checkApi != "api" || ! -z $(echo "$response" | jq -c "if .error.code then . else empty end" 2>/dev/null) ]]; then
			local totaltime=`echo "$response" | tail -1`
			totaltime=`echo "scale=0;$totaltime*1000/1" | bc`
			total=`TimeColor $totaltime`
			[[ $i -eq 0 ]] && echo -en "$OK\t$total\t" || echo -en "$total\t"
			[[ $totaltime -ge $((GOODTIME*2)) ]] && break
		else
			if [[ ${CYCLES:-1} -eq 1 ]]; then
				echo -en "$FAIL\t${ERROR_COLOR}Check firewall rules$NOCOLOR"
			elif [[ $i -eq 0 && $CYCLES -gt 1 ]]; then
				echo -en $YELLOW"[???]\t$FAIL\t"
			else
				echo -en "$FAIL\t"
			fi
			((RESULT++))
		fi
	done
	echo
	return $RESULT
}


# nslookup wrapper
function NSLookup {
	local host="$1"
	local api_host="$2"
	local RESULT=0
	if [[ -z $api_host ]]; then
		PrintCase2 "DNS" "$host"
	else
		local str="${WHITE}>>> $host"
		[[ "$host" == "$api_host" ]] && str+=" (CURRENT)"
		color_printf $CASE_WIDTH "$str$NOCOLOR"
	fi
	[[ "$host" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]] && echo "" && return $RESULT
	#`nslookup $1 > /dev/null 2>&1`
	res=(`dig -4 +noall +answer +retry=0 +time=15 "$host" | grep -oP ".*\sA\s\K[^\s]+$"`)
	RESULT=$?
	# resolve using query DNS servers (without DoH only)
	[[ $RESULT -eq 0 ]] && echo -e "$OK\t${CYAN}${res[@]}$NOCOLOR" && return $RESULT
	if [[ $DOH_ENABLED -ne 1 && $DOH -eq 0 && ! -z "$query_servers" ]]; then
		res=(`dig -4 +noall +answer +retry=0 +time=15 $query_servers "$host" | grep -oP ".*\sA\s\K[^\s]+$"`)
		RESULT=$?
	fi
	[[ $RESULT -eq 0 ]] && echo -e "$WARN\t${BYELLOW}${res[@]}$NOCOLOR" || echo -e "$FAIL\t${ERROR_COLOR}Check DNS or connection settings$DNS"$NOCOLOR
	return $RESULT
}


# traceroute wrapper
function Trace {
	local host="$1"
	echo -en $WHITE
	traceroute --resolve-hostnames --max-hop=32 "$host"
	echo -en $NOCOLOR
}


# check host with dns resolving, ip, http, https cases
function Connection {
	local host="$1"
	local checkApi=$2
	local api_host="$3"
	local RESULT=0

	NSLookup "$host" "$api_host" || return 4
	Ping "$host" || ((RESULT++))
	Curl "http://$host" $checkApi || ((RESULT++))
	Curl "https://$host" $checkApi || ((RESULT++))

	return $RESULT
}


# connections test cases
function Connections {
	local RESULT=0

	echo -e "${YELLOW}> Checking network connection:${NOCOLOR}"
	Ping "$LOCAL_GW" "$LOCAL_GW ${NOCOLOR}(Gateway)" "Check your router and physical connection"

	for host in "${ALL_DNS[@]}"
	do
		[[ "$host" == "$LOCAL_GW" ]] && continue;
		Ping "$host" "$host ${NOCOLOR}(DNS)"
	done

	#[[ ! -z $HIVE_HOST_URL ]] && API_HOST=`echo $HIVE_HOST_URL | awk -F'://' '{print $2}'` || API_HOST=
	local API_HOST=`echo "$HIVE_HOST_URL" | grep -oP "://\K[^:^/]+"`
	[[ -z "$API_HOST" ]] && echo -e $ERROR_COLOR"No API server defined in rig.conf"$NOCOLOR && exit 1

	local HOSTS=("$API_HOST")
	if [[ "$API_HOST" =~ hiveos.farm$ ]]; then
		echo -e "${YELLOW}> Checking connection to HIVE OS server:${NOCOLOR}"
		Connection "hiveos.farm" "" " "
		[[ $? -eq 4 && $DOH -ne 1 ]] && echo -e "$(PrintCase "")	${BRED}Or try to enable DoH service: ${WHITE}dnscrypt -i${NOCOLOR}"
		# add default hosts
		for host in "${API_HOSTS[@]}"; do
			[[ ! " ${HOSTS[@]} " =~ " $host " ]] && HOSTS+=("$host")
		done
	fi

	# Repo urls
	if [[ -f $HIVEREPO_LIST ]]; then
		repo=`grep -oP "deb\s*\K[^\s]+" $HIVEREPO_LIST`
		if [[ ! -z "$repo" ]]; then
			echo -e "${YELLOW}> Checking connection to repository server:${NOCOLOR}"
			for url in $repo; do
				Curl "$url"
			done
		fi
	fi

	echo -e "${YELLOW}> Checking connection to worker API servers:${NOCOLOR}"
	# API URL hosts: server-provided $API_HOST_URLs + baked-in failover pool
	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]}"
		[[ ! " ${HOSTS[@]} " =~ " $host " ]] && HOSTS+=("$host")
	done

	for host in "${HOSTS[@]}"
	do
		Connection "$host" api "$API_HOST"
		exitcode=$?
		# return exicode only for current API server
		[[ "$host" == "$API_HOST" ]] && RESULT=$exitcode
		if [[ $exitcode -lt 4 && "${API_URLS[@]}" =~ /$host ]]; then
			for url in "${API_URLS[@]}"
			do
				[[ ! $url =~ /$host ]] && continue
				Curl $url api
			done
		fi
		#echo
	done

	if [[ $RESULT -eq 3 || $ADV -eq 1 ]]; then
		echo -e "${YELLOW}> Trace route to current API server:${NOCOLOR}"
		Trace "$API_HOST"
	fi

	return $RESULT
}


# restarting network interfaces
#function RestartNetworking {
#    if [ $LAST_RESULT -eq 1 ];
#    then
#    	echo -ne "$WARN_COLOR    connections was fail: restart network"
#    	LAST_RESULT=0
#    	sudo ifdown -a && sudo ifup -a &&
#    	sudo service network-manager restart && echo -e $OK || { LAST_RESULT=1; echo -e $FAIL; }
#        echo -ne "$NOCOLOR"
#    else
#    	exit $LAST_RESULT
#    fi
#}


# ip route | awk '/default/ { print $3 }'


# header
#echo -e "$HEADER_COLOR=== Hive Network Test ===$NOCOLOR"

echo -e "${YELLOW}> Local Network Configuration:${NOCOLOR}"

#echo -e "IP:        ${PURPLE}$LOCAL_IP${NOCOLOR}"
#echo -e "Gateway:   $LOCAL_GW"
#echo -e "DNS:       ${LOCAL_DNS[@]}"
networkctl status

if [[ $DOH -eq 1 ]]; then
	mode="${BGREEN}Enabled globally"
else
	case "$DOH_ENABLED" in
		2) mode="${BRED}Enabled, but not working";;
		1) mode="${BPURPLE}Enabled for system services";;
		0) mode="${BYELLOW}Disabled";;
		*) mode="Not enabled";;
	esac
fi
echo -e "DNS-over-HTTPS: $mode$NOCOLOR"

merge_all_dns

# test connections to hive hosts
Connections

# if connections was fail try restart network interfaces
#RestartNetworking
# test connections again
#Connections

# Traceroute to hosts for debuging if connections was not fixed
#if [ $LAST_RESULT -eq 1 ];
#then
#	sleep 5
#	echo
#    echo -e "$ERROR_COLOR network was not fixed"
#    echo
#	Trace "api.hiveos.farm       "
#	echo
#	Trace "amster.hiveos.farm    "
#else
#	exit $LAST_RESULT
#fi


exit
