#!/bin/bash
###############################################################################
# Copyright 2017 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: creatediskimage                                                  #
#                                                                             #
# Creates an image of the disk at the specified channel ID on the specified   #
# z/VM guest system.                                                          #
###############################################################################

source /opt/zthin/lib/zthinshellutils
version="2.0"

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

function printCMDUsage {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Prints usage help text using a list generated by previous requests made
  #   to parse the command-line argument list.
  # @Override in zthinshellutils
  # @Code:
  echo "USAGE: $CMDNAME [OPTIONS] USER_ID CHANNEL_ID IMAGE_FILE"
  echo "       $CMDNAME [OPTIONS] CHANNEL_ID WWPN LUN IMAGE_FILE"

  echo "${optionHelp}"
  echo "NAMED ARGUMENTS: ${namedArgListing}"
} #printCMDUsage{}

###############################################################################

function printCMDDescription {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Prints a short description of this command.
  # @Overrides:
  #   printCMDDescription{} in "xcatshellutils".
  # @Code:
  echo -n "Creates an image of the disk at the specified channel ID on the "
  echo    "specified z/VM guest system."
} #printCMDDescription{}

###############################################################################

function parseArgs {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Parses and checks command-line arguments.
  # @Code:
  # Apply any defaults before handling operands
  
  if [[ -z "$gzipCompression" ]]; then
    gzipCompression=6
  fi
  
  # Non-local variables in this function are intentionally non-local.
  isOption -h --help '     Print this help message.'   && printHelp='true'
  isOption -V --version '  Print the version number of this script.'   && printVersion='true'
  isOption -v --verbose '  Print verbose output.'      && verbose='-v'
  isOption -x --debug '    Print debugging output.'    && debug='-x'
  
  # Assume the request is for SCSI disk capture.
  getPositionalArg 2 wwpn
  if [[ $wwpn != 0x* ]]; then
    unset wwpn
  fi
  
  if [[ -n $wwpn ]]; then
    getPositionalArg 1 fcpChannel
    #getPositionalArg 2 wwpn
    getPositionalArg 3 lun
    getPositionalArg 4 imageFile

    # Make channel, WWPN, and LUN lower case
    fcpChannel=$(echo ${fcpChannel} | tr '[:upper:]' '[:lower:]')
    wwpn=$(echo ${wwpn} | tr '[:upper:]' '[:lower:]')
    lun=$(echo ${lun} | tr '[:upper:]' '[:lower:]')
  else
    getPositionalArg 1 userID
    getPositionalArg 2 channelID
    getPositionalArg 3 imageFile
  fi
  
  getNamedArg --compression compOper 'Compression level, 0-9'
  if [[ $? -eq 0 ]]; then
    gzipCompression=$compOper
  fi
  
  # Handle options that provide info but don't deal with locks or disks
  if [[ $printVersion ]]; then
    echo $version
  fi
  
  if [[ $printHelp ]]; then
    printHelp
  fi
  
  if [[ $printHelp || $printVersion  ]]; then
    exit 0
  fi
  
  # Handle rest of operands
  # Validate the compression specified on the command line or in the configuration file
  if [[ ! $gzipCompression =~ "[0-9]" || ${#gzipCompression} -ne 1 ]]; then
      printError "Compression operand is not between 0 and 9, inclusive: $gzipCompression"
      exit 1
  fi
  
  local badOptions=$(getBadOptions)
  if [[ $badOptions ]]; then
    echo "ERROR: ${badOptions}"
    printCMDUsage
    exit 1
  fi
  
  if [[ -n $wwpn ]]; then
    if [[ -z "$fcpChannel" || -z "$lun" || -z "$imageFile" ]]; then
      echo 'ERROR: Missing required parameter.'
      printCMDUsage
      exit 1
    fi
    if (( ${#wwpn} < 18 || ${#lun} < 18 )); then
      echo 'ERROR: WWPN and LUN must be a 16 byte number.'
      exit 1
    fi
  elif [[ -z $userID || -z $channelID || -z $imageFile ]]; then
    echo 'ERROR: Missing required parameter.'
    printCMDUsage
    exit 1
  fi
  
  # Remove old traces beyond the number specified in /var/opt/zthin/settings.conf. If the
  # specified number is 0, less than 0, or not actually a number, then we skip
  # this cleanup.
  if [[ $keepOldTraces -gt 0 ]]; then
    local removeTraces=$(($(ls -1 /var/log/zthin/creatediskimage_trace_* | wc -l) -
                          ${keepOldTraces}))
    if [[ $removeTraces -gt 0 ]]; then
      for trace in $(ls -1 /var/log/zthin/creatediskimage_trace_* |
                     head -${removeTraces}); do
        rm -f $trace
      done
    fi
  fi
  
  timestamp=$(date -u --rfc-3339=ns | sed 's/ /-/;s/\.\(...\).*/.\1/')
  logFile=/var/log/zthin/creatediskimage_trace_${timestamp}.txt
  mkdir -p /var/log/zthin
  if [[ $debug ]]; then
    exec 2> >(tee -a $logFile)
    set -x
  else
    exec 2> $logFile
    set -x
  fi
  inform "creatediskimage start time: ${timestamp}"
  
  if [[ -n $wwpn ]]; then
    echo "FCP CHANNEL:    \"$fcpChannel\""
    echo "WWPN:           \"$wwpn\""
    echo "LUN:            \"$lun\""
    echo "IMAGE FILE:     \"$imageFile\""
    echo "COMPRESSION:    \"$gzipCompression\""
    echo ""
  else
    echo "SOURCE USER ID: \"$userID\""
    echo "DISK CHANNEL:   \"$channelID\""
    echo "IMAGE FILE:     \"$imageFile\""
    echo "COMPRESSION:    \"$gzipCompression\""
    echo ""
  fi
} #parseArgs{}

###############################################################################

function checkSanity {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Performs basic checks to ensure that a successful capture can reasonably
  #   be expected.
  # @Code:
  # Make sure the specified image file name is not already taken.
  if [[ -e $imageFile ]]; then
    printError 'The specified image file already exists.'
    exit 3
  fi
  
  if [[ $userID && $(isSystemActive $userID) ]]; then
    printError 'The specified source system is currently running.'
    exit 3
  fi
  
  # Intentionally non-local variable.
  passedSanity='true'
} #checkSanity{}

###############################################################################

function createDiskImage {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Creates an image of the specified disk.
  # @Code:
  
  function captureFCP {
    : SOURCE: ${BASH_SOURCE}
    : STACK:  ${FUNCNAME[@]}
    # @Description:
    #   Capture the specified SCSI/FCP disk with `dd`.
    # @Code:
    local out
    local rc
    
    # First bytes of image file identify file as an image of xCAT
    # and indicates the type of source disk.
    local header="xCAT FCP Disk Image:"
    
    # Next 16 bytes of image file are an ASCII representation of the disk's
    # size and the unit (BLK or CYL) used for this measurement:
    size=`/usr/bin/sg_readcap /dev/disk/by-path/ccw-0.0.${fcpChannel}-zfcp-${wwpn}:${lun} |
      egrep -i "Device size:" | 
        awk '{printf("%0.0f BLK", $3/512)}'`
    local nextPart=`echo -n $size |
          sed 's/^/                /;s/.*\(.\{16\}\)$/\1/'`
    header="$header$nextPart"
    
    # Next 9 bytes are header length in ASCII (version 2 and later images).
    header="$header HLen: 0055"
    
    # Next 8 bytes indicate GZIP compression level in ASCII.
    header="$header GZIP: $gzipCompression"
    
    # Write the header to the file and fill up the rest with blanks 
    # to make a reasonable blocksize
    if (( ${#header} <= $blockSize )); then
      printf "%-512s" "$header" > $imageFile
    else
      # Handle writing the beginning full block portion of header
      local blocks=$(( ${#header} / $blockSize ))
      local partialLength=$(( $blocks * $blockSize ))
      echo -n "${header:0:$partialLength}" > $imageFile
      # Handle writing rest of header
      printf "%-512s" "${header:$partialLength}" >> $imageFile
    fi
    
    # Now read and compress the data on this fixed-block disk and write the
    # result to the end of our image file.
    inform "Creating $(basename $imageFile) image file for ${wwpn}/${lun} disk\
      at channel ${fcpChannel} with disk size ${size}."
    inform "Compression level: $gzipCompression"
    if (( gzipCompression == 0 )); then
      out=`dd if=/dev/disk/by-path/ccw-0.0.${fcpChannel}-zfcp-${wwpn}:${lun} 2>&1 >> $imageFile`
      rc=$?
    else
      out=`gzip -$gzipCompression 2>&1 < /dev/disk/by-path/ccw-0.0.${fcpChannel}-zfcp-${wwpn}:${lun} >> $imageFile`
      rc=$?
    fi
    if (( rc )); then
      out=`echo $out | tr '\n' ' '`
      printError "An error was encountered while creating disk image. $out"
      exit 3
    fi
  } #captureFCP{}
  
  function captureFBA {
    : SOURCE: ${BASH_SOURCE}
    : STACK:  ${FUNCNAME[@]}
    # @Description:
    #   Capture the specified FBA disk with `dd`.
    # @Code:
    local out
    local rc
    
    # First bytes of image file identify file as an image of xCAT
    # and indicates the type of source disk.
    local header="xCAT FBA Part Image:"
    
    # Next 16 bytes of image file are an ASCII representation of the disk's
    # size and the unit (BLK or CYL) used for this measurement:
    size=`vmcp q v $alias | awk '{printf $6" "$7}'`
    local nextPart=`echo -n $size |
      sed 's/^/                /;s/.*\(.\{16\}\)$/\1/'`
    header="$header$nextPart"
    
    # Next 9 bytes are header length in ASCII (version 2 and later images).
    header="$header HLen: 0055"
    
    # Next 8 bytes indicate GZIP compression level in ASCII.
    header="$header GZIP: $gzipCompression"
    
    # Write the header to the file and fill up the rest with blanks 
    # to make a reasonable blocksize
    if (( ${#header} <= $blockSize )); then
      printf "%-512s" "$header" > $imageFile
    else
      # Handle writing the beginning full block portion of header
      local blocks=$(( ${#header} / $blockSize ))
      local partialLength=$(( $blocks * $blockSize ))
      echo -n "${header:0:$partialLength}" > $imageFile
      # Handle writing rest of header
      printf "%-512s" "${header:$partialLength}" >> $imageFile
    fi
    
    # Now read and compress the data on this fixed-block disk and write the
    # result to the end of our image file.
    inform "Creating $(basename $imageFile) image file for ${userID}'s\
      disk at channel ${channelID} with disk size ${size}."
    inform "Compression level: $gzipCompression"
    if (( gzipCompression == 0 )); then
      out=`dd if=/dev/disk/by-path/ccw-0.0.${alias}-part1 2>&1 >> $imageFile`
      rc=$?
    else
      out=`gzip -$gzipCompression < /dev/disk/by-path/ccw-0.0.${alias}-part1 2>&1 >> $imageFile`
      rc=$?
    fi
    
    if (( rc )); then
      out=`echo $out | tr '\n' ' '`
      printError "An error was encountered while creating disk image. $out"
      exit 3
    fi
  } #captureFBA{}
  
  function captureCKD {
    : SOURCE: ${BASH_SOURCE}
    : STACK:  ${FUNCNAME[@]}
    # @Description:
    #   Capture the specified CKD disk with `ckdencode`.
    # @Code:
    local out
    local rc
    
    # First bytes of image file identify file as an image of xCAT
    # and indicates the type of source disk.
    local header="xCAT CKD Disk Image:"
    
    # Next 16 bytes of image file are an ASCII representation of the disk's
    # size and the unit (BLK or CYL) used for this measurement:
    size=`vmcp q v $alias | awk '{printf $6" "$7}'`
    local nextPart=`echo -n $size |
      sed 's/^/                /;s/.*\(.\{16\}\)$/\1/'`
    header="$header$nextPart"
    
    # Next 12 bytes are header length in ASCII (version 2 and later images).
    header="$header HLen: 0055"
    
    # Next 8 bytes indicate GZIP compression level in ASCII.
    header="$header GZIP: $gzipCompression"
    
    # Write the header to the file and fill up the rest with blanks 
    # to make a reasonable blocksize
    if (( ${#header} <= blockSize )); then
      printf "%-512s" "$header" > $imageFile
    else
      # Handle writing the beginning full block portion of header
      local blocks=$(( ${#header} / $blockSize ))
      local partialLength=$(( $blocks * $blockSize ))
      echo -n "${header:0:$partialLength}" > $imageFile
      # Handle writing rest of header
      printf "%-512s" "${header:$partialLength}" >> $imageFile
    fi
    
    # Now read, encode, and compress the data on this count-key-data disk and
    # write the result to the end of our image file.
    local reason=""
    errorFile=`mktemp -p /var/log/zthin -t createStderr.XXXXXXXXXX`
    inform "Creating $(basename $imageFile) image file for ${userID}'s\
      disk at channel ${channelID} with disk size ${size}."
    inform "Compression level: $gzipCompression"
    if (( gzipCompression == 0 )); then
      ckdencode /dev/disk/by-path/ccw-0.0.${alias} 2>>$errorFile >> $imageFile
      rc=$?
      if (( rc )); then
        out=`cat $errorFile | tr '\n' ' '`
        reason="rc: $rc $out"
      fi
    else
      ckdencode /dev/disk/by-path/ccw-0.0.${alias} 2>>$errorFile | gzip -$gzipCompression 2>>$errorFile >> $imageFile
      declare -a pipeRC=($PIPESTATUS ${PIPESTATUS[@]})
      rc=${pipeRC[0]}
      if (( rc )); then
        out=`cat $errorFile | tr '\n' ' '`
        declare -a stages=('overall_placeholder' 'ckdencode' 'gzip' )
        getStageFailures stages[@] pipeRC[@]
        reason="at stage(rc): $getStageFailuresOut $out"
      fi
    fi
    
    if (( rc )); then
      printError "An error was encountered while creating disk image $reason"
      exit 3
    fi
  } #captureCKD{}

  if [[ ${userID} ]]; then
    connectDisk $userID $channelID 'rr' 0
  else
    connectFcp ${fcpChannel} ${wwpn} ${lun}
    if (( $? )); then
      printError "An error was encountered while attaching disk."
      exit 3
    fi
  fi

  local alias=$(getDiskAlias $userID $channelID)
  if [[ ${fcpChannel} && $(vmcp q v ${fcpChannel} | grep 'ON FCP') ]]; then
    # Device is an FCP disk. Capture with `dd`. 
    captureFCP
  elif [[ $(vmcp q v $alias | grep 'BLK ON DASD') ]]; then
    # Device is an FBA disk. Capture with `dd`. 
    captureFBA
  elif [[ $(vmcp q v $alias | grep 'CYL ON DASD') ]]; then
    if [[ -e /sys/bus/ccw/devices/0.0.${alias}/raw_track_access ]]; then
      # Device is a CKD disk with raw_track_access as an available option.
      # Reconnect in raw_track_access mode and capture with ckdencode.
      disconnectDisk $userID $channelID 0
      connectDisk $userID $channelID 'rr' 1
      alias=$(getDiskAlias $userID $channelID)
      captureCKD
    else
      printError "Attempting to capture CKD disk failed; the raw_track_access\
                  DASD driver option is unavailable on this system." 
      exit 3
    fi
  else
    printError "Specified disk does not exist or is of an unknown type." 
    exit 3
  fi
} #createDiskImage{}


###############################################################################
### SET TRAP FOR CLEANUP ON EXIT ##############################################
###############################################################################

function cleanup {
  : SOURCE: ${BASH_SOURCE}
  : STACK:  ${FUNCNAME[@]}
  # @Description:
  #   Clean up lock files, disk links, and (if we're exiting with an error)
  #   the unfinished package.
  # @Code:
  # Nothing to do for help or version options.
  if [[ $printHelp || $printVersion ]]; then
    return
  fi
  
  if [[ -e $errorFile ]]; then
    rm $errorFile
  fi
  
  if [[ $successful ]]; then
    inform "Image creation successful."
    # Only keep traces of failed disk-image creation attempt unless overriden 
    # by a configuration property.
    if [[ -e $logFile ]]; then
      if [[ $saveAllLogs ]]; then
        inform "A detailed trace can be found at: ${logFile}"
      else
        rm -f $logFile
      fi
    fi
  else
    echo -e '\nIMAGE CREATION FAILED.'
    [[ $logFile ]] && inform "A detailed trace can be found at: ${logFile}"
    # If our image-creation wasn't successful, we shouldn't leave a broken
    # image sitting around in the staging repository.
    [[ $passedSanity ]] && rm -f $imageFile
  fi

  # Also, let's make sure we've released our connection to the source disk.
  if [[ $userID ]]; then
    disconnectDisk $userID $channelID 0
  else
    disconnectFcp ${fcpChannel} ${wwpn} ${lun}
  fi
  
  timestamp=$(date -u --rfc-3339=ns | sed 's/ /-/;s/\.\(...\).*/.\1/')
  inform "creatediskimage end time: ${timestamp}"
} #cleanupProperties{}

trap 'cleanup' EXIT

trap "echo -e '\nExecution interrupted. Exiting...\n'; exit" SIGINT

###############################################################################
### START EXECUTION ###########################################################
###############################################################################
parseArgs
checkSanity
createDiskImage
successful='true'

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