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

# Load the setting
DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"

. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
#
check_if_root

# functions
USAGE() {
  echo "Usage: $0 [Option] [on|off]"
  echo "Options:"
  language_help_prompt_by_idx_no
  echo "-n, --no-check-clients    Do NOT check if clients' directories exist in $drblroot."
  echo "-c, --no-create-ssi-template  Do NOT create DRBL single system image template again. Default is to create SSI template when on."
  echo "-d, --drbl-mode [0|1|2]  Assign DRBL mode (0: Full DRBL mode, 1: DRBL SSI mode, 2: Do NOT provide diskless Linux service to clients)."
  echo "-t, --template HOST  Use HOST as the template to create etc and var tarballs. If not assigned, $0 will find the first client one in dir $drblroot."
  echo "-x, --exclude-X-config    When creating template, exclude X config."
  echo "-s, --no-stop-ocs    do NOT stop ocs. (It's better to stop OCS so that the unnecessary services in rc1.d will be removed!)."
  echo "-z, --clonezilla-mode  [0|1|2|3]  Assign Clonezilla mode (0: Full DRBL mode, 1: Clonezilla box mode, 2: Do NOT provide clonezilla service to clients, 3: Use Clonezilla live as the OS of clients)"
  echo "-h, --help           Display this help and exit"
  echo "Example:"
  echo "To turn on the DRBL SSI mode, run:"
  echo "$0 -d 1"
  echo "To turn on the DRBL SSI mode and Clonezilla box mode, run:"
  echo "$0 -d 1 -z 1"
} # end of USAGE
# 
set_clonezilla_mode() {
  local mode=$1
  echo -n "Setting clonezilla_mode=\"$mode\" in /etc/drbl/drbl_deploy.conf and /etc/drbl/drblpush.conf... "
  perl -pi -e "s/^[[:space:]]*clonezilla_mode=.*/clonezilla_mode=\"$mode\"/g" /etc/drbl/drbl_deploy.conf
  perl -pi -e "s/^[[:space:]]*clonezilla_mode=.*/clonezilla_mode=$mode/g" /etc/drbl/drblpush.conf
  echo "done!"
} # end of set_clonezilla_mode
#
set_drbl_mode() {
  local mode=$1
  echo -n "Setting drbl_mode=\"$mode\" in /etc/drbl/drbl_deploy.conf and /etc/drbl/drblpush.conf... "
  perl -pi -e "s/^[[:space:]]*drbl_mode=.*/drbl_mode=\"$mode\"/g" /etc/drbl/drbl_deploy.conf
  perl -pi -e "s/^[[:space:]]*drbl_mode=.*/drbl_mode=$mode/g" /etc/drbl/drblpush.conf
  echo "done!"
} # end of set_drbl_mode
#
#
create_drbl_ssi_clonezilla_box_template() {
  # if template is not assigned, find one.
  if [ -z "$template" ]; then
    for ih in $drblroot/*; do
      # use the 1st one drbl client we found as template
      if [ -d "$ih" ]; then
        template="$(basename $ih)"
        break
      fi
    done
  fi
  echo "Creating template tarball for DRBL SSI and Clonezilla box..."
  [ "$exclude_template_X_conf" = "yes" ] && gen_ssi_extra_opt="-x"
  [ "$stop_ocs" = "no" ] && gen_ssi_extra_opt="$gen_ssi_extra_opt --no-stop-ocs"
  drbl-gen-ssi-files -t $template $gen_ssi_extra_opt
}

#
check_full_drbl_mode_clients() {
  client_ip_lists="$(LC_ALL=C get-client-ip-list)"
  for ih in $client_ip_lists; do
    if [ ! -d $drblroot/$ih/etc ]; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
      echo "$drblroot/$ih does NOT exist!!! You have to run \"drblpush -i\" again and choose to use Full DRBL or Full Clonezilla mode. We can NOT turn on Full DRBL or Full Clonezilla mode!"
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      echo "Program terminated!!!"
      exit 1
    fi
  done
}

# default setting
template=""
check_clients_files="yes"
create_ssi_template="yes"
stop_ocs="yes"

# Parse command-line options
while [ $# -gt 0 ]; do
  case "$1" in
    -l|--language)
	shift
        if [ -z "$(echo $1 |grep ^-.)" ]; then
          # skip the -xx option, in case 
	  specified_lang="$1"
	  shift
        fi
	;;
    -s|--no-stop-ocs)
	    stop_ocs="no"
	    shift;;
    -n|--no-check-clients)
	    check_clients_files="no"
	    shift;;
    -c|--no-create-ssi-template)
	    create_ssi_template="no"
	    shift;;
    -x|--exclude-X-config)
	    exclude_template_X_conf="yes"
	    shift;;
    -t|--template)
 	    shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
	      template="$1"
	      shift
            fi
	    ;;
    -d|--drbl-mode)
 	    shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
	      drbl_mode="$1"
	      shift
            fi
	    [ -z "$drbl_mode" ] && USAGE && exit 1
	    ;;
    -z|--clonezilla-box-mode)
 	    shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
	      clonezilla_mode="$1"
	      shift
            fi
	    [ -z "$clonezilla_mode" ] && USAGE && exit 1
	    ;;
    -h|--help)  
            USAGE >& 2
            exit 2 ;;
    -*)     echo "${0}: ${1}: invalid option" >&2
            USAGE >& 2
            exit 2 ;;
    *)      break ;;
  esac
done

#
ask_and_load_lang_set $specified_lang

# Set the config files for all the process.
PXE_CONF=$pxecfg_pd/pxelinux.cfg/default
GRUB_CONF=$GRUB_EFINB_DIR/grub.cfg

case "$drbl_mode" in
  0) # full drbl mode
    [ "$check_clients_files" = "yes" ] && check_full_drbl_mode_clients
    echo "Full DRBL mode. Remove clientdir opt for label drbl in pxelinux config... "
    # For PXE
    hide_reveal_pxe_img drbl reveal $PXE_CONF
    remove_opt_in_pxelinux_cfg_block drbl clientdir
    # For uEFI
    # Older version of GNU/Linux, like CentOS 6 does not support grub uEFI network boot.
    if [ -e "$GRUB_CONF" ]; then
      hide_reveal_grub_efi_ent drbl-client reveal $GRUB_CONF
      remove_opt_in_grub_efi_cfg_block drbl-client clientdir
    fi

    set_drbl_mode full_drbl_mode
    ;;
  1) # drbl ssi mode
    echo "DRBL SSI mode. Set clientdir opt for label drbl in pxelinux config... "
    # For PXE
    hide_reveal_pxe_img drbl reveal $PXE_CONF
    add_opt_in_pxelinux_cfg_block drbl clientdir node_root
    # For uEFI
    # Older version of GNU/Linux, like CentOS 6 does not support grub uEFI network boot.
    if [ -e "$GRUB_CONF" ]; then
      hide_reveal_grub_efi_ent drbl-client reveal $GRUB_CONF
      add_opt_in_grub_efi_cfg_block drbl-client clientdir node_root
    fi

    set_drbl_mode drbl_ssi_mode
    ;;
  2) # none
    echo "DRBL service is set as unavailable."
    create_ssi_template="no"
    # Remove clientdir then hide it
    # For PXE
    remove_opt_in_pxelinux_cfg_block drbl clientdir
    hide_reveal_pxe_img drbl hide $PXE_CONF
    # For uEFI
    # Older version of GNU/Linux, like CentOS 6 does not support grub uEFI network boot.
    if [ -e "$GRUB_CONF" ]; then
      remove_opt_in_grub_efi_cfg_block drbl-client clientdir
      hide_reveal_grub_efi_ent drbl-client hide $GRUB_CONF
    fi

    set_drbl_mode none
    ;;
esac

case "$clonezilla_mode" in
  0) # full clonezilla mode
    [ "$check_clients_files" = "yes" ] && check_full_drbl_mode_clients
    echo "Full clonezilla mode. Remove clientdir opt for label clonezilla in pxelinux config... "
    # For PXE
    remove_opt_in_pxelinux_cfg_block clonezilla clientdir
    # For uEFI
    # Older version of GNU/Linux, like CentOS 6 does not support grub uEFI network boot.
    if [ -e "$GRUB_CONF" ]; then
      remove_opt_in_grub_efi_cfg_block clonezilla-se-client clientdir
    fi

    # TODO: how do we know to use it's as full_clonezilla_mode or none ?
    set_clonezilla_mode full_clonezilla_mode
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "$msg_run_dcs_clonezilla_start_to_show_pxemenu"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    ;;
  1) # clonezilla box mode
    echo "Clonezilla box mode. Set clientdir opt for label clonezilla in pxelinux config... "
    # For PXE
    add_opt_in_pxelinux_cfg_block clonezilla clientdir node_root
    # For uEFI
    # Older version of GNU/Linux, like CentOS 6 does not support grub uEFI network boot.
    if [ -e "$GRUB_CONF" ]; then
      add_opt_in_grub_efi_cfg_block clonezilla-se-client clientdir node_root
    fi

    set_clonezilla_mode clonezilla_box_mode
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "$msg_run_dcs_clonezilla_start_to_show_pxemenu"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    ;;
  2) # none
    echo "Clonezilla service is set as unavailable."
    create_ssi_template="no"
    # Remove clientdir then hide it
    # For PXE
    remove_opt_in_pxelinux_cfg_block clonezilla clientdir
    # For uEFI
    # Older version of GNU/Linux, like CentOS 6 does not support grub uEFI network boot.
    if [ -e "$GRUB_CONF" ]; then
      remove_opt_in_grub_efi_cfg_block clonezilla-se-client clientdir
    fi

    set_clonezilla_mode none
    ;;
  3) # clonezilla-live mode
    echo "Use Clonezilla live as the OS of clients. Set clientdir opt for label clonezilla in pxelinux config... "
    create_ssi_template="no"
    # No need to add any options in this mode. Later drbl-ocs-live-prep will add the required options.

    set_clonezilla_mode clonezilla_live_mode
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "$msg_run_dcs_clonezilla_start_to_show_pxemenu"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    ;;
esac
# No matter which mode of Clonezilla, we hide the Clonezilla menu first. Only dcs will reveal or hide the menu
# For PXE
hide_reveal_pxe_img clonezilla hide $PXE_CONF
hide_reveal_pxe_img Clonezilla-live hide $PXE_CONF

# For uEFI
# Older version of GNU/Linux, like CentOS 6 does not support grub uEFI network boot.
if [ -e "$GRUB_CONF" ]; then
  hide_reveal_grub_efi_ent clonezilla-se-client hide $GRUB_CONF
  hide_reveal_grub_efi_ent "clonezilla-live-client" hide $GRUB_CONF
fi

#
if  [ "$create_ssi_template" = "yes" ]; then
  if [ -n "$(echo "$drbl_mode" | grep -Ew "(1|2)")" -o \
       -n "$(echo "$clonezilla_mode" | grep -Ew "(1|2)")" ]; then
       create_drbl_ssi_clonezilla_box_template
  fi
fi

exit 0