#!/bin/bash
# Copyright (c) Huawei Technologies Co., Ltd. 2021-2022. All rights reserved.
# Upgrade Mellnaox card firmware
# Usage 1: upgrade all matching devices
# 1. chmod -R +x *
# 2. ./install.sh upgrade
#
# Usage 2: specified device bdf upgrade specified device
# 1. chmod -R +x *
# 2. ./install.sh upgrade BDF "0000:xx:xx.x,0000:yy:yy.y" or "xx:xx.x,yy:yy.y"
# Note: use the bdf value of device port 1,one or more devices can be specified.
set -e
# get work path
path=$(dirname "$0")
path="${path/\./$(pwd)}"
pwd="${path}"
cd "${pwd}"

# delete history files 
if [ -f "${pwd}/work.log" ]; then
	rm -f "${pwd}/work.log" 2>/dev/null
fi

# define device number 
dev_num=0
# define upgrade result flag,0-success,1-failure
flag_fail=0
# define result for result.xml,OK or Fail
result=""
# define error_code for result.xml,0-OK,1-Fail,4-No device,5-no version.xml,6-no tools
error_code=""
description=""

# create result.xml to record the upgrade result
create_result_xml()
{	
	echo "<?xml version=\"1.0\" encoding=\"utf-8\"?>" >"${pwd}/result.xml"
	echo "<ConfigureResult>" >>"${pwd}/result.xml"
	echo "    <FwUpgrade>" >>"${pwd}/result.xml"
	echo "        <Result>${result}</Result>" >>"${pwd}/result.xml"
	echo "        <ErrorCode>${error_code}</ErrorCode>" >>"${pwd}/result.xml"
	echo "        <Description>${description}</Description>" >>"${pwd}/result.xml"
	echo "    </FwUpgrade>" >>"${pwd}/result.xml"
	echo "</ConfigureResult>" >>"${pwd}/result.xml"
}

# read version xml to obtaining upgrade package information
read_version_xml()
{
	if [ ! -f "${pwd}/version.xml" ]; then
		# version xml not found
		result="Fail"
		error_code="5"
		description="version.xml not found"
		create_result_xml
		exit 5
	fi

	# get fw bin file name 
	fw_file=$(cat "${pwd}/version.xml" |grep -i "<FileName>"|awk 'BEGIN{FS=">"}{print $2}'|awk 'BEGIN{FS="<"}{print $1}')	
	# get fw version
	fw_version=$(cat "${pwd}/version.xml" |grep -i "<Version>"|awk 'BEGIN{FS=">"}{print $2}'|awk 'BEGIN{FS="<"}{print $1}')	
	# get fw 4-tuple information
	fw_uid=$(cat "${pwd}/version.xml" |grep -i "<SupportModelUID>"|awk 'BEGIN{FS=">"}{print $2}'|awk 'BEGIN{FS="<"}{print $1}' |sed 's/;/ /g')
	
	# check file name
	if [ $(ls |grep -w ${fw_file}|wc -l) -eq 0 ]; then
		# firmware file error
		result="Fail"
		error_code="5"
		description="FileName Error"
		create_result_xml
		exit 5
	fi
}

# preparing the mst and flint upgrade tools
prepare_tools()
{
	local tool_flag=0
	
    if [ -d "${pwd}/tools/" ]; then       
		cd "${pwd}/tools/"
		chmod +x ${pwd}/tools/*
		# force cp
		/bin/cp -arf ${pwd}/tools/* /usr/bin/ 2>/dev/null
    else
        tool_flag=1
    fi
	
	cd "${pwd}/"
	# failed to prepare the tool. exit the upgrade process.
	if [ ${tool_flag} == 1 ]; then
		result="Fail"
		error_code="6"
		description="Prepare tools failed"
		create_result_xml
		exit 6;
	fi
}

# obtaining Mellanox card information in the environment,and collect the list of cards that match the firmware package.
get_device_info()
{	
	# to avoid the scenario where the mst start has been performed, you need to perform the mst stop operation first.
	if [ $(lsmod|grep -w mst_pci|wc -l) -eq 1 ]; then
		rmmod mst_pci 2>/dev/null
	fi
	if [ $(lsmod|grep -w mst_pciconf|wc -l) -eq 1 ]; then
		rmmod mst_pciconf 2>/dev/null
	fi
	# to avoid the scenario where SP132 does not support the tool version package, forcibly stop the mst command.
	mst stop >work.log 2>/dev/null | true
	# record information
	flint -version >>work.log 2>/dev/null	
	mst status >>work.log 2>/dev/null
	
	# define the number of devices that do not need to be upgraded.
	local fw_dev=0
	# define the number of devices that match the firmware package.
	local num_index=0	
	
	# check whether the 4-tuple information of the device matches the version.xml file. 
	for pci_device in $(mst status |grep ":00"); do
	
		# run the lspci command to query device information
		lspci -s ${pci_device} -xx >baseinfo.txt
		# get the 4-tuple information of the device
		vendor_id=$(sed -n '2p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $3}')$(sed -n '2p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $2}')
		device_id=$(sed -n '2p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $5}')$(sed -n '2p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $4}')
		sub_vid=$(sed -n '4p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $15}')$(sed -n '4p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $14}')
		sub_did=$(sed -n '4p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $17}')$(sed -n '4p' baseinfo.txt |awk 'BEGIN{FS=" "}{print $16}')
		
		# check whether the device matches.
		if [[ ${fw_uid} =~ "0x${vendor_id}.0x${device_id}.0x${sub_vid}.0x${sub_did}" ]]; then	
			let num_index=${num_index}+1
		else
			# the card does not match the firmware package. upgrade is not supported. skip this step.
			continue
		fi			
		
		# obtains the firmware version information of the matching device. 
		flint -d ${pci_device} q >dev_info.txt
		# synchronizing Information to work.log
		echo -e >>work.log
		cat baseinfo.txt >>work.log
		cat dev_info.txt >>work.log
		# check whether the current firmware version of the device is consistent with the target version
		cur_fw="$(cat dev_info.txt |grep -i "FW Version"|awk 'BEGIN{FS=" "}{print $3}')"
		if [ "${cur_fw}" == "${fw_version}" ]; then
			let fw_dev=${fw_dev}+1
			# the firmware version is the same as the target version and does not need to be upgraded.
			continue
		fi
		
		# need upgrade device list		
		device_index[${dev_num}]="${pci_device}"
		# get bdf number
		local bdf=${pci_device}
		bus_num[${dev_num}]=${bdf:0:2}
		device_num[${dev_num}]=${bdf:3:2}
		function_num[${dev_num}]=${bdf:6:1}
		
		#init bdf flag,the default value is 0, indicating that the device need to upgrade
		bdf_select[${dev_num}]=0		
		let dev_num=${dev_num}+1			
	done
	
	# delete temporary files if they exist.
	if [ -f "${pwd}/baseinfo.txt" ]; then
		rm -f "${pwd}/baseinfo.txt" 2>/dev/null
	fi
	
	if [ -f "${pwd}/dev_info.txt" ]; then
		rm -f "${pwd}/dev_info.txt" 2>/dev/null
	fi

	# the current version of all device are the same as the target version,no upgrade
	if [ ${fw_dev} -eq ${num_index} ] && [ ${num_index} != 0 ]; then
		result="OK"
		error_code="0"
		description="already upgrade"
		create_result_xml
		exit 0;
	fi
	
	# no matching device exists. exit the upgrade process.
	if [ ${dev_num} == 0 ]; then
		result="Fail"
		error_code="4"
		description="No available device found"
		create_result_xml
		exit 4;
	fi
}

# upgrade firmware command 
upgrade_by_tools()
{	
	flint -y -d ${device_index[$index]} -i "${pwd}/${fw_file}"  b >"${pwd}/result.txt" | true
	# synchronizing Information to work.log
	echo "${device_index[$index]} upgrade info" >>"${pwd}/work.log"
	cat "${pwd}/result.txt" >>"${pwd}/work.log"
	# checking the upgrade result
    if [ $(cat result.txt|grep "Writing Boot image component "|grep "OK"|wc -l) -ge 1 ] || [ $(cat result.txt|grep "signature"|grep "OK"|wc -l) -ge 2 ]; then
		:
	else
		flag_fail=1
	fi		
}

# upgrade device firmware
upgrade_firmware()
{
	# upgrade in loop
	for ((i=0;i<${dev_num};i++)); do
		index=${i}	
		# if device is not in select bdfs,then skipping upgrade
		if [ ${bdf_select[$index]} == 1 ]; then
			continue
		fi		
		upgrade_by_tools
		# if upgrade failed,stop upgrade,and the result.xml file is recorded.
		if [ ${flag_fail} == 1 ]; then
			# delete tmp file
			if [ -f "${pwd}/result.txt" ]; then
				rm -f "${pwd}/result.txt" 2>/dev/null
			fi
			result="Fail"
			error_code="1"
			description="Upgrade failed"
			create_result_xml
			exit 1;
		fi
	done
	# delete tmp file
	if [ -f "${pwd}/result.txt" ]; then
		rm -f "${pwd}/result.txt" 2>/dev/null
	fi
	# the loop exits normally, and the result.xml file is recorded.
	if [ ${flag_fail} == 0 ]; then
		result="OK"
		error_code="NA"
		description="NA"
		create_result_xml
		exit 0;
	fi
}

# filter devices by bdf
filter_by_bdf()
{	
	# defines the filtering flag. The default value is 0, indicating that the entered bdf does not match.
	flag_filter=0	
	# check bdf number
	for ((i=0;i<${dev_num};i++)); do
		# skip if bdf is empty
		if [ -z "${bus_num[${i}]}" ] && [ -z "${device_num[${i}]}" ] && [ -z "${function_num[${i}]}" ]; then
			continue
		fi
		# check whether the filtered bdf device information is within the range of the card		
		if [[ ${select_bdfs} =~ "${bus_num[${i}]}:${device_num[${i}]}.${function_num[${i}]}" ]]; then 
			# indicates that the bdf needs to be upgraded.
			flag_filter=1
		else 
			# if the mismatch flag is 1, the upgrade is not performed.
			bdf_select[${i}]=1
		fi
	done
	
	# no device found by filter
	if [ ${flag_filter} == 0 ]; then 
		result="Fail"
		error_code="2"
		description="Error input"
		create_result_xml
		exit 2;
	fi
}

# set select bdf
select_bdfs=""
if [ -n "$2" ] && [ "$2" == "BDF" ]; then
	# $3 bdf format: "0000:xx:xx.x,0000:yy:yy.y" or "xx:xx.x,yy:yy.y"
	select_bdfs=$(echo $3 |sed 's/;/ /g')
fi

case $1 in
	upgrade)
		# read version xml
		read_version_xml		
		# preparing the mst and flint upgrade tools
		prepare_tools		
		# get device info by tools
		get_device_info
		# filter bdfs by input
		if [ -n "${select_bdfs}" ]; then
			filter_by_bdf
		fi		
		# upgrade firmware
		upgrade_firmware
		;;
	*)
		;;
esac

exit 0
