#!/bin/sh
# Copyright (C) 2022 iopsys Software Solutions AB
# Author: AMIN Ben Romdhane <amin.benromdhane@iopsys.eu>

. /usr/share/libubox/jshn.sh
. /usr/share/bbfdm/scripts/bbf_api

# Shared across functions. serverselection_launch sets these from the ubus
# input (or a computed default); probe_host, forked once per host, and
# serverselection_reduce read them back.
hostlist=""
port=""
iface=""
ip_proto=""
nbr_of_rep=""
timeout=""
ping_timeout=""
protocol_used=""
proto=""
cancel=""
device=""
family=""
micros=""
TMPDIR=""

# Set by serverselection_reduce once every host has been probed, read back by
# serverselection_launch to build the result.
fasthost=""
min_time=0
avg_time=0
max_time=0
ip_addr_used=""
answered=0
resolve_failures=0

serverselection_list() {
	json_add_object "serverselection"
	json_add_string "hostlist" "str"
	json_add_string "port" "str"
	json_add_string "iface" "str"
	json_add_string "ip_proto" "str"
	json_add_string "nbr_of_rep" "str"
	json_add_string "timeout" "str"
	json_add_string "protocol_used" "str"
	json_add_string "proto" "str"
	json_add_string "cancel" "str"
	json_close_object
}

serverselection_error() {
	json_init
	json_add_string "Status" "$1"
	json_add_string "FastestHost" ""
	json_add_string "IPAddressUsed" ""
	json_add_int "MinimumResponseTime" "0"
	json_add_int "AverageResponseTime" "0"
	json_add_int "MaximumResponseTime" "0"
	json_dump

	# Store data in dmmap_diagnostics for both protocols (cwmp/usp)
	[ "$3" = "both_proto" ] && {
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.DiagnosticState="$1"
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.FastestHost=""
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.MinimumResponseTime=0
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.AverageResponseTime=0
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.MaximumResponseTime=0
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.IPAddressUsed=""
		$UCI_COMMIT_BBF_DMMAP
	}
}

# Probes one host and writes "status successes min avg max host address" to $2.
#
# The repetitions run one at a time on purpose: the data model requires a
# response, or a timeout, before the next probe is sent. A single 'ping -c N'
# would pace its sends on its own one second timer instead, and measurably
# slower -- three packets take about two seconds that way against nothing at all
# for three separate one packet runs.
#
# Response times are reported in microseconds, and both probe tools print
# milliseconds, hence the scaling.
probe_host() {
	local ph_host="$1"
	local ph_out="$2"
	local ph_status="ok"
	local ph_success=0
	local ph_sum=0
	local ph_min=0
	local ph_max=0
	local ph_rep=0
	local ph_mean=0
	local ph_addr ph_res ph_times
	local ph_count ph_lo ph_avg ph_hi

	ph_addr=$(get_ip_addr_used "${ph_host}" "${ip_proto}" "${iface}")

	while [ "${ph_rep}" -lt "${nbr_of_rep}" ]; do
		ph_rep=$((ph_rep+1))
		ph_count=""
		ph_lo=""
		ph_avg=""
		ph_hi=""

		if [ "${protocol_used}" = "ICMP" ]; then
			# family is deliberately unquoted: "--protocol 4" has to split in two
			# shellcheck disable=SC2086
			if [ -z "${device}" ]; then
				ph_res=$(ping -q ${family} -c 1 -W "${ping_timeout}" "${ph_host}" 2>&1)
			else
				ph_res=$(ping -q ${family} -c 1 -W "${ping_timeout}" -I "${device}" "${ph_host}" 2>&1)
			fi

			case "${ph_res}" in
				*"bad address"*|*"unknown host"*)
					ph_status="resolve"
					break
					;;
			esac

			echo "${ph_res}" | grep -q "received" || {
				ph_status="other"
				break
			}

			ph_times=$(echo "${ph_res}" | grep "min/avg/max")
			# A lost reply is one lost repetition, not the end of this host
			[ -z "${ph_times}" ] && continue

			ph_count=$(echo "${ph_res}" | grep "received" | awk '{print $4}')
			ph_times=$(echo "${ph_times}" | awk -F'=' '{ print $2 }')
			ph_lo=$(echo "${ph_times}" | awk -F'[=/ ]' '{ print $2 }')
			ph_avg=$(echo "${ph_times}" | awk -F'[=/ ]' '{ print $3 }')
			ph_hi=$(echo "${ph_times}" | awk -F'[=/ ]' '{ print $4 }')
		else
			# shellcheck disable=SC2086
			ph_res=$(udpechoclientd -c 1 -t "${timeout}" --host "${ph_host}" --port "${port}" ${family} -i "${device}" 2>&1)

			case "${ph_res}" in
				*"Can't Resolve Host Name"*)
					ph_status="resolve"
					break
					;;
			esac

			# No reply is a lost repetition. It used to be reported as a name
			# resolution failure, which it is not.
			echo "${ph_res}" | grep -q "RCVD" || continue

			ph_times=$(echo "${ph_res}" | grep "rtt")
			[ -z "${ph_times}" ] && continue

			ph_count=$(echo "${ph_res}" | grep "Rcvd" | awk -F': ' '{print $3}' | awk -F'(' '{ print $1 }')
			ph_hi=$(echo "${ph_times}" | awk -F': ' '{ print $2 }' | awk -F'ms' '{ print $1 }')
			ph_lo=$(echo "${ph_times}" | awk -F': ' '{ print $3 }' | awk -F'ms' '{ print $1 }')
			ph_avg=$(echo "${ph_times}" | awk -F': ' '{ print $4 }' | awk -F'ms' '{ print $1 }')
		fi

		[ "${ph_count:-0}" -gt 0 ] 2>/dev/null || continue

		ph_lo=$(echo "${ph_lo:-0}" "${micros}" | awk '{printf "%.0f", $1*$2}')
		ph_avg=$(echo "${ph_avg:-0}" "${micros}" | awk '{printf "%.0f", $1*$2}')
		ph_hi=$(echo "${ph_hi:-0}" "${micros}" | awk '{printf "%.0f", $1*$2}')

		ph_success=$((ph_success + ph_count))
		ph_sum=$((ph_sum + ph_avg))
		{ [ "${ph_min}" -eq 0 ] || [ "${ph_lo}" -lt "${ph_min}" ]; } && ph_min="${ph_lo}"
		[ "${ph_hi}" -gt "${ph_max}" ] && ph_max="${ph_hi}"
	done

	if [ "${ph_success}" -gt 0 ]; then
		ph_mean=$((ph_sum / ph_success))
	else
		ph_mean=0
		ph_min=0
		ph_max=0
	fi

	echo "${ph_status} ${ph_success} ${ph_min} ${ph_mean} ${ph_max} ${ph_host} ${ph_addr}" > "${ph_out}"
}

# Reads the $1 per-host result files and elects the fastest host, setting
# fasthost, min_time, avg_time, max_time, ip_addr_used, answered and
# resolve_failures. All the reported times belong to the elected host, which is
# what the data model asks for.
serverselection_reduce() {
	local sr_hosts="$1"
	local sr_idx=0
	local sr_status sr_success sr_min sr_avg sr_max sr_host sr_addr

	fasthost=""
	min_time=0
	avg_time=0
	max_time=0
	ip_addr_used=""
	answered=0
	resolve_failures=0

	while [ "${sr_idx}" -lt "${sr_hosts}" ]; do
		sr_idx=$((sr_idx+1))
		[ -f "${TMPDIR}/host_${sr_idx}" ] || continue
		read -r sr_status sr_success sr_min sr_avg sr_max sr_host sr_addr < "${TMPDIR}/host_${sr_idx}"

		[ "${sr_status}" = "resolve" ] && resolve_failures=$((resolve_failures+1))

		# A host that never answered cannot be the fastest. One unusable host no
		# longer discards the results of the others.
		[ "${sr_success:-0}" -gt 0 ] || continue
		answered=$((answered+1))

		if [ "${avg_time}" -eq 0 ] || [ "${sr_avg}" -lt "${avg_time}" ]; then
			fasthost="${sr_host}"
			min_time="${sr_min}"
			avg_time="${sr_avg}"
			max_time="${sr_max}"
			ip_addr_used="${sr_addr}"
		fi
	done
}

serverselection_launch() {
	local input="$1"
	local old_pid cmd n host

	json_load "${input}"

	json_get_var hostlist hostlist
	json_get_var port port
	json_get_var iface iface
	json_get_var ip_proto ip_proto
	json_get_var nbr_of_rep nbr_of_rep
	json_get_var timeout timeout
	json_get_var protocol_used protocol_used
	json_get_var proto proto
	json_get_var cancel cancel

	if [ "${proto}" = "both_proto" ]; then
		old_pid=$(cat /tmp/serverselection_pid)

		[ -n "${old_pid}" ] && {
			cmd=$(cat /proc/"${old_pid}"/cmdline)
		}

		if [[ "${cmd}" = *serverselection* ]]; then
			kill -9 "${old_pid}"
		fi

		if [ "${cancel}" -eq "1" ]; then
			$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.DiagnosticState="None"
			$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.FastestHost=""
			$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.MinimumResponseTime=0
			$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.AverageResponseTime=0
			$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.MaximumResponseTime=0
			$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.IPAddressUsed=""
			$UCI_COMMIT_BBF_DMMAP

			json_init
			json_add_string "Status" "None"
			json_add_string "FastestHost" ""
			json_add_string "IPAddressUsed" ""
			json_add_int "MinimumResponseTime" "0"
			json_add_int "AverageResponseTime" "0"
			json_add_int "MaximumResponseTime" "0"
			json_dump

			return
		else
			echo $$ > /tmp/serverselection_pid
			$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.DiagnosticState="Requested_running"
			$UCI_COMMIT_BBF_DMMAP
		fi
	fi

	# Assign default value
	[ -z "${nbr_of_rep}" ] && nbr_of_rep=3
	[ -z "${port}" ] && port=7
	[ -z "${protocol_used}" ] && protocol_used="ICMP"
	if [ -z "${timeout}" ]; then
		timeout=1
		ping_timeout=1
	else
		ping_timeout=$(awk -v ms="${timeout}" 'BEGIN{printf "%.3f", ms/1000}')
		timeout=$((timeout/1000))
		[ "${timeout}" = "0" ] && timeout=1
	fi
	# empty means either family, which is what get_ip_addr_used expects
	ip_proto="${ip_proto:-}"

	# Fail if hostlist is empty
	[ -z "${hostlist}" ] && {
		serverselection_error "Error_Internal" "${nbr_of_rep}" "${proto}"
		return
	}

	if [ -n "${iface}" ]; then
		device=$(ifstatus "${iface}" | jsonfilter -e '@.l3_device')

		# If no device was found, return error
		[ -z "${device}" ] && {
			serverselection_error "Error_NoRouteToHost" "${nbr_of_rep}" "${proto}"
			return
		}
	else
		device=""
	fi

	micros=1000

	# Convert the address family once, before the host loop. Doing it per host
	# overwrote ip_proto with the flag itself, so every host after the first
	# matched neither branch and was probed with no family selection at all --
	# and get_ip_addr_used, called with the flag instead of the family, stopped
	# resolving anything.
	family=""
	if [ "${protocol_used}" = "ICMP" ]; then
		[ "${ip_proto}" = "IPv4" ] && family="-4"
		[ "${ip_proto}" = "IPv6" ] && family="-6"
	else
		[ "${ip_proto}" = "IPv4" ] && family="--protocol 4"
		[ "${ip_proto}" = "IPv6" ] && family="--protocol 6"
	fi

	TMPDIR=$(mktemp -d /tmp/bbf_ss.XXXXXX 2>/dev/null)
	[ -z "${TMPDIR}" ] && {
		serverselection_error "Error_Internal" "${nbr_of_rep}" "${proto}"
		return
	}
	trap 'rm -rf "${TMPDIR}"' EXIT TERM INT

	# Probe every host in parallel, then pick the winner. A host that resolves but
	# never answers costs a full nbr_of_rep x timeout, so probing the hosts one
	# after another needs hosts x reps x timeout: at the ten hosts the data model
	# allows that overruns the ubus block carrying the result, and the caller is
	# told the ubus method does not exist. Repetitions within a host stay
	# sequential, which is what the data model requires.
	n=0
	for host in $(echo "${hostlist}" | tr "," "\n"); do
		n=$((n+1))
		probe_host "${host}" "${TMPDIR}/host_${n}" &
	done
	wait

	serverselection_reduce "${n}"

	# Nothing answered anywhere. The response times have to read zero, and a run
	# that measured nothing is not Complete.
	if [ "${answered}" -eq 0 ]; then
		if [ "${resolve_failures}" -gt 0 ]; then
			serverselection_error "Error_CannotResolveHostName" "${nbr_of_rep}" "${proto}"
		else
			serverselection_error "Error_Other" "${nbr_of_rep}" "${proto}"
		fi
		return
	fi

	json_init
	json_add_string "Status" "Complete"
	json_add_string "FastestHost" "${fasthost}"
	json_add_string "IPAddressUsed" "${ip_addr_used}"
	json_add_int "MinimumResponseTime" "${min_time}"
	json_add_int "AverageResponseTime" "${avg_time}"
	json_add_int "MaximumResponseTime" "${max_time}"
	json_dump

	# Store data in dmmap_diagnostics for both protocols (cwmp/usp). These are the
	# same values the JSON carries: the two sinks used to disagree, the JSON
	# reporting whichever host happened to be probed last.
	[ "${proto}" = "both_proto" ] && {
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.DiagnosticState="Complete"
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.FastestHost="${fasthost}"
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.IPAddressUsed="${ip_addr_used}"
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.MinimumResponseTime="${min_time}"
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.AverageResponseTime="${avg_time}"
		$UCI_SET_BBF_DMMAP dmmap_diagnostics.serverselection.MaximumResponseTime="${max_time}"
		$UCI_COMMIT_BBF_DMMAP
	}
}

if [ "$1" = "list" ]; then
	serverselection_list
elif [ -n "$1" ]; then
	serverselection_launch "$1"
else
	serverselection_error "Error_Internal" "1"
fi
