#!/bin/bash
###############################################################################
# Copyright 2020 IBM Corp.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
#                            #
###############################################################################
# COMPONENT: refresh_bootmap                                                  #
#                                                                             #
# Refresh bootmap info of the specific device, the device can be FCP, eckd,   #
# etc. Currently, the script support FCP device only.                         #
###############################################################################

source /opt/zthin/lib/zthinshellutils

###############################################################################
### FUNCTIONS #################################################################
###############################################################################

function printCMDDescription {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Prints a short description of this command.
  # @Overrides:
  #   printCMDDescription{} in "zthinshellutils".
  # @Code:
  echo -n "Refresh bootmap info of the specific device. "
  echo    "Currently, the script only support FCP device."
} #printCMDDescription{}

function printCMDUsage {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Prints a short description of this command.
  # @Overrides:
  #   printCMDUsage{} in "zthinshellutils".
  # @Code:
  echo -n "Usage: $CMDNAME OPTIONS"
  [[ $namedArgListing ]] && echo -n " ${namedArgListing}" |
                            sed 's/\[--/ \\\n       '"${CMDNAME//?/ }"' [--/g'
  if [[ $positionalArgListing ]]; then
    echo " ${positionalArgListing}" |
      sed 's/ / \\\n       '"${CMDNAME//?/ }"' /g'
  else
    echo ''
  fi
  echo "${optionHelp}"
} #printCMDUsage

function cleanup_umountdir_and_disconnect_fcp {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Umount dir and cleanup all FCP devices.
  inform "Begin to cleanup umount dir."
  umount $deviceMountPoint/proc > /dev/null
  umount $deviceMountPoint/sys > /dev/null
  umount $deviceMountPoint/run > /dev/null
  umount $deviceMountPoint/dev > /dev/null
  umount $deviceMountPoint > /dev/null
  rm -rf $deviceMountPoint
  inform "Finish cleanuping umount dir."

  cleanup_fcpdevices
} # cleanup_umountdir_and_disconnect_fcp

function cleanup_fcpdevices {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  # Disconnect all FCPs
  inform "Begin to disconnect FCPs"
  for i in "${!fcps[@]}"
  do
    if [[ -z "$wwpns" ]];then
      output=`/opt/zthin/bin/smcli Image_Device_Undedicate -T $zthinUserID -v ${fcps[$i]} 2>&1`
      if [[ $? != 0 ]]; then
          inform "Undedicate FCP device ${fcps[$i]} fails, reason: $output"
      else
          inform "Undedicate FCP device ${fcps[$i]} successfully."
      fi
    else
      for j in "${!wwpns[@]}"
      do
        inform "Disconnect FCP device: ${fcps[$i]}"
        disconnectFcp ${fcps[$i]} 0x${wwpns[$j]} ${lun}
      done
    fi
  done
  inform "Finish disconnecting FCPs."
  inform "Begin to refresh multipath bindings."
  # Don't need to judge if the multipath_enabled exist,
  # if it doesn't exist, the $multipath_enabled will be None.
  if [[ $multipath_enabled == 1 ]]; then
    inform "Cleaning up multipath bindings and wwids."
    # delete the wwid and refresh multipath map
    multipath -F -W 1>/dev/null
    if [[ -f /etc/multipath/bindings ]]; then
        sed -i "/$wwid/d" /etc/multipath/bindings
    fi
  fi
  inform "Finish refreshing multipath bindings."
} #cleanup_fcpdevices

function printCMDExamples {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Prints a short description of this command.
  # @Overrides:
  #   printCMDDescription{} in "zthinshellutils".
  # @Code:
  echo "Example:
  ./refresh_bootmap.sh --fcpchannel="5d71" --wwpn="5005076802100C1B" --lun=0000000000000000
  ./refresh_bootmap.sh --fcpchannel="5d71,5d72" --wwpn="5005076802100C1B,5005076802200C1B,5005076802300C1B,5005076802400C1B,5005076802400C1A,5005076802300C1A,5005076802200C1A,5005076802100C1A" --lun=0000000000000000" 
} #printCMDDescription

function parseArgs {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Parses and checks command-line arguments.
  # @Code:
  # Non-local variables in this function are intentionally non-local.
  isOption -h --help "         Print this help message."   && printHelp='true'
  isOption -v --verbose "      Print verbose output."      && verbose='-v'
  isOption -x --debug "        Print debugging output."    && debug='-x'
  isOption -t --type "         Device type. The default device typs is fcp and fcp is the only currently supported type."
  isOption -f --fcpchannel "   FCP channel IDs. Support multi fcpchannel, split by ','."
  isOption -w --wwpn "         World Wide Port Name IDs. Support multi wwpn, split by ',' Example: 5005076802400c1b"
  isOption -l --lun "          Logical Unit Number ID. Example: 0000000000000000"
  isOption -s --skipzipl "     Whether skip ZIPL, if YES, only return physical World Wide Port Name IDs"

  if [[ $printHelp ]]; then
    printHelp
    exit 0
  fi

  for i in "$@"
  do
  case $i in
      -t|--type=*)
      device_type="${i#*=}"
      shift # past argument=value
      ;;
      -f|--fcpchannel=*)
      fcpchannel="${i#*=}"
      shift # past argument=value
      ;;
      -w|--wwpn=*)
      wwpn="${i#*=}"
      shift # past argument=value
      ;;
      -l|--lun=*)
      lun="${i#*=}"
      shift # past argument=value
      ;;
      -s|--skipzipl=*)
      skipzipl="${i#*=}"
      shift # past argument=value
      ;;
      --default)
      DEFAULT=YES
      shift # past argument with no value
      ;;
      *)
      printError "Unknow option: $i." 
      exit 3
      ;;
  esac
  done

  # Set default device type to FCP if not specific.
  device_type=${device_type:-fcp}
  device_type=$(echo ${device_type} | tr '[:upper:]' '[:lower:]')
} #parseArgs

function checkSanity {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Performs basic checks to ensure that a successful deploy can reasonably
  #   be expected.
  if [[ $device_type = 'fcp' ]]; then
    if [[ ! $lun || ! $fcpchannel ]];then
        printError "Please specify lun and fcpchannel."
        exit 4
    fi
    # Intentionally non-local variable.
    format='FCP'
  elif [[ $device_type = 'fba' ]]; then
    # We don't support fba now.
    :
  else
    printError "Unknown device type."
    exit 5
  fi
} #checkSanity

function checkMultipath {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Check whether multipath is enabled by checking the multipathd service status
  #   Set the multipath_enabled variable according to the result.
  multipath_enabled=0
  service_output=`systemctl status multipathd | grep "active (running)"`
  if [[ $? -eq 0 ]]; then
      inform "multipathd service is active and running."
      multipath_enabled=1
  fi
} #checkMultipath

function refreshZIPL {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Retrieves initial information for the IPL from disk
  # @Parameters:
  #   devNode: the path to access the disk
  # @Returns:
  #   None
  local devNode=$1

  inform "Begin to refreshZIPL."

  if [[ ! -e "$devNode" ]]; then
    printError "devNode ${devNode} doesn't exist. Get fcps: ${fcps[*]} and wwpns: ${wwpns[*]}"
    exit 6
  fi
  
  if [[ $skipzipl = 'YES'  ]]; then
    inform "WWPNs: ${wwpns[*]}"
    inform "FCPs: ${fcps[*]}"
    return
  fi

  # Create a mount dir.
  deviceMountPoint=$(/usr/bin/mktemp -d /mnt/XXXXX)
  if [[ $? != 0 ]]; then
    printError "Create mount dir fails."
    exit 7
  fi

  # Register EXIT sigual for umount dir and disconnect fcp devices
  trap 'cleanup_umountdir_and_disconnect_fcp' EXIT

  # Try to mount fcp device.
  inform "devNode is: $devNode point: $deviceMountPoint"
  mount $devNode $deviceMountPoint
  mount -t proc /proc $deviceMountPoint/proc
  mount -o bind /sys $deviceMountPoint/sys
  mount -o bind /run $deviceMountPoint/run
  mount -o bind /dev $deviceMountPoint/dev
  if [[ $? != 0 ]]; then
    printError "Mount fails."
    exit 8
  fi

  #Get target os version
  local osRelease="$deviceMountPoint/etc/os-release"
  local slesRelease="$deviceMountPoint/etc/SuSE-release"
  local rhelRelease="$deviceMountPoint/etc/redhat-release"

  if [[ -e $osRelease ]]; then
    os=`cat $osRelease | grep "^ID=" | sed \
      -e 's/ID=//' \
      -e 's/"//g'`
    version=`cat $osRelease | grep "^VERSION_ID=" | sed \
      -e 's/VERSION_ID=//' \
      -e 's/"//g' \
      -e 's/\.//'`
    os=$os$version

  #The /etc/redhat-release file will be deprecated in rhel7 and later release
  elif [[ -e $rhelRelease ]]; then
    os='rhel'
    version=`cat $rhelRelease | grep -i "Red Hat Enterprise Linux Server" | sed \
      -e 's/[A-Za-z\/\.\(\)]//g' \
      -e 's/^ *//g' \
      -e 's/ *$//g' \
      -e 's/\s.*$//'`
    os=$os$version
  fi
  # Prepare the new rd.zfcp parameter value
  fcps_len=${#fcps[@]}
  wwpns_len=${#wwpns[@]}
  res=$(( $fcps_len * $wwpns_len ))
  rd_zfcp=""

  # For rhel, the max support items of "rd.zfcp" is 8, so we need to judge the
  # result of FCP devices multiply WWPN.
  #
  # If the result of fcps_len multiply wwpns_len less than 8, we can do a fullmesh of
  # all FCP devices and WWPNs.
  if [ $res -le 8 ] # less than and equal 8
  then
    for fcp in ${fcps[@]}
    do
      for w in ${wwpns[@]}
      do
        rd_zfcp+="rd.zfcp=0.0."$fcp,0x$w,$lun" "
      done
    done
    inform "Do a fullmesh of FCP devices: ${fcps[*]} and WWPN: ${wwpns[*]} for rd.zfcp: $rd_zfcp"
  # If the result of fcps_len multiply wwpns_len greater than 8, we need to use
  # every FCP device and as much as possible to use WWPN.
  else
    paths=0
    fcp_index=0
    wwpn_index=0
    while [ $paths -lt 8 ]
    do
        if [ $fcp_index -ge $fcps_len ]; then
          fcp_index=0
        fi

        if [ $wwpn_index -ge $wwpns_len ]; then
          wwpn_index=0
        fi

        x="rd.zfcp=0.0."${fcps[$fcp_index]},0x${wwpns[$wwpn_index]},$lun" "

        if [[ $rd_zfcp =~ $x ]]
        then # the path already exist in rd_zfcp
          wwpn_index=$(( $wwpn_index + 1 ))
          continue
        else # the path doesn't exist in rd_zfcp
          rd_zfcp+=$x
          fcp_index=$(( $fcp_index + 1 ))
          wwpn_index=$(( $wwpn_index + 1 ))
        fi
        paths=$(( $paths + 1 ))
    done
    inform "Use as much as possible of FCP devices: ${fcps[*]} and WWPN: ${wwpns[*]} rd.zfcp: $rd_zfcp"
  fi

  # Exec zipl command to prepare device for initial problem load
  if [[ $os == rhel7* ]]; then
    inform "Refresh $os zipl."
    zipl_conf='etc/zipl.conf'
    
    ### Delete all items start with "rd.zfcp="
    sed -ri 's/rd.zfcp=\S+\s*[[:space:]]//g' ${deviceMountPoint}/$zipl_conf
    
    # Remove quote
    sed -i 's/\"$//g' ${deviceMountPoint}/$zipl_conf
    
    # Remove quote
    sed -i 's/[ \t]*$//g' ${deviceMountPoint}/$zipl_conf
    
    # Append rd.zfcp= string to "parameters=" line.
    sed -i "/^[[:space:]]parameters=/ s/$/ $rd_zfcp/" ${deviceMountPoint}/$zipl_conf
    
    # Append quote to "parameters=" line
    sed -i "/^[[:space:]]parameters=/ s/$/\"/" ${deviceMountPoint}/$zipl_conf

  elif [[ $os == rhel8* ]]; then
    inform "Refresh $os zipl."
    kernel_version_conf_files=`find $deviceMountPoint/boot/loader/entries/ -name '*.conf'`
    diskPath="/dev/disk/by-path/ccw-0.0.${fcps[0]}-zfcp-0x${wwpns[0]}:${lun}"
    wwid=`/usr/lib/udev/scsi_id --whitelisted $diskPath 2> /dev/null`
    rootPath="root=\/dev\/disk\/by-id\/dm-uuid-part1-mpath-$wwid "
    
    # The Volume may be contains several kernel version conf files.
    # So every conf file needs to be changed.
    for confFile in $kernel_version_conf_files; do
       # Delete all items start with "rd.zfcp="
       sed -ri 's/rd.zfcp=\S+\s*[[:space:]]//g' $confFile
       # Remove trailing space
       sed -i 's/[ \t]*$//g' $confFile
       # Update the root file system target path
       sed -ri "s/root=\S+\s*[[:space:]]/$rootPath /g" $confFile
       # Append rd.zfcp= string to "options root=" line.
       sed -i "/^options root=/ s/$/ $rd_zfcp/" $confFile
    done
  elif [[ $os == "" ]]; then
    inform "This is not the root disk, zipl will not be executed on it"
  else
    inform "The os version is: $os, this is not a supported linux distro"
  fi

  # Remove dirty data in /etc/fstab left by the previous vm deploy
  # remove lines containing "/mnt/ephemeral" or " swap swap "
  # TODO: need to update the implementation method when later the ephemeral
  # mount point maybe changed to user configurable, to not hardcode to /mnt/ephemeralxxx
  inform "Updating ${deviceMountPoint}/etc/fstab"
  sed -i '/\/mnt\/ephemeral/d' ${deviceMountPoint}/etc/fstab
  sed -i '/ swap swap /d' ${deviceMountPoint}/etc/fstab

  # Remove stale /etc/multipath/bindings and /etc/multipath/wwids
  # These files may contain the old info of the captured base vm,
  # which disrupts the volume bindings of the new vm deployed from above captured image, resulting
  #  a. mpath name inconsistent between 'df' and 'multipath'
  #  b. in some scenairo, boot-volume extend partition fail due to 'a'
  # By removing them, they can be recreated later with correct volume bindings
  # when multipathd start at vm boot time.
  inform "Move stale multipath bindings and wwids if any"
  if [[ -r "${deviceMountPoint}/etc/multipath/bindings" ]]; then
    mv ${deviceMountPoint}/etc/multipath/bindings ${deviceMountPoint}/etc/multipath/bindings.bak
  fi
  if [[ -r "${deviceMountPoint}/etc/multipath/wwids" ]]; then
    mv ${deviceMountPoint}/etc/multipath/wwids ${deviceMountPoint}/etc/multipath/wwids.bak
  fi

  # Refresh bootmap
  out=`chroot $deviceMountPoint /sbin/zipl 2>&1`
  rc=$?
  if (( rc != 0 )); then
    printError "Failed to execute zipl on $os due to $out"
    return 1
  fi
  # Sync to ensure cached date writen back to target disk
  blockdev --flushbufs $diskPath > /dev/null

  # Return wwpns
  inform "WWPNs: ${wwpns[*]}"

  # Return fcps
  inform "FCPs: ${fcps[*]}"
  return
} #refreshZIPL

function refreshFCPBootmap {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Refresh bootmap info of FCP device.
  # @Code:

  local out
  local rc
  local errorFile

  fcpchannels=$(echo ${fcpchannel} | tr '[:upper:]' '[:lower:]')
  lun=$(echo ${lun} | tr '[:upper:]' '[:lower:]')
  wwpn=$(echo ${wwpn} | tr '[:upper:]' '[:lower:]')

  # Add 0x to lun
  lun=0x$lun

  # Split fcpchannels and wwpn by ","
  IFS=',' read -r -a fcps <<< "$fcpchannels"
  IFS=',' read -r -a ws <<< "$wwpn"

  # Check the count of FCP devices, if greater than 8, report error and exit.
  fcps_len=${#fcps[@]}
  if [ $fcps_len -gt 8 ]
  then
      printError "ERROR: The FCP channel count $fcps_len is greater than 8."
      exit 12
  fi

  enable_zfcp_mod=`lsmod | grep zfcp`
  if [ -z "$enable_zfcp_mod" ];then
      modprobe zfcp
  fi

  local zthinUserID=$(vmcp q userid | awk '{printf $1}')
  inform "zthinUserID is $zthinUserID"

  # Check whether multipath is enabled by checking the multipathd service
  checkMultipath

  # Register EXIT sigual for cleanup FCP devices.
  trap 'cleanup_fcpdevices' EXIT

  # Try to find which physical wwpns are being used.
  for fcp in "${fcps[@]}"
  do
    # Ignore fcp device firstly
    /sbin/cio_ignore -r $fcp > /dev/null

    # Convert the fcp device to upper case
    fcp_device=($(echo "$fcp"| tr 'a-z' 'A-Z'))

    # Execute the vmcp command, don't execute dedicate command
    # if the fcp device is already dedicate to the compute node,
    res=`vmcp q fcp | grep FCP| grep "$fcp_device"`

    # The fcp device doesn't dedicate to the compute node so
    # dedicate it to the compute node.
    if [[ -z $res ]]; then
      # Try to connect FCP device and online it.
      /opt/zthin/bin/smcli Image_Device_Dedicate -T $zthinUserID -v $fcp -r $fcp
      if [[ $? -ne 0 ]]; then
        printError "Dedicate fails, maybe other machines use this FCP: ${fcp}"
        exit 3
      fi
    fi

    chccwdevDevice $fcp "online"
    if [[ $? -ne 0 ]]; then
      # Online failed
      rc=4
    fi
    host=`lszfcp | grep "$fcp" | cut -d " " -f 2`
    echo "- - -" > /sys/class/scsi_host/$host/scan
    # Use sleep command to make sure all files are generated
    # by chccwdev command.
    sleep 1
  done

  # Read and parse "lsluns" output then check command input.
  # If command inputs, merge "lsluns" output and command input
  # If not, use "lsluns" parse output directly.
  wwpns=()
  wwpn=$(lsluns |grep -E 'at port' | awk -F 'at port' '{print $2}' | sed -e 's/^[ ]*//g' | sed -e 's/\:$//g' | tr ' ' '\n' | sort -u | sed -e 's/[\t]//g' |  tr '\n' ' ')
  IFS=' ' read -r -a lsluns_wwpns <<< "$wwpn"
  if [[ ${ws[*]} ]]; then
    for w in "${lsluns_wwpns[@]}"
    do
      w=${w##*0x}
      if [[ "${ws[@]}" =~ "$w" ]]; then
        wwpns+=($w)
      fi
    done
  else
    for i in "${!lsluns_wwpns[@]}"
    do
      x=${lsluns_wwpns[i]##*0x}
      wwpns[$i]=$x
    done
  fi

  # Remove duplicates
  wwpns=($(echo "${wwpns[@]}" | tr ' ' '\n' | sort -u | tr '\n' ' '))
  if [[ ! $wwpns[0] ]]; then
    printError "No available WWPN found."
    exit 9
  fi
  inform "These WWPNs will be operated: ${wwpns[*]}"

  # Check the count of WWPN, if greater than 8, warning and don't exit.
  wwpns_len=${#wwpns[@]}
  if [ $wwpns_len -gt 8 ]
  then
      inform "WARNING: The count of WWPN can't greater than 8."
  fi

  for i in "${!fcps[@]}"
  do
    for j in "${!wwpns[@]}"
      do
        x="/dev/disk/by-path/ccw-0.0.${fcps[i]}-zfcp-0x${wwpns[j]}:${lun}"
        if [[ -e $x ]];then
          diskPath=$x
          inform "$diskPath is: $diskPath"
          break
        fi
      done
  done

  if [[ -z $diskPath ]];then
    inform "Can not find a proper path for FCP devices: ${fcps[*]} and wwpns: ${wwpns[*]}"
    exit 10
  fi

  # Move the proper FCP device to the first position.
  tmp=${fcps[0]}
  fcps[0]=${fcps[i]}
  fcps[$i]=$tmp

  # Move the proper WWPN to the first position.
  tmp=${wwpns[0]}
  wwpns[0]=${wwpns[j]}
  wwpns[$j]=$tmp


  # Set devNode according to whether multipath is enabled on the compute node
  wwid=""
  if [[ $multipath_enabled == 1 ]]; then
      wwid=`/usr/lib/udev/scsi_id --whitelisted $diskPath 2> /dev/null`
      devNode="/dev/disk/by-id/dm-uuid-part1-mpath-$wwid"
      inform "Multipath is enabled, using multipath devNode: $devNode."
  else
      devNode="/dev/disk/by-path/ccw-0.0.${fcps[0]}-zfcp-0x${wwpns[0]}:${lun}-part1"
      # For blockdev command use.
      inform "Multipath is not enabled, using single path devNode: $devNode."
  fi
  
  refreshZIPL $devNode
} #refreshFCPBootmap

function refreshBootMap {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Refresh the bootmap of disk image.
  # @Code:

  if [[ $format = 'FCP' ]]; then
    refreshFCPBootmap
  elif [[ $format = 'FBA' ]]; then
    :
  else
    printError "Device type not recognised."
    exit 11
  fi 
} #refreshBootMap


###############################################################################
### START EXECUTION ###########################################################
###############################################################################

parseArgs $@
checkSanity
refreshBootMap

###############################################################################
### END OF SCRIPT #############################################################
###############################################################################
