#!/bin/bash
# Author: Steven Shiau <steven _at_ nchc org tw>
# License: GPL

# This program is used to restore an image to multiple disks, especially many USB flash drives in a same machine by using Clonezilla live.
# Thanks to T. C. Lin <tclin _at_ mail dfes tpc edu tw> for this idea and testing, and Alvin Su for debuging.

# For more details, please refer to:
# http://drbl.sourceforge.net/faq/fine-print.php?path=./2_System/98_one_image_to_multiple_disks.faq#98_one_image_to_multiple_disks.faq
# and
# http://drbl.sourceforge.net/screenshot/?in_path=/12_Alvin_Su_Taiwan

# ///WARNING/// This program is really dangerous! You have to know what you are doing! Especially the destination disks!

# Load functions and config file
DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/opt/drbl/}"

. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
. $DRBL_SCRIPT_PATH/conf/drbl-ocs.conf
. $DRBL_SCRIPT_PATH/sbin/ocs-functions

# Settings
batch_mode="false"

#
USAGE() {
    echo "$prog - To restore an image to multiple disks"
    echo "Usage:"
    echo "$prog [OPTION] IMAGE_NAME DEV"
    echo "IMAGE_NAME  The image name, i.e. the image dir in $ocsroot/"
    echo "DEV         The destination device name, e.g. sdc, sdd, sde..."
    echo "OPTION:"
    echo "-b, --batch-mode   Run in batch mode"
    echo "Ex:"
    echo "To restore the image \"usb-image\" to 3 USB flash drives sda, sdb, and sdc, you can run:"
    echo "  $prog usb-image sda sdb sdc"
} # end of USAGE

####################
### Main Program ###
####################
prog="$(basename $0)"

# Parse command-line options
while [ $# -gt 0 ]; do
  case "$1" in
    -b|--batch) batch_mode="true"; shift;;
    -um|--user-mode)
            shift
            # skip the -xx option, in case 
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              ocs_user_mode="$1"
              shift
            fi
            [ -z "$ocs_user_mode" ] && USAGE && exit 1
	    ;;
    -*)     echo "${0}: ${1}: invalid option" >&2
            USAGE >& 2
            exit 2 ;;
    *)      break ;;
  esac
done
#
image_name="$1"
shift
tgt_dsks=$*

check_if_root
ask_and_load_lang_set
#
[ -z "$image_name" ] && image_name="ask_user"
[ -z "$tgt_dsks" ] && tgt_dsks="ask_user"

#
if [ "$batch_mode" = "false" ]; then
  [ -z "$ocs_user_mode" ] && ask_if_beginner_or_expert_mode
fi

run_again_flag="false"
if [ "$image_name" = "ask_user" ]; then
  imagedir="$ocsroot"
  get_target_dir_name_when_restoring_disk  # get $target_dir
  image_name="$target_dir"
  run_again_flag="true"
fi
if [ "$tgt_dsks" = "ask_user" ]; then
  # To get $target_hd
  get_target_hd_name_from_local_machine "$msg_choose_the_disks_to_restore \n$msg_linux_disk_naming $msg_press_space_to_mark_selection" checklist
  tgt_dsks="$target_hd"
  run_again_flag="true"
fi
#
if [ "$run_again_flag" = "true" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_SUCCESS
  echo "PS. $msg_run_drbl_ocs_again_cmd"
  echo "$prog $image_name $tgt_dsks"
  echo "$msg_if_you_want_to_run_in_unattended_mode"
  echo "$prog -b $image_name $tgt_dsks"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
fi

#
# Check inupt values
# 1. image name
check_input_target_image "$ocsroot/$image_name"
# 2. destination disks
check_input_hd $tgt_dsks

# Give warnings. 
echo $msg_delimiter_star_line
if [ "$batch_mode" = "false" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "$msg_uppercase_Warning!!! $msg_uppercase_Warning!!! $msg_uppercase_Warning!!!"
  echo "$msg_uppercase_Warning! $msg_all_data_in_dev_will_be_overwritten: $tgt_dsks"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  echo "$msg_are_u_sure_u_want_to_continue"
  echo -n "[y/N] "
  read continue_restoring
  case "$continue_restoring" in
  	Y|y|[yY][eE][sS])
  	   echo "$msg_ok_let_do_it!"
  	   ;;
  	*)
  	   echo "$msg_program_stop!!!"
  	   exit 0
  esac
fi

# Calculate the disk number
disk_no="$(LC_ALL=C echo $tgt_dsks | wc -w)"

# Find the original disk
src_disk="$(LC_ALL=C cat $ocsroot/$image_name/disk)"
src_disk_no="$(LC_ALL=C echo $src_disk | wc -w)"
if [ "$src_disk_no" -ne 1 ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
  echo "$msg_source_disks_more_than_1"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  echo "$msg_program_stop"
  exit 1
fi

# Prepare the temp images
for i in $tgt_dsks; do
  create-ocs-tmp-img -t $ocsroot $image_name ${image_name}_cnvt_${i} $src_disk ${i}
done

# Create partition table first. If we do not do this first, sfdisk will wait for the lock file to be released, so the 2nd clone won't start in any case.
for i in $tgt_dsks; do
  if [ -z "$(grep -Ew "$i" /proc/partitions)" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "Disk /dev/$i not found. Skip restoring it."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    continue
  fi
  tgt_hd_file="/dev/$i"
  tgt_dir="$ocsroot/${image_name}_cnvt_${i}"
  if [ -n "$(grep -iE "^Partition Table:" $tgt_dir/${i}-pt.parted 2>/dev/null | grep -iE "gpt")" ]; then
    pt_type="gpt"
  else
    pt_type="mbr"
  fi
  echo "Creating the partition table for $i..."
  case "$pt_type" in
  mbr)
    # If the partition table is tag as "gpt", change it as msdos
    if [ -n "$(LC_ALL=C parted -s $tgt_hd_file print | grep -iE "^Partition Table:" | grep -iE "gpt")" ]; then
      echo "The partition table in $tgt_hd_file is for GPT, now make it for MBR."
      LC_ALL=C parted -s $tgt_hd_file mklabel msdos
    fi
    sfdisk --force $tgt_hd_file < $tgt_dir/${i}-pt.sf
    inform_kernel_partition_table_changed mbr $tgt_hd_file
    ;;
  gpt)
    clean_mbr_partition_table_before_gen_gpt_table $tgt_hd_file
    create_gpt_table_if_no_table $tgt_hd_file
    echo "Running: sgdisk -l $tgt_dir/${i}-gpt.gdisk $tgt_hd_file"
    LC_ALL=C sgdisk -l $tgt_dir/${i}-gpt.gdisk $tgt_hd_file
    sgdisk_rc=$?
    if [ "$sgdisk_rc" -ne 0 ]; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
      echo "Failed to create partition table on $i."
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    fi
    echo "Informing kernel that the OS that partition table has changed..."
    inform_kernel_partition_table_changed gpt $tgt_hd_file
    ;;
  esac
done

for i in $tgt_dsks; do
  [ -z "$(grep -Ew "$i" /proc/partitions)" ] && continue
  echo "Restoring image for disk $i..."
  # use -k to skip sfdisk since we have done that in the previous step.
  ocs-sr -g auto -nogui -b -k -e1 auto -e2 -r -j2 -p true restoredisk "${image_name}_cnvt_${i}" "$i" &
done

# When can we remove the converted temp image dir?
# Clean the temp image dir
#for i in $tgt_dsks; do
#  rm -rf $ocsroot/${image_name}_${i}
#done
