#!/bin/bash

# ctwx v4/v6 qos scripts
# Modified to support IPv6 marking

version="0.0.2"

usage()
{
	echo "Usage: $0 <-t type> <-i nic> <-r rate> <-a appid> [-d] [-v] [-h]"
	echo "options:"
	echo " -t type          operation type, e.g. set, get, del, reset"
	echo " -i nic           the nic name, e.g. eth0"
	echo " -r rate          the bandwidth value, e.g. 20mbit"
	echo " -a appid         the appid"
	echo " -d               enable debug message"
	echo " -v               show qos version"
	echo " -h               show help message"
}

type=
nic=
rate=
appid=
aid=
debug=0

while getopts "t:i:r:a:dvh" opt; do
    case $opt in
    t)
        type=$OPTARG
        ;;
	i)
		nic=$OPTARG
        ;;
	r)
		rate=$OPTARG
        ;;
	a)
		appid=$OPTARG
        ;;
	d)
		debug=1
		;;
	v)
		echo "${version}"
		exit 0
		;;
	h)
		usage
		exit 0
		;;
	?)
		echo "invalid option, try help"
		exit 1
		;;
	esac
done

validate_appid()
{
	if [ -z "${appid}" ]; then
		echo "appid is nil"
		exit 1
	fi

	if [ ${#appid} -lt 8 ]; then
		echo "appid is too short"
		exit 1
	fi
}

validate_args()
{
	if [ -z "${type}" ]; then
		echo "type is nil"
		exit 1
	fi

	if [ "${type}" != "set" -a "${type}" != "get" -a "${type}" != "del" -a "${type}" != "reset" ]; then
		echo "type is invalid"
		exit 1
	fi

	if [ -z "${nic}" ]; then
		echo "nic is nil"
		exit 1
	fi

	local res=$(ip -br a show ${nic})
	if [ -z "${res}" ]; then
		echo "the nic ${nic} is invalid"
		exit 1
	fi

	if [ "${type}" = "set" ]; then
		validate_appid

		if [ -z "${rate}" ]; then
			echo "operation '${type}' need 'rate' parameter"
			exit 1
		fi
	elif [ "${type}" = "get" ]; then
		validate_appid
	fi
}

log()
{
	if [ ${debug} -eq 1 ]; then
		local t=$(date '+%Y-%m-%d %H:%M:%S')
		echo "${t} $1" 1>&2
	fi
}

aid=${appid:0:8}
log "parameters, type=${type}, nic=${nic}, rate=${rate}, appid=${appid}, aid=${aid}"

validate_args

# singleton
lock_file="/run/.aid2fw.lock"
exec 200>${lock_file}
flock -n 200 || {
	echo "another qos script is running, exit"
	exit 1
}
log "acquired singleton lock"

aid_file=/run/.aid2fw.cfg
fw=

appid_to_fw()
{
	fw=
	local num=$(grep "=${aid}$" ${aid_file} 2>/dev/null | awk -F= '{ print $1}')
	if [ -n "${num}" ]; then
		fw=${num}
		return
	fi
	num=$(tail -n 1 ${aid_file} 2>/dev/null | awk -F= '{ print $1 }')
	if [ -z "${num}" ]; then
		num=1
	else
		if [ ${num} -gt 14 ]; then
			echo "too many appids"
			exit 1
		fi
		num=$((num + 1))
	fi

	log "generate new record, ${num}=${aid}"
	echo "${num}=${aid}" >> ${aid_file}
	fw=$(printf %x ${num})
	return
}

set_container_qos()
{
	local cid=$1
	local netcls

	log "process appid ${appid}, container id ${cid}"

	netcls=$(cat /sys/fs/cgroup/net_cls/docker/${cid}/net_cls.classid 2>/dev/null)
	if [ -z "${netcls}" ]; then
		echo "not found container ${cid} netcls id"
		return
	fi

	# 1. Setup IPv4 iptables Chain
	iptables -t mangle -S mqos 1 1>/dev/null 2>/dev/null
	if [ $? -ne 0 ]; then
		log "cmd: iptables -t mangle -N mqos"
		iptables -t mangle -N mqos
	fi

	# 2. Setup IPv6 ip6tables Chain (NEW)
	ip6tables -t mangle -S mqos 1 1>/dev/null 2>/dev/null
	if [ $? -ne 0 ]; then
		log "cmd: ip6tables -t mangle -N mqos"
		ip6tables -t mangle -N mqos 2>/dev/null
	fi

	# 3. Add Output Jump Rule IPv4
	iptables -t mangle -C OUTPUT -o ${nic} -j mqos 2>/dev/null
	if [ $? -ne 0 ]; then
		log "cmd: iptables -t mangle -A OUTPUT -o ${nic} -j mqos"
		iptables -t mangle -A OUTPUT -o ${nic} -j mqos
	fi

	# 4. Add Output Jump Rule IPv6 (NEW)
	ip6tables -t mangle -C OUTPUT -o ${nic} -j mqos 2>/dev/null
	if [ $? -ne 0 ]; then
		log "cmd: ip6tables -t mangle -A OUTPUT -o ${nic} -j mqos"
		ip6tables -t mangle -A OUTPUT -o ${nic} -j mqos 2>/dev/null
	fi

	appid_to_fw
	if [ -z "${fw}" ]; then
		echo "get fw value failed"
		exit 1
	else
		log "get appid fw value, ${aid}=${fw}"
	fi

	# 5. Mark Packets IPv4
	iptables -t mangle -C mqos -m cgroup --cgroup ${netcls} -m comment --comment ${aid} -j MARK --set-xmark 0x${fw}00/0xff00 2>/dev/null
	if [ $? -ne 0 ]; then
		log "cmd: iptables -t mangle -A mqos -m cgroup --cgroup ${netcls} -m comment --comment ${aid} -j MARK --set-xmark 0x${fw}00/0xff00"
		iptables -t mangle -A mqos -m cgroup --cgroup ${netcls} -m comment --comment ${aid} -j MARK --set-xmark 0x${fw}00/0xff00
	fi

	# 6. Mark Packets IPv6 (NEW)
	ip6tables -t mangle -C mqos -m cgroup --cgroup ${netcls} -m comment --comment ${aid} -j MARK --set-xmark 0x${fw}00/0xff00 2>/dev/null
	if [ $? -ne 0 ]; then
		log "cmd: ip6tables -t mangle -A mqos -m cgroup --cgroup ${netcls} -m comment --comment ${aid} -j MARK --set-xmark 0x${fw}00/0xff00"
		ip6tables -t mangle -A mqos -m cgroup --cgroup ${netcls} -m comment --comment ${aid} -j MARK --set-xmark 0x${fw}00/0xff00 2>/dev/null
	fi

	# TC logic remains same as it matches on fw mark (protocol independent)
	local res=$(tc qdisc show dev ${nic} 2>/dev/null | grep 'qdisc htb ffee: root' | grep 'default 10')
	if [ -z "${res}" ]; then
		tc qdisc del dev ${nic} root 2>/dev/null
		log "cmd: tc qdisc add dev ${nic} root handle ffee: htb default 10"
		tc qdisc add dev ${nic} root handle ffee: htb default 10
	fi

	res=$(tc class show dev ${nic} parent ffee: classid ffee:1 2>/dev/null | grep 'htb' | grep 'rate 10Gbit ceil 10Gbit')
	if [ -z "${res}" ]; then
		log "cmd: tc class add dev ${nic} parent ffee: classid ffee:1 htb rate 10gbit"
		tc class add dev ${nic} parent ffee: classid ffee:1 htb rate 10gbit
	fi

	res=$(tc class show dev ${nic} parent ffee:1 classid ffee:10 2>/dev/null | grep 'htb' | grep 'rate 10Gbit ceil 10Gbit')
	if [ -z "${res}" ]; then
		log "cmd: tc class add dev ${nic} parent ffee:1 classid ffee:10 htb rate 10gbit"
		tc class add dev ${nic} parent ffee:1 classid ffee:10 htb rate 10gbit
	fi

	res=$(tc class show dev ${nic} parent ffee:1 classid ffee:1${fw} 2>/dev/null | grep 'htb' | grep -i "rate ${rate} ceil ${rate}")
	if [ -z "${res}" ]; then
		log "cmd: tc class replace dev ${nic} parent ffee:1 classid ffee:1${fw} htb rate ${rate}"
		tc class replace dev ${nic} parent ffee:1 classid ffee:1${fw} htb rate ${rate}
	fi

	res=$(tc filter show dev ${nic} protocol all parent ffee: 2>/dev/null | grep fw | grep "handle 0x${fw}00/0xff00 classid ffee:1${fw}")
	if [ -z "${res}" ]; then
		log "cmd: tc filter add dev ${nic} protocol all parent ffee: prio 1 handle 0x${fw}00/0xff00 fw classid ffee:1${fw}"
		tc filter add dev ${nic} protocol all parent ffee: prio 1 handle 0x${fw}00/0xff00 fw classid ffee:1${fw}
	fi
}

set_qos()
{
	# get containers by appid
	log "get containers by appid ${appid}"
	local containers=$(docker ps --no-trunc | grep "${appid}" | awk '{ printf("%s ", $1) }')
	if [ -z "${containers}" ]; then
		log "no container found by appid ${appid}"
		exit 0
	fi

	for container in ${containers}
	do
		set_container_qos ${container}
	done
}

get_qos()
{
	# Just checking iptables is enough to know configuration exists, as we sync both.
	# -A mqos -m cgroup --cgroup 112233 -m comment --comment 00000000 -j MARK --set-xmark 0x100/0xff00
	local mark=$(iptables -t mangle -S mqos 2>/dev/null | grep "${aid}" | awk '{ print $14 }')
	if [ -z "${mark}" ]; then
		log "not found appid in iptables"
		exit 1
	fi
	# filter pref 49152 fw chain 0 handle 0x100/0xff00 classid ffee:11
	local classid=$(tc filter show dev ${nic} protocol all parent ffee: 2>/dev/null | grep fw | grep "${mark}" | awk '{ print $10 }')
	if [ -z "${classid}" ]; then
		log "not found appid in filter"
		exit 1
	fi
	# class htb ffee:11 parent ffee:1 prio 0 rate 100Mbit ceil 100Mbit burst 1600b cburst 1600b
	local res=$(tc class show dev ${nic} parent ffee:1 2>/dev/null | grep "htb ${classid}" | awk '{ print $9}')
	if [ -z "${res}" ]; then
		log "not found appid in class"
		exit 1
	fi
	echo "${res}"
	exit 0
}

del_qos()
{
	if [ -z "${rate}" -a -z "${appid}" ]; then
		log "cmd: tc qdisc del dev ${nic} root"
		tc qdisc del dev ${nic} root
		exit 0
	fi

	# -A mqos -m cgroup --cgroup 112233 -m comment --comment 00000000 -j MARK --set-xmark 0x100/0xff00
	local mark=$(iptables -t mangle -S mqos 2>/dev/null | grep "${aid}" | awk '{ print $14 }')
	if [ -z "${mark}" ]; then
		log "not found appid in iptables"
		appid_to_fw
		if [ -z "${fw}" ]; then
			log "not found appid's fw"
			exit 0
		fi
		mark="0x${fw}00/0xff00"
	fi

	# Clean up iptables rule matching this AppID
	log "cleaning up iptables rules for ${aid}"
	# Using loop to remove all containers rules for this appid
	iptables -t mangle -S mqos 2>/dev/null | grep "${aid}" | while read -r line; do
		# Remove "-A mqos" prefix to form delete command
		rule=${line#-A mqos}
		iptables -t mangle -D mqos $rule
	done

	# Clean up ip6tables rules matching this AppID (NEW)
	log "cleaning up ip6tables rules for ${aid}"
	ip6tables -t mangle -S mqos 2>/dev/null | grep "${aid}" | while read -r line; do
		rule=${line#-A mqos}
		ip6tables -t mangle -D mqos $rule 2>/dev/null
	done

	# filter pref 49152 fw chain 0 handle 0x100/0xff00 classid ffee:11
	local classid=$(tc filter show dev ${nic} protocol all parent ffee: 2>/dev/null | grep fw | grep "${mark}" | awk '{ print $10 }')
	if [ -z "${classid}" ]; then
		log "not found appid in filter"
	else
		log "cmd: tc filter del dev ${nic} protocol all parent ffee: prio 1 handle 0x${fw}00/0xff00 fw classid ffee:1${fw}"
		tc filter del dev ${nic} protocol all parent ffee: prio 1 handle 0x${fw}00/0xff00 fw classid ffee:1${fw}
	fi
	# class htb ffee:11 parent ffee:1 prio 0 rate 100Mbit ceil 100Mbit burst 1600b cburst
	local res=$(tc class show dev ${nic} parent ffee:1 2>/dev/null | grep "htb ${classid}" | awk '{ print $9}')
	if [ -z "${res}" ]; then
		log "not found appid in class"
	else
		if [ -n "${rate}" ]; then
			log "cmd: tc class del dev ${nic} parent ffee:1 classid ffee:1${fw} htb rate ${rate}"
			tc class del dev ${nic} parent ffee:1 classid ffee:1${fw} htb rate ${rate} 2>/dev/null
		else
			log "cmd: tc class del dev ${nic} parent ffee:1 classid ffee:1${fw} htb rate ${rate}"
			tc class del dev ${nic} parent ffee:1 classid ffee:1${fw} htb 2>/dev/null
		fi
	fi

	exit 0
}

reset_qos()
{
	log "cmd: tc qdisc del dev ${nic} root"
	tc qdisc del dev ${nic} root 2>/dev/null

	log "cmd: iptables -t mangle -A OUTPUT -o ${nic} -j mqos"
	for((i=0;i<10;i++)); do
		iptables -t mangle -D OUTPUT -o ${nic} -j mqos 2>/dev/null
		if [ $? -ne 0 ]; then
			break
		fi
	done

	# Clean IPv6 OUTPUT jumps (NEW)
	log "cmd: ip6tables -t mangle -A OUTPUT -o ${nic} -j mqos"
	for((i=0;i<10;i++)); do
		ip6tables -t mangle -D OUTPUT -o ${nic} -j mqos 2>/dev/null
		if [ $? -ne 0 ]; then
			break
		fi
	done

	log "cmd: iptables -t mangle -F mqos 2>/dev/null"
	iptables -t mangle -F mqos 2>/dev/null

	log "cmd: iptables -t mangle -X mqos 2>/dev/null"
	iptables -t mangle -X mqos 2>/dev/null

	# Clean IPv6 Chain (NEW)
	log "cmd: ip6tables -t mangle -F mqos 2>/dev/null"
	ip6tables -t mangle -F mqos 2>/dev/null

	log "cmd: ip6tables -t mangle -X mqos 2>/dev/null"
	ip6tables -t mangle -X mqos 2>/dev/null

	log "cmd: rm -f ${aid_file}"
	rm -f ${aid_file}

	exit 0
}

if [ "${type}" = "set" ]; then
	set_qos
elif [ "${type}" = "get" ]; then
	get_qos
elif [ "${type}" = "del" ]; then
	del_qos
elif [ "${type}" = "reset" ]; then
	reset_qos
fi