#!/bin/env bash

set -o pipefail

export LC_ALL=C.UTF-8

NAME='tunet-bash'
VERSION='1.2.9'

AUTH4_LOG_URL='https://auth.tsinghua.edu.cn/cgi-bin/srun_portal'
AUTH4_CHALLENGE_URL='https://auth.tsinghua.edu.cn/cgi-bin/get_challenge'
AUTH6_LOG_URL='https://auth6.tsinghua.edu.cn/cgi-bin/srun_portal'
AUTH6_CHALLENGE_URL='https://auth6.tsinghua.edu.cn/cgi-bin/get_challenge'
TUNET_USER_INFO_URL='https://login.tsinghua.edu.cn/cgi-bin/rad_user_info'
TUNET_USER_INFO_JSON_URL="${TUNET_USER_INFO_URL}?callback=any"
REGEX_USER_INFO_JSON='"billing_name":"([^"]+)".*"online_device_total":"([^"]+)"[^}]*"products_name":"([^"]+)"[^}]*"sysver":"([^"]+)"[^}]*"user_balance":([^,]+)[^}]*"user_mac":"([^"]+)"'
REDIRECT_URL='http://info.tsinghua.edu.cn/'
REGEX_AC_ID='//auth[46]\.tsinghua\.edu\.cn/index_([0-9]+)\.html'

KEY_ARRAY_LENGTH=4

CACHE_DIR="$HOME/.cache/$NAME"
LOG_LEVEL=${TUNET_LOG_LEVEL:-'info'}

key_array=()
data_array=()

verbose=0
ipv=4
date_format='--rfc-3339 s'
op='whoami'

USERNAME="$TUNET_USERNAME"
PASSWORD="$TUNET_PASSWORD"
PASSNAME="$TUNET_PASSNAME"

fill_key() {
	local key="$1"
	local array=()
	for ((i = 0; i < ${#key}; i++)); do
		local char=${key:$i:1}
		local value=$(printf "%d" "'$char")
		array+=($value)
	done
	while [ ${#array[@]} -lt 16 ]; do
		array+=(0)
	done
	for ((i = 0; i < KEY_ARRAY_LENGTH; i++)); do
		local temp_u32=0
		for ((j = 0; j < 4; j++)); do
			temp_u32=$((temp_u32 | (array[$((i * 4 + j))] << (8 * j))))
		done
		key_array[i]=$temp_u32
	done
}

fill_data() {
	local data="$1"
	local data_len=${#data}
	for ((i = 0; i < data_len; i++)); do
		local char=${data:$i:1}
		data_array+=($(printf "%d" "'$char"))
	done
	local n=$(((data_len + 3) / 4))
	local out_len=$(((n + 1) * 4))
	while [ ${#data_array[@]} -lt $out_len ]; do
		data_array+=(0)
	done
	data_array[$((out_len - 4))]=$(($data_len & 0xFF))
	data_array[$((out_len - 3))]=$(($data_len >> 8 & 0xFF))
	data_array[$((out_len - 2))]=$(($data_len >> 16 & 0xFF))
	data_array[$((out_len - 1))]=$(($data_len >> 24 & 0xFF))
}

encode() {
	local n=$((${#data_array[@]} / 4 - 1))
	local q=$((6 + 52 / (n + 1)))
	local d=0
	local z=0
	for ((i = 0; i < 4 && $((n * 4 + i)) < ${#data_array[@]}; i++)); do
		z=$((z | (data_array[$((n * 4 + i))] << (8 * i))))
	done
	for ((i = 0; i < q; i++)); do
		d=$((d + 0x9E3779B9))
		local e=$(((d >> 2) & 3))
		for ((p = 0; p <= n; p++)); do
			local y_index=$((((p + 1) % (n + 1)) * 4))
			local y=0
			for ((j = 0; j < 4 && $((y_index + j)) < ${#data_array[@]}; j++)); do
				y=$((y | (data_array[$((y_index + j))] << (8 * j))))
			done
			local m=$(((z >> 5) ^ (y << 2)))
			m=$((m + ((y >> 3) ^ (z << 4) ^ (d ^ y))))
			m=$((m + (key_array[$(((p & 3) ^ e))] ^ z)))
			m_index=$((p * 4))
			local temp_m=$((data_array[$m_index] | (data_array[$((m_index + 1))] << 8) | (data_array[$((m_index + 2))] << 16) | (data_array[$((m_index + 3))] << 24)))
			m=$((m + temp_m))
			m=$((m & 0xFFFFFFFF))
			data_array[$((m_index + 0))]=$((m & 0xFF))
			data_array[$((m_index + 1))]=$(((m >> 8) & 0xFF))
			data_array[$((m_index + 2))]=$(((m >> 16) & 0xFF))
			data_array[$((m_index + 3))]=$(((m >> 24) & 0xFF))
			z=$m
		done
	done
	for ((i = 0; i < ${#data_array[@]}; i++)); do
		local hex=$(printf "%x" ${data_array[$i]})
		echo -e -n "\x$hex"
	done
}

tea() {
	fill_key "$1"
	fill_data "$2"
	encode
}

log_date() {
	echo "[$(date $date_format)]"
}

config_log() {
	if [ $verbose == 1 ]; then
		$LOG_LEVEL="debug"
	fi
}

log_error() {
	if [ $LOG_LEVEL == "info" ] || [ $LOG_LEVEL == "debug" ] ||
		[ $LOG_LEVEL == "error" ]; then
		echo "$(log_date) ERROR $1" >&2
	fi
}

log_debug() {
	[ $LOG_LEVEL == "debug" ] && echo -e "$(log_date) DEBUG $1" >&2
}

log_info() {
	if [ $LOG_LEVEL == "info" ] || [ $LOG_LEVEL == "debug" ]; then
		echo -e "$(log_date) INFO $1" >&2
	fi
}

check_user() {
	USERNAME=${USERNAME:-"$TUNET_USERNAME"}
	if [ -z $USERNAME ]; then
		log_error "TUNET_USERNAME is not set"
		exit 1
	fi
}

check_pass() {
	PASSNAME=${PASSNAME:-"$TUNET_PASSNAME"}
	PASSWORD=${PASSWORD:-"$TUNET_PASSWORD"}
	if [ -z $PASSNAME ]; then
		if [ -z $PASSWORD ]; then
			log_error "TUNET_PASSWORD is not set"
			exit 1
		fi
	else
		PASSWORD="$(pass show "$PASSNAME")"
	fi
}

fetch_ac_id() {
	log_debug "fetch ac_id"
	local res=$(curl -s $REDIRECT_URL)
	[[ $res =~ $REGEX_AC_ID ]]
	local ac_id=${BASH_REMATCH[1]}
	if [ -z $ac_id ]; then
		log_debug "ac_id not found, using 1 as default"
		echo "1"
	else
		log_debug "get ac_id from $REDIRECT_URL: $ac_id"
		echo $ac_id
	fi
}

fetch_challenge() {
	log_debug "fetch challenge v$ipv"
	local AUTH_CHALLENGE_URL=$([ $ipv == 6 ] && echo $AUTH6_CHALLENGE_URL || echo $AUTH4_CHALLENGE_URL)
	local res=$(curl -s $AUTH_CHALLENGE_URL --data-urlencode "username=$USERNAME" --data-urlencode "double_stack=1" --data-urlencode "ip=" --data-urlencode "callback=callback")
	local REGEX_CHALLENGE='"challenge":"([^"]+)"'
	[[ $res =~ $REGEX_CHALLENGE ]]
	local challenge=${BASH_REMATCH[1]}
	echo $challenge
}

gen_hmacmd5() {
	echo -n $1 | openssl dgst -md5 -hmac "" -r | cut -d ' ' -f 1
}

post_info() {
	local challenge=$1
	local json="{\"acid\":\"$2\",\"enc_ver\":\"srun_bx1\",\"ip\":\"\",\"password\":\"$PASSWORD\",\"username\":\"$USERNAME\"}"
	local data=$(echo -n $json | sed 's/ //g' | sed 's/"acid":"\([0-9]\+\)"/"acid":\1/g')
	echo $(tea $challenge $data | base64 | tr -d '\n' | tr \
		'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/' \
		'LVoJPiCN2R8G90yg+hmFHuacZ1OWMnrsSTXkYpUq/3dlbfKwv6xztjI7DeBE45QA')
}

login() {
	[ -f "$CACHE_DIR/passwd" ] && source "$CACHE_DIR/passwd"
	check_user
	check_pass
	log_debug "begin auth$ipv login"
	local ac_id=$(fetch_ac_id)
	local challenge=$(fetch_challenge "$ipv")
	log_debug "challenge: $challenge"
	local info="{SRBX1}$(post_info $challenge $ac_id)"
	log_debug "info: $info"
	local password_md5=$(gen_hmacmd5 $challenge)
	log_debug "password_md5: {MD5}$password_md5"
	local checksum="$challenge$USERNAME$challenge$password_md5$challenge$ac_id$challenge${challenge}200${challenge}1$challenge$info"
	checksum=$(echo -n $checksum | openssl sha1 -hex -r | cut -d ' ' -f 1)
	log_debug "checksum: $checksum"
	log_debug "make login request"
	local AUTH_LOG_URL=$([ $ipv == 6 ] && echo $AUTH6_LOG_URL || echo $AUTH4_LOG_URL)
	local response=$(curl -s -X POST $AUTH_LOG_URL \
		-H "Content-Type: application/x-www-form-urlencoded" \
		--data-urlencode "action=login" \
		--data-urlencode "ac_id=$ac_id" \
		--data-urlencode "double_stack=1" \
		--data-urlencode "n=200" \
		--data-urlencode "type=1" \
		--data-urlencode "username=$USERNAME" \
		--data-urlencode "password={MD5}$password_md5" \
		--data-urlencode "info=$info" \
		--data-urlencode "chksum=$checksum" \
		--data-urlencode "callback=callback")
	log_debug "post $AUTH_LOG_URL: $response"
	local REGEX_SUC_MSG='"suc_msg":"([^"]+)"'
	[[ $response =~ $REGEX_SUC_MSG ]]
	local suc_msg=${BASH_REMATCH[1]}
	if [ "$suc_msg" != "login_ok" ]; then
		if [ -z $suc_msg ]; then
			local REGEX_ERR_MSG='"error_msg":"([^"]+)"'
			[[ $response =~ $REGEX_ERR_MSG ]]
			local err_msg=${BASH_REMATCH[1]}
			log_error "$err_msg"
		else
			log_error "$suc_msg"
		fi
		exit 1
	else
		log_info "$suc_msg"
		exit 0
	fi
}

logout() {
	[ -f "$CACHE_DIR/passwd" ] && source "$CACHE_DIR/passwd"
	check_user
	log_debug "begin auth$ipv logout"
	local AUTH_LOG_URL=$([ $ipv == 6 ] && echo $AUTH6_LOG_URL || echo $AUTH4_LOG_URL)
	local response=$(curl -s -X POST $AUTH_LOG_URL \
		-H "Content-Type: application/x-www-form-urlencoded" \
		--data-urlencode "action=logout" \
		--data-urlencode "ac_id=1" \
		--data-urlencode "double_stack=1" \
		--data-urlencode "username=$USERNAME" \
		--data-urlencode "callback=callback")
	log_debug "post $AUTH_LOG_URL: $response"
	local REGEX_SUC_MSG='"error":"([^"]+)"'
	[[ $response =~ $REGEX_SUC_MSG ]]
	local suc_msg=${BASH_REMATCH[1]}
	if [ "$suc_msg" != "ok" ]; then
		local REGEX_ERROR_MSG='"error_msg":"([^"]+)"'
		[[ $response =~ $REGEX_ERROR_MSG ]]
		local error_msg=${BASH_REMATCH[1]}
		log_error "$error_msg"
		exit 1
	else
		log_info "$suc_msg"
		exit 0
	fi
}

whoami() {
	local res=$(curl -s $TUNET_USER_INFO_URL)
	log_debug "get $TUNET_USER_INFO_URL: $res"
	local cnt=$(echo $res | tr ',' '\n' | wc -l)
	if [ $cnt != 22 ]; then
		log_error "expect 22 fields, get $cnt: $res"
		exit 1
	else
		local user=$(echo $res | cut -d ',' -f1)
		if [ $verbose -eq 1 ]; then
			local login=$(echo $res | cut -d ',' -f2)
			local online=$(echo $res | cut -d ',' -f3)
			local online=$((online - login))
			local online=$(awk "BEGIN {printf \"%.2f\n\", $online / 3600}")
			local login=$(date -d "@$login" --rfc-3339 s)
			local in=$(echo $res | cut -d ',' -f4)
			local out=$(echo $res | cut -d ',' -f5)
			local tot=$(echo $res | cut -d ',' -f7)
			local sum=$((in + out))
			local in=$(awk "BEGIN {printf \"%.2f\n\", $in / 1048576}")
			local out=$(awk "BEGIN {printf \"%.2f\n\", $out / 1048576}")
			local sum=$(awk "BEGIN {printf \"%.2f\n\", $sum / 1048576}")
			local tot=$(awk "BEGIN {printf \"%.2f\n\", $tot / 1073741824}")
			local ip=$(echo $res | cut -d ',' -f9)
			local res=$(curl -s $TUNET_USER_INFO_JSON_URL)
			[[ $res =~ $REGEX_USER_INFO_JSON ]]
			log_debug "get $TUNET_USER_INFO_JSON_URL: $res"
			local billing_name=${BASH_REMATCH[1]}
			local device=${BASH_REMATCH[2]}
			local products_name=${BASH_REMATCH[3]}
			local sysver=${BASH_REMATCH[4]}
			local balance=${BASH_REMATCH[5]}
			local mac=${BASH_REMATCH[6]}
			local mac=$(echo -n "$mac" | tr -- '-ABCDEF' ':abcdef')
			local label_width=18
			printf "%-${label_width}s %s\n" "Username:" "$user"
			printf "%-${label_width}s %s\n" "Login Time:" "$login"
			printf "%-${label_width}s %s h\n" "Age:" "$online"
			printf "%-${label_width}s %s\n" "Billing Name:" "$billing_name"
			printf "%-${label_width}s %s\n" "Products Name:" "$products_name"
			printf "%-${label_width}s %s\n" "Device Online:" "$device"
			printf "%-${label_width}s %s CNY\n" "User Balance:" "$balance"
			printf "%-${label_width}s %s Mi\n" "Traffic In:" "$in"
			printf "%-${label_width}s %s Mi\n" "Traffic Out:" "$out"
			printf "%-${label_width}s %s Mi\n" "Traffic Sum:" "$sum"
			printf "%-${label_width}s %s Gi\n" "Traffic Total:" "$tot"
			printf "%-${label_width}s %s\n" "MAC Address:" "$mac"
			printf "%-${label_width}s %s\n" "IP Address:" "$ip"
			if command -v jq >/dev/null 2>&1; then
				local res="${res:4:-1}"
				local device_detail=$(echo "$res" | jq -r '.online_device_detail // empty' 2>/dev/null)
				if [ -n "$device_detail" ] && [ "$device_detail" != "null" ]; then
					echo
					printf "%-${label_width}s\n" "Device Details:"
					local device_num=1
					echo "$device_detail" | jq -r 'to_entries[] | "\(.key)|\(.value.ip)|\(.value.ip6)|\(.value.class_name)|\(.value.os_name)"' 2>/dev/null | while IFS='|' read -r device_id device_ip device_ip6 device_class device_os; do
						[ -z "$device_id" ] && continue
						local device_indent=$(printf "%*s" 2 "")
						local field_indent=$(printf "%*s" 4 "")
						local field_width=$((label_width - 4))
						printf "${device_indent}Device %d:\n" "$device_num"
						printf "${field_indent}%-${field_width}s %s\n" "Rad Online ID:" "$device_id"
						printf "${field_indent}%-${field_width}s %s\n" "IPv4:" "${device_ip:-N/A}"
						printf "${field_indent}%-${field_width}s %s\n" "IPv6:" "${device_ip6:-N/A}"
						[ -n "$device_class" ] && [ "$device_class" != "" ] && printf "${field_indent}%-${field_width}s %s\n" "Class Name:" "$device_class"
						[ -n "$device_os" ] && [ "$device_os" != "" ] && printf "${field_indent}%-${field_width}s %s\n" "OS Name:" "$device_os"
						echo
						((device_num++))
					done
				else
					echo
					printf "%-${label_width}s %s\n" "Device Details:" "${dim}No details available${reset}"
				fi
			else
				log_debug "jq not found, skipping device details parsing"
			fi
			printf "%-${label_width}s %s\n" "System Version:" "$sysver"
		else
			log_info "$user"
		fi
		exit 0
	fi
}

help() {
	man 1 "$NAME"
}

config() {
	while [[ -z $USERNAME ]]; do
		read -p "username: " USERNAME
	done
	echo "export TUNET_USERNAME=$USERNAME" >$CACHE_DIR/passwd
	if [[ $use_passname = "yes" ]]; then
		while [[ -z $PASSNAME ]]; do
			read -p "passname: " PASSNAME
		done
		echo "export TUNET_PASSNAME=$PASSNAME" >>$CACHE_DIR/passwd
	else

		while [[ -z $PASSWORD ]]; do
			read -s -p "password: " PASSWORD
			echo
		done
		echo "export TUNET_PASSWORD=$PASSWORD" >>$CACHE_DIR/passwd
	fi
	chmod 600 $CACHE_DIR/passwd
}

preprocess_args() {
	local args=("$@")
	new_args=()
	local i=0
	while ((i < ${#args[@]})); do
		local arg="${args[i]}"
		case "$arg" in
			--)
				new_args+=("--")
				((i++))
				;;
			--*)
				new_args+=("$arg")
				((i++))
				;;
			-?*)
				local opt_str="${arg:1}"
				while [[ -n "$opt_str" ]]; do
					local opt_char="${opt_str:0:1}"
					new_args+=("-${opt_char}")
					opt_str="${opt_str:1}"
				done
				((i++))
				;;
			*)
				new_args+=("$arg")
				((i++))
				;;
		esac
	done
}

preprocess_args "$@"

set -- "${new_args[@]}"
mkdir -p "$CACHE_DIR"

while [[ $# -gt 0 ]]; do
	case "$1" in
		-h | --help)
			op="help"
			shift
			;;
		-c | --config)
			op="config"
			shift
			;;
		-i | --login)
			op="login"
			shift
			;;
		-o | --logout)
			op="logout"
			shift
			;;
		-w | --whoami)
			op="whoami"
			shift
			;;
		-v | --verbose)
			verbose=1
			shift
			;;
		-a | --auth)
			ipv="$2"
			shift 2
			;;
		--date-format)
			date_format="$2"
			shift 2
			;;
		--version)
			echo "$NAME $VERSION"
			exit 0
			;;
		--pass)
			use_passname="yes"
			shift
			;;
		--)
			shift
			break
			;;
		*)
			echo "Unknown option: $1" >&2
			exit 1
			;;
	esac
done

if [ "$ipv" = "auto" ]; then
	REGEX_IPV='//auth([46])\.tsinghua\.edu\.cn'
	res=$(curl -s $REDIRECT_URL)
	[[ $res =~ $REGEX_IPV ]]
	ipv=${BASH_REMATCH[1]}
	if [ -z $ipv ]; then
		ipv=4
	fi
	log_debug "ipv: $ipv"
fi

if [[ "$ipv" != "4" ]] && [[ $ipv != "6" ]]; then
	echo "Unknown auth method: $ipv" >&2
	exit 1
fi

case $op in
	help)
		help
		;;
	config)
		config
		;;
	whoami)
		whoami
		;;
	login)
		config_log
		login
		;;
	logout)
		config_log
		logout
		;;
esac
