#!/usr/bin/env bash
. colors

VERS="0.2.0"
POOLS=()
#for param in "$@"; do
#	POOLS+=("$param")
#done

print_help(){
	echo -e "${CYAN}Stratum-ping helper v$VERS"
	echo -e "${WHITE}This tool tests pool latency. Based on 2miners stratum-ping (https://github.com/2miners/stratum-ping)"
	echo
	echo -e "${YELLOW}Usage:"
	echo -e "\t${CYAN}pool-ping                ${LGRAY}- print this help ${NOCOLOR}"
	echo -e "\t${CYAN}pool-ping -h             ${LGRAY}- print this help ${NOCOLOR}"
	echo -e "\t${CYAN}pool-ping \"pool list\"    ${LGRAY}- test latancy of given pools${NOCOLOR}"
	echo -e "\t${CYAN}pool-ping -s \"pool list\" ${LGRAY}- sort given pool by latency${NOCOLOR}"
	echo -e "\t${CYAN}pool-ping -j \"pool list\" ${LGRAY}- print pools latancy in JSON format${NOCOLOR}"
	}

wallet=
password=
need_sort=0
print_json=0
stratum_type="stratum1"
while [ -n "$1" ]; do
	case "$1" in
		-u)
			wallet="-u $2 "
			shift
		;;
		-p)
			password="-p $2 "
			shift
		;;
		-s)
			need_sort=1
		;;
		-j)
			need_sort=1
			print_json=1
		;;
		-t)
			stratum_type=$2
			shift
		;;
		-h)
			print_help
			exit 0
		;;
		*)
			POOLS+=("$1")
		;;
	esac
	shift
done

#if [[ ${#POOLS} -eq 0 ]]; then
	# POOLS+=("eu-eth.hiveon.net:4444" "eu-eth.hiveon.net:14444" "ssl://eu-eth.hiveon.net:24443")
	# POOLS+=("ru-eth.hiveon.net:4444" "ru-eth.hiveon.net:14444" "ssl://ru-eth.hiveon.net:24443")
	# POOLS+=("naw-eth.hiveon.net:4444" "naw-eth.hiveon.net:14444" "ssl://naw-eth.hiveon.net:24443")
	# POOLS+=("aspac1-eth.hiveon.net:4444" "aspac1-eth.hiveon.net:14444" "ssl://aspac1-eth.hiveon.net:24443")
	# POOLS+=("eu1-etc.ethermine.org:4444" "ssl://eu1-etc.ethermine.org:5555")
#fi

CYCLES=3
#[[ $need_sort -eq 1 ]] && CYCLES=1

GOODTIME=100
BADTIME=200
TIMEOUT=5

CASE_WIDTH=56

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

OK=$GREEN"[OK]"$NOCOLOR

[[ $need_sort -eq 1 ]] && FAIL=999 || FAIL=$RED"[FAIL]"$NOCOLOR

unset GREP_OPTIONS

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

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=`sed 's/\x1b\[[0-9;]*m//g' <<< "$str"`
	local len=$(( $pad + ${#str} - ${#wocolors} ))
	printf "%-${len}b" "$str"
}


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


function TimeColor {
	if [[ $1 -le $GOODTIME ]]; then
		echo "${WHITE}$1 ms${NOCOLOR}"
	elif [[ $1 -gt $BADTIME ]]; then
		echo "${RED}$1 ms${NOCOLOR}"
	else
		echo "${YELLOW}$1 ms${NOCOLOR}"
	fi
}


# 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 $TIMEOUT $host 2>&1`
		RESULT=$?
		if [[ $RESULT -eq 0 ]]; then
			local totaltime=`echo "$response" | awk -F '/' 'END {print $5}'`
			totaltime=`bc <<< "scale=0;($totaltime+0.99)/1"`
			total=`TimeColor $totaltime`
			[[ $i -eq 0 ]] && echo -en "\t$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 "\t${YELLOW}[???]\t$FAIL\t"
			else
				echo -en "$FAIL\t"
			fi
			((RESULT++))
		fi
	done
	echo
	return $RESULT
}


function Stratum {
	local server="$1"
	local host=
	local RESULT=0
	local response=
	local tls=

	local Print="PrintCase2"
	local PrintTime="TimeColor"

	if [[ ${server,,} =~ (stratum\+tcp://|stratum\+ssl://|^)([^:]+):([0-9]+) ]]; then
		host="${BASH_REMATCH[2]}:${BASH_REMATCH[3]}"
		stratum=${BASH_REMATCH[1]}
		if [[ $stratum =~ ssl* ]]; then
			tls=1
			$Print "STRATUM+SSL" "$stratum$host"
		else
			$Print "STRATUM" "$stratum$host"
			stratum=
		fi
	else
		echo -e "$FAIL\t${ERROR_COLOR}Bad host specified$NOCOLOR"
		return 1
	fi

	for((i=0; i<${CYCLES:-1}; i++)); do
		response=`stratum-ping ${wallet}${password}-c 1 $stratum$host`
		RESULT=$?
		local error=`echo "$response" | jq '.[0].error'`
		if [[ $error == 'null' ]]; then
			totaltime=`jq '.[0].max_time' <<< $response`
			total=`$PrintTime $totaltime`
			[[ $i -eq 0 ]] && echo -en "\t$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 "\t${YELLOW}[???]\t$FAIL\t"
			else
				echo -en "$FAIL\t"
			fi
			((RESULT++))
		fi
	done
	echo
	return $RESULT
}


function Stratum_avg {
	local stratum=$1
	local host=$2
	local port=$3
	local RESULT=0
	local response=
	local tls=

	[[ ! $stratum =~ ssl* ]] && stratum=

	echo -en ${stratum}${host}:${port}

	#echo "stratum-ping ${wallet}${password}-c 3 ${tls:+-tls} $host" >> stratum.log
	response=`stratum-ping ${wallet}${password}-c $CYCLES ${stratum}${host}:${port}`
	RESULT=$?
	echo "$response" >> response.txt
	#min/avg/max = 27.349962ms, 28.86982ms, 29.773557ms
	local avg_time=`echo "$response" | grep 'min/avg/max' | awk '{print $4}' | tr -d 'ms,'`
	if [[ ! -z $avg_time ]]; then
		echo -en " $avg_time"
	else
		echo -en " $FAIL"
		((RESULT++))
	fi
	echo
	return $RESULT
}

# nslookup wrapper
function NSLookup {
	local host=$1
	local RESULT=0
	local str="${WHITE}>>> $host"
	color_printf $CASE_WIDTH "$str$NOCOLOR"
	#PrintCase2 "DNS" "$host"
	res=(`dig -4 +noall +answer +retry=0 +time=15 $host | grep -oP ".*\sA\s\K[^\s]+$"`)
	RESULT=$?
	[[ $RESULT -eq  0 ]] && echo -e "\t$OK\t${CYAN}${res[@]}$NOCOLOR" || echo -e "\t$FAIL\t${ERROR_COLOR}Check DNS or connection settings$DNS"$NOCOLOR
	return $RESULT
}

function QuickStratumPing() {
#	time=$FAIL
#	for pool in ${POOLS[@]}; do
#		if [[ ${pool,,} =~ (stratum\+tcp://|stratum\+ssl://|^)([^:]+):([0-9]+) ]]; then
#			stratum=${BASH_REMATCH[1]}
#			host=${BASH_REMATCH[2]}
#			port=${BASH_REMATCH[3]}

#			[[ "$last_host" == "$host" && $last_time != 'null' ]] && echo "${stratum}${host}:${port} $time" && continue
#			response=`Stratum_avg "$stratum" "$host" "$port"`
#			result=$?
#			time=`echo "$response" | awk '{print $2}'`
#			last_host=$host
#			last_time=$time
#			echo "$response"
#		else
#			echo "$pool $FAIL"
#			host=
#		fi
#	done

stratum-ping ${POOLS[@]} | jq ['.[] | {url: .url, time: .avg_time}']
}

if [[ ${#POOLS} -eq 0 ]]; then
	print_help
elif [[ $need_sort -eq 1 ]]; then
	res=`QuickStratumPing`
	if [[ $print_json -eq 0 ]]; then
		echo "$res" | jq -r 'sort_by(.time) | .[].url'
	else
		echo "$res"
	fi
else
	last_host=
	next=0
	for pool in ${POOLS[@]}; do
		host=`grep -oP "(ssl://|^)\K[^:]+" <<< "$pool"`
		if [[ $last_host != $host ]]; then
			next=0
			last_host=$host
			# skip ip hosts
			if [[ ! $host =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
				NSLookup "$host"
				[[ $? -ne 0 ]] && next=1 && continue
			else
				echo "${WHITE}>>> $host$NOCOLOR"
			fi
			Ping "$host"
		elif [[ $next -eq 1 ]]; then
			continue
		fi

		Stratum "$pool"
	done
fi

exit
