#!/bin/sh

######################################################################
# AutoBW
# Set QoS upload/download speeds using Ookla speedtests and spdMerlin
# Requires bc and spdMerlin (v3.3.2 or later)
# Some of the functionality here was either taken directly or
#  adapted from the FreshJr_QOS script
#
# v1.0 (3/25/2020): Initial Release
# v1.1 (4/01/2020): Complete rewrite
#	     	    -No longer requires bash
#		    -Checks that bc is installed & QoS enabled &
#                     upload traffic is directed to unthrottled class
#		    -No need to stop/start QoS
#		    -Class rates/ceils are scaled
# v1.2 (4/03/2020): More changes
#		    -Check for spdMerlin
#		    -Formatted logger message concerning new BW
#		    -Check for appropriate iptable rules that ensure
#             spdMerlin traffic in unaffected by QoS
# v1.3 (4/03/2020): Changed an important check
#		    -Check for spdMerlin v3.3.2 or greater to ensure
#             that QoS will not interfere with calculated speeds
# v1.4 (4/03/2020): Added option for automatic updates
#		    -If turned on, check for updates. If found then update.
# v1.5 (4/05/2020): Now handles VPN and Gbps|Mbps|Kbps rates/ceils
#			-Use VPN values, if defined (thanks maghuro)
#			-Handle rate/ceil units other than just Kbps (thanks peepsnet)
######################################################################
readonly SCRIPT_AUTHOR="squidbillyms"
readonly SCRIPT_VERSION="v1.5"
readonly SCRIPT_NAME="AutoBW"
readonly SCRIPT_BRANCH="master"
readonly SCRIPT_REPO="https://raw.githubusercontent.com/$SCRIPT_AUTHOR/$SCRIPT_NAME/$SCRIPT_BRANCH"

#Should I automatically check for and perform updates
# 0=no ( I will update manually when I am ready manually), 1=yes do it for me
# update manually using the following:
#curl "https://raw.githubusercontent.com/squidbillyms/AutoBW/master/AutoBW" -o /jffs/scripts/AutoBW && chmod 755 AutoBW
autoupdate=0

#Scale average values by this factor
# Adjust to optimize bufferbloat and quality grade at DSLreports.com
down_scale_factor=0.85
up_scale_factor=0.85

#Make sure speeds (AFTER scaling) are within these boundaries
# Mainly just to prevent bandwidth from being set too low
# (or too high)... not worried about upper limit but it is there
# for the sake of completeness
download_lower_limit=10   #Mbps
download_upper_limit=100  #Mbps
upload_lower_limit=1      #Mbps
upload_upper_limit=100    #Mbps

#Calculate the average download and upload speeds over this many
# consecutive spdMerlin runs
num_spd_tests=2

#Pause between successive spdMerlin runs (if num_spd_tests > 1)
# 0 = no pause, 1 = yes pause
pause_between_test=0
#If yes (or 1), pause for how long (will be used with sleep command)
pause_this_long=15s

#spdMerlin command
spdMer_command="/jffs/scripts/spdmerlin generate"

#Use VPN values? If no, set 0. If yes, specify VPN client number
Use_VPN=0

#Location of upload/download stats
if [ $Use_VPN -eq 0 ]; then
  upload_stats="/jffs/addons/spdmerlin.d/csv/Uploaddaily_WAN.htm"
  download_stats="/jffs/addons/spdmerlin.d/csv/Downloaddaily_WAN.htm"
else
  upload_stats="/jffs/addons/spdmerlin.d/csv/Uploaddaily_VPNC$Use_VPN.htm"
  download_stats="/jffs/addons/spdmerlin.d/csv/Downloaddaily_VPNC$Use_VPN.htm"
fi

#Test Mode (run and calculate but do not make any changes)
# meaning no "nvram commit" or "tc class change" commands 
# will be executed
# 1 = test mode (no changes), 0 = make permenant changes 
test_mode=1
######################################################################

update() {	#Adapted from code by Jack Yaz
	printf "\n AutoUpdate is enabled. Checking for update ...\n"
	doupdate="false"
	localver=$(grep "SCRIPT_VERSION=" /jffs/scripts/"$SCRIPT_NAME" | grep -m1 -oE 'v[1-9]([.][0-9])')
	/usr/sbin/curl -fsL --retry 3 "$SCRIPT_REPO/$SCRIPT_NAME" | grep -qF $SCRIPT_AUTHOR || { echo -e " ERROR: Failed to check new version, continuing anyway."; return 1; }
	serverver=$(/usr/sbin/curl -fsL --retry 3 "$SCRIPT_REPO/$SCRIPT_NAME" | grep "SCRIPT_VERSION=" | grep -m1 -oE 'v[1-9]([.][0-9])')
	if [ "$localver" != "$serverver" ]; then
		doupdate="version"
		printf " Found new version: this is %s ---> latest version is %s.\n" $localver $serverver
	else
		localmd5="$(md5sum "/jffs/scripts/$SCRIPT_NAME" | awk '{print $1}')"
		remotemd5="$(curl -fsL --retry 3 "$SCRIPT_REPO/$SCRIPT_NAME" | md5sum | awk '{print $1}')"
		if [ "$localmd5" != "$remotemd5" ]; then
			#doupdate="md5"
			printf " Something is different about this copy of %s. Just letting you know. :)\n" $SCRIPT_NAME
		fi
	fi
		
	if [ "$doupdate" != "false" ]; then
		/usr/sbin/curl -fsL --retry 3 "$SCRIPT_REPO/$SCRIPT_NAME" -o "/jffs/scripts/$SCRIPT_NAME" &> /dev/null
		chmod 0755 /jffs/scripts/"$SCRIPT_NAME"
		printf " $SCRIPT_NAME was successfully updated to %s. Please run it again.\n\n" $serverver
        exit 0
	else
		printf " Based on versioning alone, you have the latest version: %s\n" $serverver
	fi
}

check() {	
	#Check to make sure bc and spdMerlin (v3.3.2 or later) is installed,
	# QoS is enabled, and FreshJr_QoS is not in delayed start
	bc -v &> /dev/null
        if [ "$?" -ne 0 ]; then
		echo " ERROR: 'bc' is required but does not appear to be available."
		exit 1;
	elif [ ! -x "/jffs/scripts/spdmerlin" ]; then
		echo " ERROR: spdMerlin does not exist as an executable script in '/jffs/scripts/' " 
		exit 1; 
	#Check to make sure spdMerlin version is 3.3.2 or later (>=3.3.2 prevents QoS from interfering with speedtests)
	elif [ $(echo -e "$(grep -m1 'SCRIPT_VERSION=' /jffs/scripts/spdmerlin | awk -F"\"v|\"" '{print $2}')\n3.3.2" | sort -r | tail -n1) != "3.3.2" ]; then
		echo " ERROR: Need spdMerlin version >= 3.3.2 to ensure QoS doesn't interfere with calculated speeds."
		echo " ERROR: Try running 'spdmerlin update'"
		exit 1;	
	elif [ "$(nvram get qos_enable)" == "0" ]; then
		echo " ERROR: QoS does not appear to be enabled."
                exit 1;
	else
		for pid in $(pidof FreshJR_QOS); do
			if [ $pid != $$ ]; then
				if ! [ "$(ps -w | grep "${pid}.*\(install\|menu\|rules\|rates\)" | grep -v "grep")" ] ; then
					kill $pid
				fi
			fi
		done
	fi

	if [ "$test_mode" -eq "1" ]; then
		echo -e "\n TEST MODE: No changes will be made."
	fi
}

get_new_speeds() {	#Get new down/up speeds using spdMerlin
	echo
	echo "               ---------------- SPDMERLIN -------------"
	echo "                 Download (Kbps)       Upload (Kbps)"
	echo "               -------------------  -------------------"

	#Calculate average download and upload speeds
	sum_down=0; sum_up=0; x=1
	while [ "$x" -le "$num_spd_tests" ]
	do
		printf "\rRunning %2.0f of %2.0f speedtest, please wait... " $x $num_spd_tests 
                ${spdMer_command} &> /dev/null
                Mbps_down=$(tail -n1 $download_stats | awk -F ',' '{print $NF}')
                Mbps_up=$(tail -n1 $upload_stats | awk -F ',' '{print $NF}')
                Kbps_down=$(echo "$Mbps_down*1024" | bc -l)
                Kbps_up=$(echo "$Mbps_up*1024" | bc -l)

                sum_down=$(echo $sum_down+$Kbps_down | bc -l)
                sum_up=$(echo $sum_up+$Kbps_up | bc -l)

		printf "\rTest %2.0f of %2.0f    %10.1f           %10.1f\n" $x $num_spd_tests $Kbps_down $Kbps_up

  		x=$(( $x + 1 ))
		if [ "$pause_between_test" -eq 1 ] && [ "$x" -le "$num_spd_tests" ]; then 
  			printf "Sleeping for %s ..." $pause_this_long  
			sleep $pause_this_long
                fi
	done

	#Calculate the average values
	Kbps_down=$(echo "$sum_down/$num_spd_tests" | bc -l)
	Kbps_up=$(echo "$sum_up/$num_spd_tests" | bc -l)
	if [ "$num_spd_tests" -gt "1" ]; then
		printf "Average          %10.1f           %10.1f\n" $Kbps_down $Kbps_up
	fi

	#Apply user-defined scale factors
	printf "Scale Factors         %5.2f                %5.2f\n" $down_scale_factor $up_scale_factor
	Kbps_down=$(echo "$Kbps_down*$down_scale_factor" | bc -l)
	Kbps_up=$(echo "$Kbps_up*$up_scale_factor" | bc -l)
	printf "Scaled Speeds    %10.1f           %10.1f\n" $Kbps_down $Kbps_up

	#Make sure download and uploads speeds are within defined user-defined limits above
	download_lower_limit=$(echo "$download_lower_limit*1024" | bc -l)
	download_upper_limit=$(echo "$download_upper_limit*1024" | bc -l)
	upload_lower_limit=$(echo "$upload_lower_limit*1024" | bc -l)
	upload_upper_limit=$(echo "$upload_upper_limit*1024" | bc -l)
	outta_bounds=0
	if [ $(echo "$Kbps_down < ($download_lower_limit)" | bc -l) -eq 1 ]; then
   		printf "* Download speed (%8.1f Kbps) < lower limit (%8.1f Kbps)\n" $Kbps_down $download_lower_limit | logger -t $SCRIPT_NAME
   		printf "* Download speed (%8.1f Kbps) < lower limit (%8.1f Kbps)\n" $Kbps_down $download_lower_limit 
    		Kbps_down=$download_lower_limit
		outta_bounds=1
	elif [ $(echo "$Kbps_down > $download_upper_limit" | bc -l) -eq 1 ]; then
   		printf "* Download speed (%8.1f Kbps) > upper limit (%8.1f Kbps)\n" $Kbps_down $download_upper_limit | logger -t $SCRIPT_NAME
   		printf "* Download speed (%8.1f Kbps) > upper limit (%8.1f Kbps)\n" $Kbps_down $download_upper_limit 
   		Kbps_down=$download_upper_limit
		outta_bounds=1
	fi
	if [ $(echo "$Kbps_up < $upload_lower_limit" | bc -l) -eq 1 ]; then
   		printf "* Upload speed (%8.1f Kbps) < lower limit (%8.1f Kbps)\n" $Kbps_up $upload_lower_limit | logger -t $SCRIPT_NAME
   		printf "* Upload speed (%8.1f Kbps) < lower limit (%8.1f Kbps)\n" $Kbps_up $upload_lower_limit 
   		Kbps_up=$upload_lower_limit
		outta_bounds=1
	elif [ $(echo "$Kbps_up > $upload_upper_limit" | bc -l) -eq 1 ]; then
   		printf "* Upload speed (%8.1f Kbps) > upper limit (%8.1f Kbps)\n" $Kbps_up $upload_upper_limit | logger -t $SCRIPT_NAME
   		printf "* Upload speed (%8.1f Kbps) > upper limit (%8.1f Kbps)\n" $Kbps_up $upload_upper_limit 
   		Kbps_up=$upload_upper_limit
		outta_bounds=1
	fi

	if [ "$outta_bounds" -eq "1" ]; then
		printf "Corrected Speeds   %10.1f            %10.1f\n" $Kbps_down $Kbps_up
	fi	

	#Get current QoS down/up speeds
	old_Kbps_up=$( nvram get qos_obw )
	old_Kbps_down=$( nvram get qos_ibw )

	if [ "$test_mode" -eq "0" ]; then
		#Set Upload/Download Limit
		printf " -----> Setting QoS Download Speed to %8.1f Kbps ..." $Kbps_down | logger -t $SCRIPT_NAME
		printf " -----> Setting QoS Upload Speed to %8.1f Kbps ..." $Kbps_up | logger -t $SCRIPT_NAME
		nvram set qos_ibw=$(echo $Kbps_down | cut -d. -f1)
		nvram set qos_obw=$(echo $Kbps_up | cut -d. -f1)
		nvram commit
	fi 

	echo
	echo "               ----------------- QOS ----------------"
	echo "                 Download (Kbps)      Upload (Kbps)"
	echo "               -------------------  -----------------"
	printf "%s           %10.1f          %10.1f\n" "Previous" $old_Kbps_down $old_Kbps_up
	printf "%s   %10.1f          %10.1f\n" "New (from above)" $Kbps_down $Kbps_up
	
	
}

update_via_TC() {
	#read parameters for tc
	PARMS=""
	OVERHEAD=$(nvram get qos_overhead)
	if [ ! -z "$OVERHEAD" ] && [ "$OVERHEAD" -gt "0" ] ; then
			ATM=$(nvram get qos_atm)
			if [ "$ATM" == "1" ] ; then
					PARMS="overhead $OVERHEAD linklayer atm "
			else
					PARMS="overhead $OVERHEAD linklayer ethernet "
			fi
	fi

	#Calculate scale factor based on new/old speeds. This factor will be used to scale class rates & ceils
	#If old speeds are 0, just set the scale factor to 1 (cannot divide by 0)
	if [ "$old_Kbps_up" -eq "0" ]; then
		uscale_factor=1.0
	else
		uscale_factor=$( echo "$Kbps_up/$old_Kbps_up" | bc -l )
	fi
	if [ "$old_Kbps_down" -eq "0" ]; then
		dscale_factor=1.0
	else
		dscale_factor=$( echo "$Kbps_down/$old_Kbps_down" | bc -l ) 
	fi

	printf "%s           %8.3f            %8.3f\n" "Scaling by" $dscale_factor $uscale_factor

	echo
	echo "              ====================================== DOWNLOAD ======================================"
	echo "               ----------------- RATE -----------------    ----------------- CEIL -----------------"
	echo "Class            Previous (Kbps)        New (Kbps)           Previous (Kbps)         New (Kbps) "
	echo "-----          -------------------  -------------------    -------------------  -------------------"
	while read -r line;
	do
		class=$( echo ${line} | sed -n -e 's/.*class htb 1://p' | cut -d' ' -f1 )
		prio=$( echo ${class} | cut -c 2 )

		prevRate=$( echo ${line} | awk -F"rate |K?M?G?bit" '{print $2}' )
		prevCeil=$( echo ${line} | awk -F"ceil |K?M?G?bit" '{print $3}' )
		rateUnit=$( echo ${line} | cut -d' ' -f11)
		rateUnit=${rateUnit//[0-9]/}
		ceilUnit=$( echo ${line} | cut -d' ' -f15)
		ceilUnit=${ceilUnit//[0-9]/}
		if [ "$rateUnit" = "bit" ]; then
			prevRate=$( echo "$prevRate*0.001" | bc -l)
		fi
		if [ "$rateUnit" = "Mbit" ]; then
			prevRate=$( echo "$prevRate*1000" | bc -l)
		fi
		if [ "$rateUnit" = "Gbit" ]; then
			prevRate=$( echo "$prevRate*1000000" | bc -l)
		fi
		
		if [ "$ceilUnit" = "bit" ]; then
			prevCeil=$( echo "$prevCeil*0.001" | bc -l)
		fi
		if [ "$ceilUnit" = "Mbit" ]; then
			prevCeil=$( echo "$prevCeil*1000" | bc -l)
		fi
		if [ "$ceilUnit" = "Gbit" ]; then
			prevCeil=$( echo "$prevCeil*1000000" | bc -l)
		fi

		Burst=$( echo ${line} | sed -n -e 's/.* burst \([a-zA-Z0-9]*\).*/\1/p' )
		Cburst=$( echo ${line} | sed -n -e 's/.*cburst \([a-zA-Z0-9]*\).*/\1/p' )

		Rate=$(echo "$prevRate*$dscale_factor" | bc -l)
		Ceil=$(echo "$prevCeil*$dscale_factor" | bc -l)

		printf "%2.0f                 %10.1f         %10.1f                %10.1f         %10.1f\n" $class $prevRate $Rate $prevCeil $Ceil
	
		#Debug
		#echo -e "$line\n class=$class\n prio=$prio\n Rate=$Rate\n Ceil=$Ceil\n Burst=$Burst\n Cburst=$Cburst\n"

		#echo -e "tc class change dev br0 parent 1:1 classid 1:${class} htb ${PARMS}prio ${prio} rate ${Rate}Kbit ceil ${Ceil}Kbit burst ${Burst} cburst ${Cburst}"
		if [ "$test_mode" -eq "0" ]; then
		Rate=$(echo $Rate | cut -d. -f1)
		Ceil=$(echo $Ceil | cut -d. -f1)
		
		tc class change dev br0 parent 1:1 classid 1:${class} htb ${PARMS}prio ${prio} rate ${Rate}Kbit ceil ${Ceil}Kbit burst ${Burst} cburst ${Cburst}
		
		fi

	done <<EOF
                $( tc class show dev br0 | grep "parent 1:1 " )
EOF
 	
	echo
	echo "              ======================================= UPLOAD ======================================="
	echo "               ----------------- RATE -----------------    ----------------- CEIL -----------------"
	echo "Class            Previous (Kbps)        New (Kbps)           Previous (Kbps)         New (Kbps) "
	echo "-----          -------------------  -------------------    -------------------  -------------------"
	while read -r line;
	do
		class=$( echo ${line} | sed -n -e 's/.*class htb 1://p' | cut -d' ' -f1 )
		prio=$( echo ${class} | cut -c 2 )

		prevRate=$( echo ${line} | awk -F"rate |K?M?G?bit" '{print $2}' )
		prevCeil=$( echo ${line} | awk -F"ceil |K?M?G?bit" '{print $3}' )
		rateUnit=$( echo ${line} | cut -d' ' -f11)
		rateUnit=${rateUnit//[0-9]/}
		ceilUnit=$( echo ${line} | cut -d' ' -f15)
		ceilUnit=${ceilUnit//[0-9]/}
		if [ "$rateUnit" = "bit" ]; then
			prevRate=$( echo "$prevRate*0.001" | bc -l)
		fi
		if [ "$rateUnit" = "Mbit" ]; then
			prevRate=$( echo "$prevRate*1000" | bc -l)
		fi
		if [ "$rateUnit" = "Gbit" ]; then
			prevRate=$( echo "$prevRate*1000000" | bc -l)
		fi
		
		if [ "$ceilUnit" = "bit" ]; then
			prevCeil=$( echo "$prevCeil*0.001" | bc -l)
		fi
		if [ "$ceilUnit" = "Mbit" ]; then
			prevCeil=$( echo "$prevCeil*1000" | bc -l)
		fi
		if [ "$ceilUnit" = "Gbit" ]; then
			prevCeil=$( echo "$prevCeil*1000000" | bc -l)
		fi
		
		Burst=$( echo ${line} | sed -n -e 's/.* burst \([a-zA-Z0-9]*\).*/\1/p' )
		Cburst=$( echo ${line} | sed -n -e 's/.*cburst \([a-zA-Z0-9]*\).*/\1/p' )

		Rate=$(echo "$prevRate*$uscale_factor" | bc -l)
		Ceil=$(echo "$prevCeil*$uscale_factor" | bc -l)

		printf "%2.0f                 %10.1f         %10.1f                %10.1f         %10.1f\n" $class $prevRate $Rate $prevCeil $Ceil
	
		#Debug
		#echo -e "$line\n"
		#echo -e "$line\n class=$class\n prio=$prio\n Rate=$Rate\n Ceil=$Ceil\n Burst=$Burst\n Cburst=$Cburst\n"

		if [ "$test_mode" -eq "0" ]; then
		Rate=$(echo $Rate | cut -d. -f1)
		Ceil=$(echo $Ceil | cut -d. -f1)
		tc class change dev eth0 parent 1:1 classid 1:${class} htb ${PARMS}prio ${prio} rate ${Rate}Kbit ceil ${Ceil}Kbit burst ${Burst} cburst ${Cburst}
		fi

	done <<EOF
                $( tc class show dev eth0 | grep "parent 1:1" )
EOF
	if [ "$test_mode" -eq "1" ]; then
		echo -e "\n TEST MODE: No changes were made."
	fi
	echo
}

########### START HERE ###########
if [ "$autoupdate" -ne "0" ]; then update; fi;
check
get_new_speeds
update_via_TC