#!/bin/sh
##########################
# History:
#	2017/08/01 - [Tao Wu] Create file
#
# Copyright (c) 2017 Ambarella International LP
#
# This file and its contents ("Software") are protected by intellectual
# property rights including, without limitation, U.S. and/or foreign
# copyrights. This Software is also the confidential and proprietary
# information of Ambarella International LP and its licensors. You may not use, reproduce,
# disclose, distribute, modify, or otherwise prepare derivative works of this
# Software or any portion thereof except pursuant to a signed license agreement
# or nondisclosure agreement with Ambarella International LP or its authorized affiliates.
# In the absence of such an agreement, you agree to promptly notify and return
# this Software to Ambarella International LP.
#
# This file includes sample code and is only for internal testing and evaluation.  If you
# distribute this sample code (whether in source, object, or binary code form), it will be
# without any warranty or indemnity protection from Ambarella International LP or its affiliates.
#
# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
# INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF NON-INFRINGEMENT,
# MERCHANTABILITY, AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL AMBARELLA INTERNATIONAL LP OR ITS AFFILIATES BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; COMPUTER FAILURE OR MALFUNCTION; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
###############

WIFI_STRESS_VERSION="0.1.2"
mode=$1
pm=$2

SSID=$3
PASSWORD=$4

SVR_IP=$5
SVR_PORT=$6

KA_ITRVL=$7
DTIM_ITRVL=$8

RECOVERY_AP_CONNECT=$9

############# Configuration it ###############

DEVICE=wlan0
MODULE_ID=UNKNOWN
HIB_STATE_CONF=/tmp/hib_state
DEBUG_WIFI_LOG=0
############# Generate config ###############

## 0: AP6234, 1: AP6212A, after etc.
## 10: SD8977
CHIP_ID=1

############# Exit Error Number ###############
ERRNO_OK=0
ERRNO_PARAM=1
ERRNO_ENV=2
#ERRNO_SSID_NOT_FOUND=3
#ERRNO_PASSWORD_WRONG=4
#ERRNO_SSID_REPEATED=5

#############  Basic Function ##################
usages()
{
	echo "Version: ${WIFI_STRESS_VERSION}"
	echo "This script used to stress test WiFi KeepAlive and Wakeup function with wowl tool"
	echo "usage:                  $0 [mem|disk] [shut|off|no] [ssid] [password|0] [server_ip] [server_port] [ka_interval] [dtim_interval] [recovery_if_dissconnect]"
	echo ""
	echo "Example:"
	echo "Autorun in DDR-Refresh: # $0 mem off AMBA-taowu 0 10.0.0.4 7877 10 600 0"
	echo "Autorun in Hibernation: # $0 disk shut AMBA-taowu 0 10.0.0.4 7877 10 600 0"

	echo ""
	echo "NOTE: AMBA-taowu is SSID, insert 0 if open."

	echo ""
	echo "Server CMD:"
	echo "1. Longtime Test (KeepAlive)  $ net_socket -s -p 7877 -r -e"
	echo "2. Stress Test (Sleep/Wakeup) $ net_socket -s -p 7877 -a 22/9999/0/0 -e"
	exit $ERRNO_OK
}

do_cmd()
{
	local command="$1"
	echo "CMD=[${command}]"
	${command}
}

kill_apps()
{
	#the max times to kill app
	local KILL_NUM_MAX=10

	for app in "$@"
	do
		local kill_num=0
		while [ "$(pgrep "${app}")" != "" ]
		do
			if [ $kill_num -ge $KILL_NUM_MAX ]; then
				echo "Please try execute \"killall ${app}\" by yourself"
				exit $ERRNO_ENV
			else
				killall -9 "${app}"
				sleep 1
			fi
			kill_num=$((kill_num + 1));
		done
	done
}

################# GPIO Function #######################

# sysfs_gpio():
# Function Parameters: (id ; dir [out|in]; val)
# Funciton Return: return val if dir equal in
sysfs_gpio()
{
	id=$1
	dir=$2
	val=$3

	echo "${id}" > /sys/class/gpio/export

	if [ "${dir}" = "out" ]; then
		echo out > /sys/class/gpio/gpio"${id}"/direction
		echo "${val}" > /sys/class/gpio/gpio"${id}"/value
	elif [ "${dir}" = "in" ]; then
		echo in > /sys/class/gpio/gpio"${id}"/direction
		val=$(cat /sys/class/gpio/gpio"${id}"/value)
	else
		echo "Invalid direction ${dir}, choose out or in"
	fi

	echo "${id}" > /sys/class/gpio/unexport

	if [ "${dir}" = "in" ]; then
		echo "Read GPIO[${id}] value [${val}]"
		return "${val}"
	fi
}

################ System Env #####################

check_env()
{
	if [ "$(which wowl)" = "" ]; then
		echo "Not Found wowl tools"
		exit $ERRNO_ENV
	fi

	if [ "$(ifconfig ${DEVICE})" = "" ]; then
		echo "Not Found Interface ${DEVICE}"
		exit $ERRNO_ENV
	fi

	if [ -f /etc/ambarella.conf ]; then
		. /etc/ambarella.conf
	fi

	if [ -e /sys/module/bcmdhd ]; then
		if [ -e /lib/firmware/broadcom ] ; then
			MODULE_ID=BCM
			if [ -e /lib/firmware/broadcom/ap6234_ps ] ; then
				CHIP_ID=0
			else
				CHIP_ID=1
			fi
		elif [ -e /lib/firmware/cypress ] ; then
			MODULE_ID=CYW
			if [ -e /lib/firmware/cypress/cyw43340 ] ; then
				CHIP_ID=0
			else
				CHIP_ID=1
			fi
		else
			MODULE_ID=BCM
			CHIP_ID=1
		fi
	elif [ -e /sys/module/mlan ]; then
		MODULE_ID=MRVL
		CHIP_ID=10
	fi

	if [ "${MODULE_ID}" = "BCM" ]; then
		if [ "$(which wl)" = "" ]; then
			echo "Not Found wl tools"
			exit $ERRNO_ENV
		fi
		if [ "$(which wl_hostsleep)" = "" ]; then
			echo "Not Found wl_hostsleep tools"
			exit $ERRNO_ENV
		fi
	elif [ "${MODULE_ID}" = "CYW" ]; then
		if [ "$(which wl)" = "" ]; then
			echo "Not Found wl tools"
			exit $ERRNO_ENV
		fi
	elif [ "${MODULE_ID}" = "MRVL" ]; then
		if [ "$(which mlanutl)" = "" ]; then
			echo "Not Found mlanutl tools"
			exit $ERRNO_ENV
		fi
	fi

	echo "MODULE_ID=${MODULE_ID}, CHIP_ID=${CHIP_ID}"
}

################ System PM #####################

do_hibernation()
{
	echo disk > /sys/power/state
}

do_dram_refresh()
{
	echo mem > /sys/power/state
}

do_poweroff()
{
	do_cmd "poweroff -f"
}

do_shutdown()
{
	if [ "$SYS_BOARD_BSP" = "s2lm_elektra" ]; then
		sysfs_gpio 112 out 1
	elif [ "$SYS_BOARD_BSP" = "s3lm_elektra" ]; then
		sysfs_gpio 26 out 1
		sysfs_gpio 26 out 0
	elif [ "$SYS_BOARD_BSP" = "s5lm_elektra_v" ]; then
		sysfs_gpio 61 out 1
		sysfs_gpio 61 out 0
	else
		echo "Do not support shutdown on $SYS_BOARD_BSP"
		exit $ERRNO_ENV
	fi
}

do_sleep_cpu()
{
	if [ -f ${HIB_STATE_CONF} ]; then
		if [ "${pm}" = "shut" ]; then
			do_shutdown
		elif [ "${pm}" = "off" ]; then
			do_poweroff
		else
			echo "Do nothing, please power off your HW by manully"
		fi
	else
		touch ${HIB_STATE_CONF}
		do_hibernation
	fi
}

################ Connect AP for Broadcom/Marvell #####################

connect_ap()
{
	echo "connect_ap start"

	if [ "${MODULE_ID}" = "MRVL" ]; then
		if [ "$(which ml_setup.sh)" = "" ]; then
			echo "Not Found ml_setup.sh tools"
			exit $ERRNO_ENV
		fi
		if [ ${PASSWORD} -eq 0 ]; then
			ml_setup.sh sta "${SSID}"
		else
			ml_setup.sh sta "${SSID}" "${PASSWORD}"
		fi

	elif [ "${MODULE_ID}" = "BCM" ]; then
		if [ "$(which wl_setup.sh)" = "" ]; then
			echo "Not Found wl_setup.sh tools"
			exit $ERRNO_ENV
		fi
		if [ ${PASSWORD} -eq 0 ];then
			wl_setup.sh sta "${SSID}"
		else
			wl_setup.sh sta "${SSID}" "${PASSWORD}"
		fi
	elif [ "${MODULE_ID}" = "CYW" ]; then
		if [ "$(which wl_setup.sh)" = "" ]; then
			echo "Not Found wl_setup.sh tools"
			exit $ERRNO_ENV
		fi
		if [ ${PASSWORD} -eq 0 ];then
			wifi_setup.sh sta nl80211 "${SSID}"
		else
			wifi_setup.sh sta nl80211 "${SSID}" "${PASSWORD}"
		fi
	fi

	ping "${SVR_IP}" -c 2
	if [ $? -ne 0 ]; then
		echo "Dissconnect with Server"
		exit $ERRNO_ENV
	fi

	if [ $DEBUG_WIFI_LOG -eq 1 ]; then
		echo 8 > /proc/sys/kernel/printk
		if [ "${MODULE_ID}" = "BCM" ]; then
			echo 20 > /sys/module/bcmdhd/parameters/dhd_console_ms
		fi
	fi

	echo "connect_ap end"
}

################ Hibernation for Broadcom/Marvell #####################

prepare_before_hib()
{
	echo "prepare_before_hib start"

	if [ "${MODULE_ID}" = "MRVL" ]; then
		#do_cmd "mlanutl ${DEVICE} hostcmd /lib/firmware/mrvl/sdio_pulldown.conf sdio_pulldown_set"
		do_cmd "mlanutl ${DEVICE} fast_link_sync_enable 1"
		do_cmd "ifconfig ${DEVICE} down"
		do_cmd "rmmod sd8xxx"
		do_cmd "rmmod mlan"
	elif [ "${MODULE_ID}" = "BCM" ]; then
		do_cmd "wl_hostsleep 1"
	fi

	echo "prepare_before_hib end"
}

resume_from_hib()
{
	echo "resume_from_hib start"

	if [ "${MODULE_ID}" = "MRVL" ]; then
		do_cmd "modprobe sd8xxx"
		do_cmd "ifconfig ${DEVICE} up"
		#do_cmd "mlanutl ${DEVICE} pscfg 28 1"
		do_cmd "mlanutl ${DEVICE} cloud_keep_alive /lib/firmware/mrvl/ka_demo.conf stop"
		do_cmd "mlanutl ${DEVICE} wakeupreason"

		ml_status=$(grep media_state /proc/mwlan/${DEVICE}/info)
		ml_status_value=${ml_status#*media_state=}
		ml_status_value=$(echo "${ml_status_value}" | sed -r 's/.*"(.+)".*/\1/')
		#echo ${ml_status_value}
		if [ "${ml_status_value}" = "Connected" ]; then
			do_cmd "udhcpc -i${DEVICE} -A 1 -b"
		else
			echo "WiFi status wrong: ${ml_status_value}"
			if [ ${RECOVERY_AP_CONNECT} -eq 1 ]; then
				echo "Recovery connection with AP"
				connect_ap
			else
				echo "Dissconnect with AP"
				exit $ERRNO_ENV
			fi
		fi
	elif [ "${MODULE_ID}" = "BCM" ]; then
		do_cmd "wl_hostsleep 0"
		do_cmd "wl PM 2"
		do_cmd "wl bcn_li_dtim 1"
		do_cmd "wl tcpka_conn_del 1"

		wakeind=$(wl wowl_wakeind)

		do_cmd "wl wowl_wakeind clear"
		do_cmd "wl wowl_clear"
		do_cmd "wl wowl_pattern clr"

		local need_reconnect_ap=1
		echo "WiFi wakeup indication="${wakeind}""

		if [ "${wakeind}" != "No wakeup indication set" ]; then
			if [ "$(echo "${wakeind}" | grep "Netpattern")" ]; then
				echo "Include wakeup netpattern"
				need_reconnect_ap=0
			fi
			if [ "$(echo "${wakeind}" | grep "TCP")" ]; then
				echo "Include Keepalive Data, Keepalive Timeout, FIN"
				need_reconnect_ap=0
			fi
			if [ ${need_reconnect_ap} -eq 1 ]; then
				if [ ${RECOVERY_AP_CONNECT} -eq 1 ]; then
					echo "Recovery connection with AP"
					connect_ap
				else
					echo "Dissconnect with AP"
					exit $ERRNO_ENV
				fi
			fi
		fi
	elif [ "${MODULE_ID}" = "CYW" ]; then
		##do_cmd "wl_hostsleep 0"
		do_cmd "wl PM 2"
		do_cmd "wl bcn_li_dtim 1"
		do_cmd "wl tcpka_conn_del 1"

		wakeind=$(wl wowl_wakeind)

		do_cmd "wl wowl_wakeind clear"
		do_cmd "wl wowl_clear"
		do_cmd "wl wowl_pattern clr"

		local need_reconnect_ap=1
		echo "WiFi wakeup indication=${wakeind}"

		if [ "${wakeind}" != "No wakeup indication set" ]; then
			if [ $(echo "${wakeind}" | grep "Netpattern") ]; then
				echo "Include wakeup netpattern"
				need_reconnect_ap=0
			fi
			if [ $(echo "${wakeind}" | grep "TCP" ) ]; then
				echo "Include Keepalive Data, Keepalive Timeout, FIN"
				need_reconnect_ap=0
			fi
			if [ ${need_reconnect_ap} -eq 1 ]; then
				if [ ${RECOVERY_AP_CONNECT} -eq 1 ]; then
					echo "Recovery connection with AP"
					connect_ap
				else
					echo "Dissconnect with AP"
					exit $ERRNO_ENV
				fi
			fi
		fi
	fi

	echo "resume_from_hib end"
}

hib_tcpka()
{
	echo "hib_tcpka start"

	ping "${SVR_IP}" -c 2
	if [ $? -ne 0 ]; then
		echo "Dissconnect with Server"
		exit $ERRNO_ENV
	fi

	retry_cnt=0
	wowl_exec_ok=0
	while [ $wowl_exec_ok -eq 0 ]
	do
		if [ "${MODULE_ID}" = "MRVL" ]; then
			local l_ka_itrvl=$((KA_ITRVL*1000))
			wowl -w ${CHIP_ID} -i ${DEVICE} -c "${SVR_IP}" -p "${SVR_PORT}" -t -s -m "${l_ka_itrvl}" -d "${DTIM_ITRVL}" -n -r -a -P 1 -q &
			wowl_exec_ok=1

		elif [ "${MODULE_ID}" = "BCM" ]; then
			wowl -w ${CHIP_ID} -i ${DEVICE} -c "${SVR_IP}" -p "${SVR_PORT}" -t -s -m "${KA_ITRVL}" -d "${DTIM_ITRVL}" -n -a -P 1 -q &

			## Used for check if wowl setup ok.
			sleep 2
			wl_hostsleep 1
			if [ $? -eq 0 ]; then
				wl_hostsleep 0
				retry_cnt=$((retry_cnt+1))
				echo "Setup wowl fail count=$retry_cnt, retry again after sleep 60s."
				killall -9 wowl
				sleep 60
			else
				wowl_exec_ok=1
			fi
		elif [ "${MODULE_ID}" = "CYW" ]; then
			wowl -w ${CHIP_ID} -i ${DEVICE} -c "${SVR_IP}" -p "${SVR_PORT}" -t -s -m "${KA_ITRVL}" -d "${DTIM_ITRVL}" -n -x -a -P 1 &
			wowl_exec_ok=1
		fi
	done

	echo "hib_tcpka end"
}

prepare_before_poweroff()
{
	echo "prepare_before_poweroff start"

	if [ "${MODULE_ID}" = "MRVL" ]; then
		do_cmd "ifconfig ${DEVICE} down"
		do_cmd "rmmod sd8xxx"
		do_cmd "rmmod mlan"
	fi

	echo "prepare_before_poweroff end"
}

################ DRAM Refresh for Broadcom/Marvell #####################

mem_tcpka()
{
	echo "mem_tcpka start"

	ping "${SVR_IP}" -c 2
	if [ $? -ne 0 ]; then
		echo "Dissconnect with Server"
		exit $ERRNO_ENV
	fi

	if [ "${MODULE_ID}" = "MRVL" ]; then
		local l_ka_itrvl=$((KA_ITRVL*1000))
		wowl -w ${CHIP_ID} -i ${DEVICE} -c "${SVR_IP}" -p "${SVR_PORT}" -t -s -m "${l_ka_itrvl}" -d "${DTIM_ITRVL}" -n -a -l -P 1 &
	elif [ "${MODULE_ID}" = "BCM" ]; then
		wowl -w ${CHIP_ID} -i ${DEVICE} -c "${SVR_IP}" -p "${SVR_PORT}" -t -s -m "${KA_ITRVL}" -d "${DTIM_ITRVL}" -e -n -l -o -a -P 1 -q &
	elif [ "${MODULE_ID}" = "CYW" ]; then
		wowl -w ${CHIP_ID} -i ${DEVICE} -c "${SVR_IP}" -p "${SVR_PORT}" -t -s -m "${KA_ITRVL}" -d "${DTIM_ITRVL}" -e -n -l -o -a -P 1 &
	fi

	echo "mem_tcpka end"
}

################ Main Function #####################

## Show usage when parameter num is not equal 4
if [ $# -ne 9 ]; then
	usages
fi

check_env
connect_ap

case ${mode} in
	mem)
		mem_tcpka
		;;
	disk)
		prepare_before_hib
		do_sleep_cpu
		sleep 3

		# loop start
		resume_from_hib
		sleep 5
		hib_tcpka
		sleep 1

		prepare_before_poweroff
		do_sleep_cpu
		sleep 1
		# loop end
		;;
	*)
		echo "Please Select PM Mode [mem|disk]"
		exit $ERRNO_PARAM
		;;
esac
