#!/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}"
mlnx_log_error_list=("failed to allocate command entry" "mlx5e_netdev_init failed, err" "Waiting for FW initialization, timeout abort in" "firmware CRC error")
smartsp_dmesg_path="/mnt/usb/spforbmc/spinfo/dmesg.log"
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 u+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 | true
		# 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
}

check_flint_status() {
    set +e
    flint_flag=0
    flint -d ${device_index[$index]} q >> "${pwd}/work.log" 2>/dev/null
    if [ $? != 0 ];then
        mst restart
        flint -d ${device_index[$index]} q >> "${pwd}/work.log" 2>/dev/null
        if [ $? != 0 ];then
            flint_flag=1
        fi
    fi
    set -e
    if [ ${flint_flag} == 1 ]; then
            result="Fail"
            error_code="9"
            description="flint tool is abnormal"
            create_result_xml
            exit 9;
    fi
}

check_mellanox_target_version() {
    check_flint_status
    if [ $? != 0 ];then
        return 1
    fi
    flint -d ${device_index[$index]} q > "${pwd}/tmp_flint_out.log"  | tee -a "${pwd}/work.log"
    if [ $? != 0 ];then
        return 1
    fi
    cat ${pwd}/tmp_flint_out.log >> "${pwd}/work.log" 2>/dev/null
    cur_fw="$(cat ${pwd}/tmp_flint_out.log |grep -i 'Fw Version:'|awk 'BEGIN{FS=" "}{print $3}')"
    if [ -f ${pwd}/tmp_flint_out.log ];then
        rm -f ${pwd}/tmp_flint_out.log
    fi
    if [ "$cur_fw" != "$fw_version" ];then
        return 1
    fi
    
    return 0
}

# 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
    check_mellanox_target_version
    if [ $? != 0 ];then
        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
}

check_dmesg_mellanox_card_status() {
    check_mlnx_result=0
    check_bdf_list=()
    error_bdf=""
    if [ -n "$1" ];then
        check_bdf_list+=("$1")
    else
    	local jx=0;
        for ((jx=0;jx<${dev_num};jx++))
        do
            check_bdf_list+=("${device_index[$jx]}")
        done
    fi
    
    #add "echo $(cmd) 2>/dev/null" to avoid exiting this script when dmesg no mlx5_core 
    echo $(dmesg -T | grep -i mlx5_core  > ${pwd}/mlnx_dmesg_check.log)  2>/dev/null
    if [ -f ${smartsp_dmesg_path} ]; then
        echo $(cat ${smartsp_dmesg_path} | grep -i mlx5_core >> ${pwd}/mlnx_dmesg_check.log)  2>/dev/null
    fi
    
    for check_bdf in "${check_bdf_list[@]}"
    do
        #delete last num of bdf, example: 81:00.0 -> 81:00
        check_bdf=${check_bdf%.*}
        for mlnx_error in "${mlnx_log_error_list[@]}"
        do
            if [ "`cat ${pwd}/mlnx_dmesg_check.log 2>/dev/null |grep ${check_bdf} |grep "${mlnx_error}" | wc -l`" != "0" ];then
                cat ${pwd}/mlnx_dmesg_check.log 2>/dev/null |grep ${check_bdf} |grep "${mlnx_error}" >> "${pwd}/work.log" 2>/dev/null
                error_bdf=$(cat ${pwd}/mlnx_dmesg_check.log |grep ${check_bdf} |grep "${mlnx_error}"|awk 'BEGIN{FS="]"}{print $2}'|awk 'BEGIN{FS=" "}{print $2}')
                check_mlnx_result=1
            fi
        done
    done
    
    if [ -f ${pwd}/mlnx_dmesg_check.log ];then
        rm -f ${pwd}/mlnx_dmesg_check.log
    fi
    
    cd "${pwd}/"
    if [ ${check_mlnx_result} != 0 ]
    then
        #prepare tools failed
        result="Fail"
        error_code="8"
        description="Mellanox card bdf:${error_bdf} check firmware staus failed!"
        create_result_xml
        exit 8;
    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		

		#check mellanox card in dmesg status is abnormal
		check_dmesg_mellanox_card_status ${select_bdfs}
		# upgrade firmware
		upgrade_firmware
		;;
	help|H|h)
		echo "Usage: ./install.sh upgrade BDF <bdf>"
		echo "Examples:"
		echo "    ./install.sh upgrade"
		echo "    ./install.sh upgrade BDF 01:00.0"
		;;
	*)
		;;
esac

exit 0
