#!/bin/bash
# License: GPL 
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to check the Clonezilla image integrity
# (1) Check the partition table files (GPT included)
# (2) Check the MBR/EFI boot loader
# (3) Check the partition image using partclone.chkimg

#
DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"
. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
. /etc/drbl/drbl-ocs.conf
. $DRBL_SCRIPT_PATH/sbin/ocs-functions

# Load the config in ocs-live.conf. This is specially for Clonezilla live. It will overwrite some settings of /etc/drbl/drbl-ocs.conf, such as $DIA...
[ -e "/etc/ocs/ocs-live.conf" ] && . /etc/ocs/ocs-live.conf

# Settings
# By default we do not check the image saved by dd.
dd_checking_method="no-check"
nogui="off"   # Default to turn on TUI

#
USAGE() {
    echo "$ocs - To check the image of Clonezilla"
    echo "Usage:"
    echo "To run $ocs:"
    echo "$ocs [OPTION] IMAGE DEVICE"
    echo "IMAGE is the image dir name, not absolute path"
    echo "DEVICE is the device name (not whole disk, only partition), e.g., sda1, sda2, or /dev/sda1, /dev/sda2..."
    echo "If \"ask_user\" is used as IMAGE or DEVICE, a dialog menu will be shown to allow selection."
    echo "If no IMAGE is specified, a dialog menu will be shown."
    echo "If no DEVICE is specified, all of the devices will be checked."
    echo
    echo "Options:"
    echo "-b, --batch-mode   Run image checking in batch mode"
    echo "-nogui, --nogui    Do not show GUI (TUI) of Partclone or Partimage, use text only"
    echo "-or, --ocsroot DIR Specify DIR (absolute path) as directory ocsroot (i.e. overwrite the ocsroot assigned in drbl.conf)"
    echo
    echo "Ex:"
    echo "To check the image \"my-image\", which is located in $ocsroot/my-image, run"
    echo "   $ocs my-image"
    echo
} # end of USAGE
#
img_check_by_partclone() {
  local file_ fs_ partclone_img_info_tmp file_basename rc
  local cmd_partclone_check split_flag link_flag check_flag run_cmd_
  local -a retcodes='()'
  # First, we find the filesystem 
  # $target_d, $img_file are globle variables
  file_="$(unalias ls &>/dev/null; ls $target_d/$img_file.*-img* 2>/dev/null | sort | head -n 1)"
  file_basename="$(basename ${file_})"
  if [ -n "${file_}" ]; then
    if [ -e "$target_d/$img_file.dd-img" -o \
	 -n "$(ls $target_d/$img_file.dd-img.aa* 2>/dev/null)" -o \
	 -n "$(ls $target_d/$img_file.dd-ptcl-img.* 2>/dev/null)" ]; then
      if [ "$dd_checking_method" = "no-check" ]; then
        [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
        echo "Image was saved by dd or partclone.dd. No need or no way to check the file system integrity: $img_file" | tee --append ${OCS_LOGFILE}
        sleep 1
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
        return 0
      fi
    fi
    if [ -n "$(echo $file_basename | grep -Eo -- "-ptcl-img")" ]; then
      # new format, image file is like: sda1.ext4-ptcl-img.gz, sda1.ext4-ptcl-img.gz.aa
      fs_="$(echo $file_basename | sed -e "s/^$img_file\.//g" -e "s/-ptcl-img.*//g")"
    else
      # old format, image file is like: sda2.hfsp-img.aa  sda2.hfsp-img.ab  sda2.hfsp-img.ac
      fs_="$(echo $file_basename | sed -e "s/^$img_file\.//g" -e "s/-img.*//g")"
    fi
  fi
  if [ -z "${fs_}" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "The file system can not be decided in function img_check_by_partclone!!!"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop!"
    my_ocs_exit 1
  fi

  if [ -h "$file_" ]; then
    link_flag="yes"
  else
    link_flag="no"
  fi
 
  if [ -n "$(echo "$file_basename" | grep -Eo -- "-ptcl-img")" ]; then
    # New format
    if [ -n "$(echo "$file_basename" | grep -Eo -- "-ptcl-img.*.aa")" ]; then
      # New format with image split, e.g. sda1.ext4-ptcl-img.gz.aa
      get_image_cat_zip_cmd ${file_basename}
      # e.g. sda1.ext4-ptcl-img.gz.aa -> sda1.ext4-ptcl-img.gz.*
      img_file_prefix="$(echo ${file_basename} | sed -r -e "s/\.aa*$//").*"
      split_flag="yes"
    else
      # New format with image not split, e.g. sda1.ext4-ptcl-img.gz
      get_image_cat_zip_cmd ${file_basename}
      # The file is NOT split, so the file name is just like "sda1.ext4-ptcl-img.gz" only, no "."
      img_file_prefix="${file_basename}"
      split_flag="no"
    fi
  else
    # Old format
    # The split suffix length for old format is only 2, so we do not have to consider >=3 (*.aaa or more).
    if [ -f "$target_d/$img_file.${fs_}-img.aa" ]; then
      # Old format with image split
      get_image_cat_zip_cmd $target_d/$img_file.${fs_}-img.aa
      # The files are split, like hda1.00, hda1.01, so we have to add "."
      img_file_prefix="$img_file.${fs_}-img.*"
      split_flag="yes"
    else
      # Old format with image not split
      get_image_cat_zip_cmd $target_d/$img_file.${fs_}-img
      # The file is NOT split, so the file name is just like "hda1" only, no "."
      img_file_prefix="$img_file.${fs_}-img"
      split_flag="no"
    fi
  fi
  echo $msg_delimiter_star_line

  if [ "$link_flag" = "no" -a "$split_flag" = "no" ]; then
    # Special case: no link + no split. To get the best performance, remove cat,
    # and avoid the error the some compression program can not read link file, e.g., pixz does not have "-f" like zstdmt.
    check_flag="direct" 
  else
    check_flag="cat" 
  fi

  cmd_partclone_check="partclone.chkimg -L $ocs_log_dir/${img_file}-img-chk.log $PARTCLONE_CHKIMG_OPT -s -"

  # //NOTE// Here we force to use LC_ALL=C for partclone since we need to use get_partclone_image_info to parse the log file to get the rate. Only the keyword in English is supported in get_partclone_image_info.
  case "$check_flag" in
  cat)
    # Could be link + split image. Use for loop to cat the source image then unzip it.
    cat <<_EOT_ | tee -a $ocs_log_dir/ocs-chk-img.log
( for img in $target_d/$img_file_prefix; do
    cat \$img
  done
) | $unzip_stdin_cmd | LC_ALL=C $cmd_partclone_check
_EOT_
    ( for img in $target_d/$img_file_prefix; do
        cat $img
      done
    ) | $unzip_stdin_cmd | LC_ALL=C $cmd_partclone_check
    retcodes=(${PIPESTATUS[@]})
    if [ "${retcodes[1]}" -eq 0 -a "${retcodes[2]}" -eq 0 ]; then
      rc=0
    else
      rc=1
    fi
    ;;
  direct)
    # Not link and not split image. Can be read directly
    case "$unzip_stdin_cmd" in
      cat)
        # No compression, so the image file can be read directly.
        # Replace stdin "-s -" as the source file.
        cmd_partclone_check="$(echo $cmd_partclone_check | sed -r -e "s|-s -|-s $target_d/$img_file_prefix|")"
        echo "LC_ALL=C $cmd_partclone_check" | tee -a $ocs_log_dir/ocs-chk-img.log
        eval LC_ALL=C $cmd_partclone_check
	rc=$?
        ;;
      *)
        # Compressed image. Use unzip command to output stdout
        run_cmd_="$unzip_stdin_cmd $target_d/$img_file_prefix | LC_ALL=C $cmd_partclone_check"
        echo "$run_cmd_" | tee -a $ocs_log_dir/ocs-chk-img.log
        eval "$run_cmd_; retcodes=(\${PIPESTATUS[@]})"
        if [ "${retcodes[0]}" -eq 0 -a "${retcodes[1]}" -eq 0 ]; then
          rc=0
        else
          rc=1
        fi
        ;;
    esac
    ;;
  esac
  return $rc
} # end of img_check_by_partclone
#
do_partition_or_lv_img_check() {
  local target_d="$1"
  local img_file="$(to_filename $2)"
  local part="$(format_dev_name_with_leading_dev $3)"  # part is like /dev/sda1, /dev/mapper/ocs_luks_UbO
  local pt_swap rc img_file partition

  pt_swap="$(grep -Eo "^$part[[:space:]]+swap[[:space:]]+" $target_d/dev-fs.list 2>/dev/null)"

  if [ -f "$target_d/$img_file.000" -o \
       -n "$(ls $target_d/$img_file.aa* 2>/dev/null)" -o \
       -f "$target_d/$img_file" ]; then
    # The saved image is from partimage
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "This partition image was saved by partimage." | tee --append ${OCS_LOGFILE}
    echo "$msg_img_check_not_supported" | tee --append ${OCS_LOGFILE}
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    rc=1
    if [ "$ocs_batch_mode" != "on" ]; then
      echo -n "$msg_press_enter_to_continue..."
      read
    fi
  elif [ -f "$target_d/$img_file.ntfs-img" -o \
	 -n "$(ls $target_d/$img_file.ntfs-img.aa* 2>/dev/null)" ]; then
    # The saved image is from ntfsclone
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "This partition image was saved by ntfsclone." | tee --append ${OCS_LOGFILE}
    echo "$msg_img_check_not_supported" | tee --append ${OCS_LOGFILE}
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    rc=1
    if [ "$ocs_batch_mode" != "on" ]; then
      echo -n "$msg_press_enter_to_continue..."
      read
    fi
  elif is_partclone_image $target_d $img_file; then
    # The saved image is from partclone
    img_check_by_partclone
    rc=$?
    if [ "$rc" -gt 0 -a "$ocs_batch_mode" != "on" ]; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE | tee --append ${OCS_LOGFILE}
      echo "$msg_this_part_in_the_img_is_broken: $img_file" | tee --append ${OCS_LOGFILE}
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL | tee --append ${OCS_LOGFILE}
      echo -n "$msg_press_enter_to_continue..."
      read
    else
      [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING | tee --append ${OCS_LOGFILE}
      echo "$msg_this_part_in_the_img_is_ok: $img_file" | tee --append ${OCS_LOGFILE}
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL | tee --append ${OCS_LOGFILE}
    fi
  elif [ -e "$target_d/$img_file.dd-img" -o \
	 -n "$(ls $target_d/$img_file.dd-img.aa* 2>/dev/null)" -o \
	 -n "$(ls $target_d/$img_file.dd-ptcl-img.* 2>/dev/null)" ]; then
    if [ "$dd_checking_method" = "no-check" ]; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
      echo "Image was saved by dd or partclone.dd. No need or no way to check the file system integrity: $img_file" | tee --append ${OCS_LOGFILE}
      sleep 1
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      rc=0
    else
      # For some of the special partition, even it's dd, we do not need to check that. E.g. VMware vmkcore partition. It only need partition ID=fc, we do not care the content of the file system. Ref: http://communities.vmware.com/thread/85788
      # The saved image is from dd
      if [ -e "$target_d/$img_file.dd-img.info" ]; then
        # Get the paritition info, we will get "$partition"
        . $target_d/$img_file.dd-img.info
      fi
      if [ "$partition" = "VMware_vmkcore_partition" ]; then
        echo "VMware vmkcore partition found. No ndeed to check the file system." | tee --append ${OCS_LOGFILE}
        rc=0
      elif [ "$partition" = "BSD_slice" ]; then
        echo "BSD slice found. No ndeed to check the file system." | tee --append ${OCS_LOGFILE}
        rc=0
      elif [ "$partition" = "BSD_swap_partition" ]; then
        echo "BSD swap partition found. No ndeed to check the file system." | tee --append ${OCS_LOGFILE}
        rc=0
      else
        [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
        echo "This partition image was saved by dd." | tee --append ${OCS_LOGFILE}
        echo "$msg_img_check_not_supported"
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
        rc=1
        if [ "$ocs_batch_mode" != "on" ]; then
          echo -n "$msg_press_enter_to_continue..."
          read
        fi
      fi
    fi
  else
    # The last possibility is the swap info file.
    if [ -n "$pt_swap" ]; then
        task_check_swap_info_file $target_d $part
        rc=$?
        if [ "$rc" -eq 0 ]; then
          return 0
        fi
    else
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
      echo "This image file for partition $img_file was not found!"
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      rc=1
      if [ "$ocs_batch_mode" != "on" ]; then
        echo -n "$msg_press_enter_to_continue..."
        read
      fi
    fi
  fi

  echo $msg_delimiter_star_line
  return $rc
} # end of do_partition_or_lv_img_check
#
task_logv_check() {
  # mode is unicast or multicast
  local tgt_parts="$1"  # tgt_parts is like: hda1 hda2 hda5
  local volg is_in_chosen_partition lvm_tmp rc rctlv
  PV_PARSE_CONF="$target_dir_fullpath/lvm_vg_dev.list"
  LOGV_PARSE_CONF="$target_dir_fullpath/lvm_logv.list"
  if [ ! -f "$PV_PARSE_CONF" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "This LVM PV file NOT found: $PV_PARSE_CONF"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo -n "$msg_press_enter_to_continue..."
    read
  fi
  if [ ! -f "$LOGV_PARSE_CONF" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "This LVM LV file NOT found: $LOGV_PARSE_CONF"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo -n "$msg_press_enter_to_continue..."
    read
  fi
  
  #
  while read lv fs; do
   # Then we process the real data partition, only those in the chosen partitions
   # Ex:
   # /dev/vg3/lvol0  Linux rev 1.0 ext3 filesystem data (large files)
   # Then lvol0 is belong to VG vg3
   volg="$(echo "$lv" | awk -F"/" '{print $3}')"
   # Find if the LV is in the chosen partition (via VG, we can determine that)
   # EX: tgt_parts: hda1, hda3, hda5...
   #     vg3 /dev/hda3 nPMQQ0-D2yN-YRHL-9fBM-0cUm-vgcw-DCUTri
   is_in_chosen_partition="no"
   for ipt in $tgt_parts; do
     if [ -n "$(grep -E "[[:space:]]+/dev/$ipt[[:space:]]+" $PV_PARSE_CONF | grep -E "\<$volg\>")" ]; then
       # Found the chosen partitions is in the VG
       is_in_chosen_partition="yes"
       break
     fi
   done
   # If not in the chosen partition, skip this, continue with the next.
   [ "$is_in_chosen_partition" = "no" ] && continue
   fn="$(to_ocs_dev_img_name $lv)"
   # create the swap if it's swap partition
   case "$fs" in 
     *[Ss][Ww][Aa][Pp]*)
        echo $msg_delimiter_star_line
        echo "Found the swap partition $lv info:"
	# read LABEL, UUID info for $partition if swappt-$(to_filename ${fn}).info exists
        uuid_opt=""
        label_opt=""
	if [ -e "$target_dir_fullpath/swappt-$(to_filename ${fn}).info" ]; then
	  cat "$target_dir_fullpath/swappt-$(to_filename ${fn}).info"
          echo "Swap partition info file found!"
        else
          [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
          echo "$msg_this_swap_part_info_not_found: ${fn}"
          [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
          echo -n "$msg_press_enter_to_continue..."
          read
        fi
        echo $msg_delimiter_star_line
	# then skip the rest.
        continue;; 
   esac
   [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
   echo "Checking the device LV $fn in the image \"$(basename $target_dir_fullpath)\"..."
   [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
   do_partition_or_lv_img_check $target_dir_fullpath $fn $lv
   rc=$?
   rctlv="$(($rctlv + $rc))"
  done < $LOGV_PARSE_CONF
  return $rctlv
} # end of task_logv_check
#
task_check_swap_info_file() {
  local tgt_dir_fp="$1"
  local idev="$2" # e.g., /dev/sda1, /dev/mapper/ocs_luks_UbO
  local fn rc
  echo $msg_delimiter_star_line
  fn="$(to_ocs_dev_img_name $idev)"
  if [ -e "$tgt_dir_fp/swappt-${fn}.info" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "Found the swap partition $idev info file \"$tgt_dir_fp/swappt-${fn}.info\"."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo $msg_delimiter_star_line
    rc=0
  else
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "The swap partition $idev info file \"$tgt_dir_fp/swappt-${fn}.info\" was not found."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    sleep 2
    rc=1
  fi
  echo $msg_delimiter_star_line
  return $rc
} # end of task_check_swap_info_file
#
task_luks_img_check() {
  local lks_maps_rt="$*"  # e.g., ocs_luks_XYO ocs_luks_ABX
  local clted_luks_devs_pv_rt rc rcp rcplv

  rcp=0
  for id in $lks_maps_rt; do
    if grep -qsEw /dev/mapper/$id "$target_dir_fullpath/lvm_vg_dev.list"; then
      # Found LVM on LUKS mapping device. Process later.
      clted_luks_devs_pv_rt="$clted_luks_devs_pv_rt mapper/$id"
      continue
    fi
    do_partition_or_lv_img_check $target_dir_fullpath "$id" "mapper/$id"
    rc=$?
    rcp="$(($rcp + $rc))"
  done

  rcplv=0
  clted_luks_devs_pv_rt="$(echo $clted_luks_devs_pv_rt)" # Remove the leading spaces.
  # process LVM over LUKS if clted_luks_devs_pv_rt exists
  if [ -n "$clted_luks_devs_pv_rt" ]; then
    task_logv_check "$clted_luks_devs_pv_rt"
    rc=$?
    rcplv="$(($rcplv + $rc))"
  fi
  echo "End of task_luks_img_check."

  return $(( $rcp + $rcplv ))
} # end of task_luks_img_check
#
task_checkimage() {
  local target_dir="$1"
  local target_parts="$2"
  local parts_included target_dir_fullpath
  local p_type rc rct part_is_lvm ecryptfs_rc pt_swap

  if [ "$target_dir" = "ask_user" ]; then
    get_target_dir_name_when_checking_img_restorable # output: target_dir
  fi
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "$msg_the_image_to_be_cheked: $target_dir"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL

  ecryptfs_rc="1"
  if is_ecryptfs_img $ocsroot/$target_dir; then
    # If it's encrypted image, we have to decrypt it.
    ocs_sr_type="restore"
    # //NOTE// If encrypt_ocs_img="yes", after this step, ocsroot and target_dir will be changed
    # The original ones will be kept as ocsroot_orig and target_dir_orig.
    prepare_ecryptfs_mount_point_if_necessary 
    ecryptfs_rc="$?"
    if [ "$ecryptfs_rc" -ne 0 ]; then
      echo "$msg_program_stop"
      my_ocs_exit 1
    fi
  fi

  #
  source_parts_no="$(get_parts_list_from_img $ocsroot/$target_dir | sed -e "s/ *$//g" | wc -w | awk '{print $1}')"

  #
  if [ "$target_parts" = "ask_user" ]; then
    ANS_TMP=`mktemp /tmp/ocs_ans.XXXXXX`
    trap "[ -f "$ANS_TMP" ] && rm -f $ANS_TMP" HUP INT QUIT TERM EXIT
    get_existing_partitions_from_img $ANS_TMP $ocsroot/$target_dir no check
    # we have to remove " (comes with checklist in dialog) 
    # so that for loop will work (Specially for FC3/4...)
    target_parts="$(cat $ANS_TMP | tr -d \")"
    [ -f "$ANS_TMP" ] && rm -f $ANS_TMP
    # target name exists, but file "parts" is empty ?
    check_target_parts $ocsroot/$target_dir/parts "$target_parts"
  elif [ -z "$target_parts" ]; then
    target_parts="$(get_parts_list_from_img $ocsroot/$target_dir)"
  else
    echo "$msg_the_requested_partition_to_be_checked: $target_parts"
  fi

  # target_hd will be extract from $target_parts, maybe we will have one more
  # find the target hd
  # maybe we will have one more hd (like hda1, hda2, hdb1, hdb3 -> hda, hdb)
  for ipart in $target_parts; do
    thd_tmp="$(get_diskname $ipart)"
    if [ -z "$target_hd" ]; then
      target_hd="$thd_tmp"
    elif [ -z "$(echo $target_hd | grep -Ew "$thd_tmp" 2>/dev/null)" ]; then
      target_hd="$target_hd $thd_tmp"
    fi
  done

  #
  if [ -e "$ocsroot/$target_dir/disk" ]; then
    image_type="disk"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "$msg_a_disk_image_saved_from_hd: $(get_disk_list_from_img $ocsroot/$target_dir)"
    echo "$msg_the_part_image_to_be_cheked: $target_parts"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  elif [ -e "$ocsroot/$target_dir/parts" ]; then
    image_type="partition"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "$msg_a_part_image_saved_from_partitions: $(get_parts_list_from_img $ocsroot/$target_dir)"
    echo "$msg_the_part_image_to_be_cheked: $target_parts"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  else
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "$msg_broken_image_dir"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop!" | tee --append ${OCS_LOGFILE}
    my_ocs_exit 1
  fi

  # Use $target_dir_fullpath as the absolute path, i.e. put leading $ocsroot, since from now on we need the full path dir to access them.
  target_dir_fullpath="$ocsroot/$target_dir"
  check_input_target_image "$target_dir_fullpath"

  screen_not_blank

  # if $create_part (global variable) is no, only some preparations 
  # in create_partition, it won't run sfdisk in fact.
  
  # strip the leading spaces
  target_parts="$(echo $target_parts | sed -e "s/^[[:space:]]*//g")"

  echo $msg_delimiter_star_line
  # Start checking partition table.
  echo "Checking the partition table in the image \"$target_dir\"..."
  for ihd in $target_hd; do
    echo "Checking $ihd..."
    # Get the partition type
    p_type="$(get_partition_table_type_from_img "$target_dir_fullpath" "${ihd}")"
    # Only check mbr or gpt, skip pv_disk.
    case "$p_type" in
     mbr|gpt)
      echo "Partition table type: $p_type" | tee --append ${OCS_LOGFILE}
      if [ ! -e "$target_dir_fullpath/$(to_filename ${ihd})-pt.sf" ]; then
        [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
        echo "$msg_this_partition_table_file_not_found: $(to_filename ${ihd})-pt.sf" | tee --append ${OCS_LOGFILE}
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
        echo -n "$msg_press_enter_to_continue..."
        read
      else
	echo "$msg_this_partition_table_file_found: ${ihd}, $target_dir_fullpath/$(to_filename ${ihd})-pt.sf" | tee --append ${OCS_LOGFILE}
      fi
      ;;
     unknown)
      if is_md_dev $ihd; then
        [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
        echo "Found Linux Software RAID: $ihd. Assume no need to check partition table."
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
	continue
      fi
      [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
      if [ -e "target_dir_fullpath/$(to_filename ${hdsk})-pt.parted" ]; then
        echo "Unknown partition table format from the file $target_dir_fullpath/$(to_filename ${hdsk})-pt.parted!" | tee --append ${OCS_LOGFILE}
      else
        echo "File $target_dir_fullpath/$(to_filename ${hdsk})-pt.parted not found!" | tee --append ${OCS_LOGFILE}
	echo "Assume it's for whole disk as a file system or RAID cases..."
      fi
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      continue
      ;;
    esac
  done

  echo $msg_delimiter_star_line
  # When Clonezilla does not save a whole disk image (i.e. use the mode saveparts), the swap partition will be skipped. Therefore only check if swap info files exist when it is savedisk.
  if [ "$image_type" = "disk" ]; then
    echo $msg_delimiter_star_line
    echo "Checking the swap partition info in the image \"$target_dir\"..."
    # Check swap partition
    for ihd in $target_hd; do
      echo "Checking $ihd..."
      p_type="$(get_partition_table_type_from_img "$target_dir_fullpath" "${ihd}")"
      # Only check mbr or gpt, skip pv_disk.
      case "$p_type" in
       mbr|gpt)
        echo "Finding swap partition(s) in file $target_dir_fullpath/$(to_filename ${ihd})-pt.sf..."
        for partition in `get_swap_partition_sf_format $target_dir_fullpath/$(to_filename ${ihd})-pt.sf`; do
          echo "Checking swap partition $partition..."
	  # read LABEL, UUID info for $partition if swappt-$(to_filename $partition).info exists
          uuid_opt=""
          label_opt=""
	  if [ -e "$target_dir_fullpath/swappt-$(to_filename $partition).info" ]; then
	    cat "$target_dir_fullpath/swappt-$(to_filename $partition).info"
            echo "Swap partition info file found!"
          else
            [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
            echo "$msg_this_swap_part_info_not_found: $partition"
            [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
            echo -n "$msg_press_enter_to_continue..."
            read
          fi
        done
        ;;
       unknown)
        [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
        echo "Unknown partition table format from the file $target_dir_fullpath/$(to_filename ${ihd})-pt.parted!" | tee --append ${OCS_LOGFILE}
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
        echo "$msg_program_stop!" | tee --append ${OCS_LOGFILE}
        my_ocs_exit 1
        ;;
      esac
    done
  else
    echo "This is not an image for whole disk. Swap partition info file might not be checked..." | tee --append ${OCS_LOGFILE}
  fi

  echo $msg_delimiter_star_line
  # Check MBR
  echo "Checking the MBR in the image \"$target_dir\"..."
  # Reinstall executable code area (first 446 bytes in MBR)
  for ihd in $target_hd; do
    p_type="$(get_partition_table_type_from_img "$target_dir_fullpath" "${ihd}")"
    # Ref: http://en.wikipedia.org/wiki/Master_boot_record
    # Master Boot Record (MBR) is the 512-byte boot sector:
    # 446 bytes (executable code area) + 64 bytes (table of primary partitions) + 2 bytes (MBR signature; # 0xAA55) = 512 bytes.
    # However, some people also call executable code area (first 446 bytes in MBR) as MBR.
    if [ "$p_type" = "gpt" ]; then
      echo "GPT disk. No need to check MBR."
      continue
    fi
    if is_md_dev $ihd; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
      echo "Found Linux Software RAID device: $ihd. Assume no need to check MBR."
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      continue
    fi
    if [ ! -e "$target_dir_fullpath/$(to_filename ${ihd})-mbr" ]; then
       [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
       echo "$msg_mbr_img_of_this_hd_not_found: ${ihd}"
       [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
       echo -n "$msg_press_enter_to_continue..."
       read
    else
       echo "MBR file for this disk was found: ${ihd}" | tee --append ${OCS_LOGFILE}
    fi
  done

  echo $msg_delimiter_star_line
  PV_PARSE_CONF="$target_dir_fullpath/lvm_vg_dev.list"
  do_LVM_check="no"
  do_LUKS_check="no"
  rct=0
  for partition in $target_parts; do
    lks_mapper=""
    # sda1 -> sda
    hd_tmp="$(get_diskname $partition)"
    part_is_lvm="no"
    # If we partition is listed in lvm_vg_dev.list, process LVM later. //NOTE// LVM might use Id=83 instead of 8e, so we can not parse it based on Id.
    if [ -e "$PV_PARSE_CONF" ]; then
      for i in $(sed -e 's!^.*/dev/\([^[:space:]]\{3,\}\)[[:space:]]*.*$!\1!g' $PV_PARSE_CONF); do
        if [ "$partition" = "$(get_master_dev_of_multipath $i)" ]; then
          part_is_lvm="yes"
          break
        fi
      done
    fi
    # If part_is_lvm is yes, we should skip the rest...
    if [ "$part_is_lvm" = "yes" ]; then
      do_LVM_check="yes"
      # This continue is to skip the rest of this partition do loop
      continue
    fi

    # LUKS
    if [ -e "$target_dir_fullpath/luks-dev.list" ]; then
      if [ -n "$(grep -Ew /dev/$partition $target_dir_fullpath/luks-dev.list)" ]; then
	if [ -e "$target_dir_fullpath/${partition}-luksHeader.bin" ]; then
          echo "$msg_luks_header_file_found: $target_dir_fullpath/${partition}-luksHeader.bin" | tee --append ${OCS_LOGFILE}
	else
          [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
          echo "$msg_luks_header_file_not_found: $target_dir_fullpath/${partition}-luksHeader.bin" | tee --append ${OCS_LOGFILE}
          [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
          echo "$msg_program_stop!"
          my_ocs_exit 1
	fi
        do_LUKS_check="yes"
        if [ -e "$target_dir_fullpath/luks-crypttab-from-OS.info" ]; then
          echo "$msg_luks_crypttab_file_found: $target_dir_fullpath/luks-crypttab-from-OS.info" | tee --append ${OCS_LOGFILE}
          lks_mapper="$(LC_ALL=C get_luks_mapped_dev_name_from_os from-img /dev/$partition $target_dir_fullpath)"
          if [ -z "$lks_mapper" ]; then
            [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
            echo "$msg_luks_mapped_name_not_found_in_this_file: $target_dir_fullpath/luks-crypttab-from-OS.info."
	    echo "$msg_name_it_arbitrarily_restered_os_boot_might_fail"
            [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
	    countdown 5
	    echo
	  fi
        else
          echo "$msg_luks_crypttab_file_not_found: $target_dir_fullpath/luks-crypttab-from-OS.info. Use $target_dir_fullpath/luks-dev.list" | tee --append ${OCS_LOGFILE}
	fi
        if [ -z "$lks_mapper" ]; then 
          lks_mapper="$(LC_ALL=C grep -Ew /dev/$partition $target_dir_fullpath/luks-dev.list | awk -F " " '{print $3}' | xargs basename)"
        fi
        clted_luks_devs_map_rt="$clted_luks_devs_map_rt $lks_mapper"
        # This continue is to skip the rest of this partition do loop
        continue
      fi
    fi
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "Checking the partition $partition in the image \"$target_dir\"..."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    do_partition_or_lv_img_check $target_dir_fullpath $partition /dev/$partition
    rc=$?
    rct="$(($rct + $rc))"
  done

  # Check LVM
  if [ "$do_LVM_check" = "yes" ]; then
    # LVM exists, check PV/VG/LV.
    task_logv_check "$target_parts"
    rc=$?
    rct="$(($rct + $rc))"
    echo $msg_delimiter_star_line | tee --append ${OCS_LOGFILE}
  fi

  # Check LUKS
  if [ "$do_LUKS_check" = "yes" ]; then
    # LUKS exists, restore it
    echo "LUKS device image exists, checking it..."
    task_luks_img_check $clted_luks_devs_map_rt
    rc=$?
    rct="$(($rct + $rc))"
    echo $msg_delimiter_star_line | tee --append ${OCS_LOGFILE}
  fi

  if [ "$rct" -gt 0 -a "$ocs_batch_mode" != "on" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE | tee --append ${OCS_LOGFILE}
    echo "$msg_broken_partition_img_found: $target_dir" | tee --append ${OCS_LOGFILE}
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL | tee --append ${OCS_LOGFILE}
    echo -n "$msg_press_enter_to_continue..."
    read
  else
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING | tee --append ${OCS_LOGFILE}
    echo "$msg_all_the_part_lv_are_ok: $target_dir" | tee --append ${OCS_LOGFILE}
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL | tee --append ${OCS_LOGFILE}
  fi
  return $rct
} # end of task_checkimage

####################
### Main program ###
####################

ocs_file="$0"
ocs=`basename $ocs_file`
#
while [ $# -gt 0 ]; do
 case "$1" in
   -b|--batch) ocs_batch_mode="on"; shift;;
   -or|--ocsroot)
           # overwrite the ocsroot in drbl.conf
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             ocsroot="$1"
             shift;
           fi
           [ -z "$ocsroot" ] && USAGE && exit 1
           ;;
   -nogui|--nogui)
           shift; 
           # -nogui is for backward compatable, better to use --nogui
           nogui="on"
           ;;
   -*)     echo "${0}: ${1}: invalid option" >&2
           USAGE >& 2
           exit 2 ;;
   *)      break ;;
 esac
done

ocs_chk_img_name="$1"
shift
ocs_chk_dev="$(strip_leading_dev $*)"

force_TERM_as_linux_if_necessary

#
check_if_root
ask_and_load_lang_set

# Create log dir
mkdir -p $ocs_log_dir

# imagedir is a variable which ask_user related function need
imagedir="$ocsroot"
if [ -z "$ocs_chk_img_name" ]; then
  ocs_chk_img_name="ask_user"
  ocs_sr_mode="interactive"
fi
# if "$ocs_chk_dev" is nothing, we check all the partitions.

#
if [ "$nogui" = "off" ]; then
  # TUI is on.
  PARTCLONE_CHKIMG_OPT="$PARTCLONE_CHKIMG_OPT -N"
fi

log_line_start="$(LC_ALL=C wc -l $OCS_LOGFILE | cut -d" " -f1)"

# Clean the previous saved env so if ocs-onthefly is cancelled, won't read the previous env and reboot/poweroff.
rm -f /var/lib/clonezilla/ocs-vars

task_checkimage "$ocs_chk_img_name" "$ocs_chk_dev"
rc=$?

# Save the mode
save_ocs_related_vars

log_line_start="$((log_line_start + 1))"
# Output the log summary.
echo "Summary of image checking:"
echo "=========================="
tail -n +$log_line_start $OCS_LOGFILE
echo "=========================="

my_ocs_exit $rc
