#!/bin/ash

OUTPUT="data.txt" # used to get selected items on checklist dialog

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

# displays an error dialog and reboots computer or exits program
# all the parameters are used as message to be displayed in dialog
perror () {
  echo "ERROR: $@" 1>&2
  if [ "$SLX_REMOVEPART_AUTO" != "yes" ]; then
    dialog --title "ERROR" --stdout --msgbox "$@" 15 60
  else
    dialog --title "ERROR" --no-cancel --stdout --pause "$@\n\nReboot in:" 15 60 600
    reboot
  fi
  exit 1
}

# displays a warning dialog. if mode = automatic holds the infobox for 5s
# all the parameters are used as message to be displayed in dialog
pwarning  () {
  echo "WARNING: $@" 1>&2
  if [ "$SLX_REMOVEPART_AUTO" != "yes" ]; then
    dialog --title "WARNING" --stdout --msgbox "$@" 15 60
  else
    dialog --title "WARNING" --stdout --infobox "$@" 15 60
    sleep 5
  fi
}

. /opt/openslx/config || perror "Could not source config."

# removes a partition from disk. uses fdisk for mbr and sgdisk for gpt
# param $1 is the disk to remove partition from
# param $2 is the format of the $1 disk
# param $3 is the number of the partition to be removed
remove_part_delete () {
  unmount -f $1$3 > /dev/null 2>&1
  if [ "$2" != "gpt" ]; then
    echo "[REMOVE] Remove partition $3 from msdos disk $1"
    (echo d; echo ${3}; echo w) | fdisk $disk &>/dev/null
  else
    echo "[REMOVE] Remove partition $3 from gpt disk $1"
    sgdisk $1 -d $3 &>/dev/null
  fi
}

# display a confirmation dialog then removes the partition (see remove_part_delete)
# param $1 is the disk to remove partition from
# param $2 is the format of the $1 disk
# param $3 is the number of the partition to be removed
# param $4 is the partition name (used to show to user which partition is getting deleted)
remove_part_confirm () {
  dialog --title "REMOVE PARTITION (USB DISK)" --yesno "Are you sure you want to delete $4 from removable disk $1?" 15 60

  ret=$?
  case $ret in
    0)
      remove_part_delete $1 $2 $3;;
    1)
      echo "User does not want to remove $4 from $1";;
    255)
      echo "ESC pressed";;
  esac
}

# handler for removing a partition. checks if usb and call apropriate functions (see remove_part_delete and remove_part_confirm).
# param $1 is the disk to remove partition from
# param $2 is the format of the $1 disk
# param $3 is the type of the $1 disk. eg. usb or ata.
# param $4 is the partition name
remove_part () {
  part_number=$(echo $4 | sed 's/.*[a-zA-Z]\(.*\)/\1/')
  if [ "$3" = "usb" ]; then
    if [ "$SLX_REMOVEPART_AUTO" != "yes" ]; then
      remove_part_confirm $1 $2 $part_number $4
    else
      remove_part_delete $1 $2 $part_number
    fi
  else
    remove_part_delete $1 $2 $part_number
  fi
}

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

################## debug info ###################
echo "AUTOMATIC: $SLX_REMOVEPART_AUTO"
echo "REMOVE ALL: $SLX_REMOVEPART_ALL"
if [ "$SLX_REMOVEPART_ALL" != "yes" ] && [ "$SLX_REMOVEPART_AUTO" = "yes" ]; then
  echo "PARTITIONS: $SLX_REMOVEPART_PART"
fi
#################################################

# get list of disk using the format /dev/sdx
DISKS=$(ls /dev/ | grep -e "^sd.*[a-zA-Z]$" | sed -e "s/^/\/dev\//")
if [ -z "$DISKS" ]; then
  perror "Can't find a hard disk."
fi

for disk in $DISKS; do
  # get disk format (used to decide which program will be used)
  if sgdisk $disk | grep -qs 'invalid GPT'; then
    disk_format="msdos"
  else
    disk_format="gpt"
  fi
  # remove "/dev/" from diskname so it can be used in udevadm
  disk_name=$(echo $disk | sed -e "s/\/dev\///")
  # get if disk is ata or usb or something else
  disk_type=$(udevadm info --query=all --name="$disk_name" | grep ID_BUS | cut -d"=" -f2)
  # get disk partition table. if gpt uses sgdisk else uses fdisk
  if [ "$disk_format" != "gpt" ]; then
    disk_partitions=$(fdisk -l $disk 2>/dev/null | grep ^"$disk" | tr -s " " | cut -d" " -f1)
  else
    disk_partitions=$(sgdisk -p "$disk" | grep -e "^[^A-Z]" | tr -s " " | cut -d" " -f2| sed -e "s|^|${disk}|")
  fi

  ################## debug info ###################
  echo "-------------------------------"
  echo "[DISK]" $disk $disk_format $disk_type
  echo "[PART]" $disk_partitions "\n"
  #################################################

  if [ "$SLX_REMOVEPART_ALL" != "yes" ]; then

    if [ "$SLX_REMOVEPART_AUTO" != "yes" ]; then
      # manual mode, with partition selection

      disk_part_count=0  # number of partitions in disk
      disk_part_cl=""    # checklist options string

      # fill checklist string using partitions list
      for disk_part in $disk_partitions; do
        # if disk is not gpt has to check for extended partitions
        if [ "$disk_format" != "gpt" ]; then
          # check wether the partition is extended or not
          disk_part_extn=$(fdisk -l $disk | grep -e "$disk_part.*.Extended")
          # display extended with a (extended) after the partition name
          if [ "$disk_part_extn" ]; then
            disk_part_cl="$disk_part_cl $disk_part $disk_part(extended) off"
          else
            # logical or primary partition. can just display partition name
            disk_part_cl="$disk_part_cl $disk_part $disk_part off"
          fi
        else
          # gpt disk. can just display partition name
          disk_part_cl="$disk_part_cl $disk_part $disk_part off"
        fi
        # +1 disk_partition in checklist
        disk_part_count=$((disk_part_count+1))
      done

      # checks if there are partitions on disk. if there are shows a checklist dialog
      if [ ${disk_part_count} -gt 0 ]; then

        title="Choose partitions to remove from $disk($disk_type):"
        dialog --checklist "$title" 15 60 ${disk_part_count} $disk_part_cl 2>$OUTPUT

        ret=$?
        case $ret in
          0)
            echo "User picked $(cat $OUTPUT)"
            for disk_part in $(cat $OUTPUT); do
              remove_part $disk $disk_format $disk_type $disk_part
            done;;
          1)
            echo "User canceled action";;
          255)
            [ -s $OUTPUT ] &&  cat $OUTPUT || echo "ESC pressed.";;
        esac
      else
        echo "No partitions found on $disk"
      fi

    else
      # automatic mode with config partitions

      # iterate through user defined partitions to be removed
      for disk_part in $SLX_REMOVEPART_PART; do
        # check wether the partition specified by the user is in this disk and remove it
        if [ $(echo "$disk_partitions" | grep "$disk_part") ]; then
          remove_part $disk $disk_format $disk_type $disk_part
        fi
      done
    fi

  else

    if [ "$SLX_REMOVEPART_AUTO" != "yes" ]; then
      # manual mode and remove all partitions
      # will go through disks and ask for permission to remove all partitions

      dialog --title "CLEAR DISK" --yesno "Are you sure you want to wipe all partitions from $disk?" 15 60

      ret=$?
      case $ret in
        0)
          for disk_part in $disk_partitions; do
            remove_part $disk $disk_format $disk_type $disk_part
          done;;
        1)
          echo "User does not want to remove partitions from $disk";;
        255)
          echo "ESC pressed";;
      esac
    else
      # automatic remove all
      # basically manual mode and remove all except it does not ask for permission
      echo "Removing all partitions "$disk_partitions" from $disk"

      pwarning "Removing all partitions from $disk"
      for disk_part in $disk_partitions; do
        remove_part $disk $disk_format $disk_type $disk_part
      done
    fi

  fi
done

rm $OUTPUT

pwarning "Removepart executed successfully"
reboot

