#!/bin/bash

# DD Options
BLOCKSIZE="4" # in megabytes

# Paths
MOUNTPATH="/mnt/mediadrive"
DEVICEDIR="/dev"
HOMEDIR="."

# Colors
RED='\033[0;31m'
LRED="\033[1;31m"
BLUE="\033[0;34m"
LBLUE="\033[1;34m"
GREEN="\033[0;32m"
LGREEN="\033[1;32m"
YELLOW="\033[1;33m"
CYAN="\033[0;36m"
LCYAN="\033[1;36m"
PURPLE="\033[0;35m"
LPURPLE="\033[1;35m"
BWHITE="\e[1m"
NC='\033[0m' # No Color

if [[ $EUID -ne 0 ]]; then
  printf "${LRED}This script must be run as root${NC}\n" 1>&2
   exit 1
fi

Options(){

  printf "${BWHITE}[Q]${NC} - Quit\n"
}
FindImages(){
  printf "${YELLOW}Searching for image files..${NC}\n"
  echo
  printf "${LCYAN}Results Found:${NC}\n"
  echo
  if [ "${1}" ];then
    find $1 -regex ".*\.\(img\|iso\)" | grep [.iso,.img]
  else
    find ${HOMEDIR} -regex ".*\.\(img\|iso\)" | grep [.iso,.img]
  fi
}
CheckFile(){
  if [ ! -f "${1}" ];then
    printf "${LRED}${1} is not a valid file path${NC}\n"
    exit 0;
  fi
}

CheckDev(){
  if [ ! -b "${1}" ];then
    printf "${LRED}${1} is not a valid block device${NC}\n"
    exit 0;
  fi
}

ConfirmWrite(){
  echo
  printf "${YELLOW}Are you sure you want to write changes to${NC} ${LRED}${2}${NC}${YELLOW}?${NC}\n"
  printf "${YELLOW}This will overwrite existing data${NC}\n"
  printf "${YELLOW}[y/n]: ${NC}"
  read choice

  if [ "${choice}" = "n" ];then
    echo
    printf "${LRED}Aborted!${NC}\n"
    exit 0;
  else
    echo
    printf "${YELLOW}Writing data:${NC}\n"
    printf "${LGREEN}${1}${NC} --> ${LPURPLE}${2}${NC}\n"
    echo
  fi
}
Quit(){
  if [ "${1}" = "q" ];then
    echo
    printf "${LRED}Exiting......${NC}\n"
    echo
    exit 0;
 elif [ "${1}" = "Q" ];then
   echo
   printf "${LRED}Exiting......${NC}\n"
   echo
   exit 0;
  fi
}
FindMgr(){

    use_which=$(which {dnf,pacman,xbps-install})
    base_dir="/usr/bin"
    printf "${BWHITE}\n\nGetting Dependencies...${NC}\n\n" 
    if [ -f "${base_dir}/apt" ] || [ -f "${base_dir}/apt-get" ];then
        apt-get update; apt-get -y install pv

    elif [ "${use_which}" = "${base_dir}/dnf" ];then
        dnf -y install pv

    elif [ "${use_which}" = "${base_dir}/pacman" ];then
        pacman -Sy pv

    elif [ "${use_which}" = "${base_dir}/xbps-install" ];then
        xbps-install -Sy pv

    else
        printf "\n${LRED}No suitible package manager was found!${NC}\n\n"
    fi
    return
}
case "$1" in

-d) if [ -f /usr/bin/pv ];then
     clear
     printf "${LCYAN}idd - interactive dd${NC}\n"
     printf "${LPURPLE}********************************${NC}\n"
     Options
     echo
     printf "${LGREEN}Image files (img, iso) in ${2}${NC}\n"
     printf "${LBLUE}-------------------------------------------${NC}\n"

     #find $2 -regex ".*\.\(img\|iso\)" | grep [.iso,.img]
     FindImages $2
     if [ ! $? = 0 ];then
       printf "${LRED}No image files found${NC}\n"
     fi
     printf "${LBLUE}-------------------------------------------${NC}\n"

     printf "${YELLOW}Copy & paste full file path here: ${NC}"
     read filepath
     Quit ${filepath}
     CheckFile ${filepath}
     clear
     Options
     echo
     printf "${LGREEN}List of DEVICES in ${DEVICEDIR}/${NC}\n"
     printf "${LBLUE}-------------------------------------------${NC}\n"
     lsblk
     printf "${LBLUE}-------------------------------------------${NC}\n"
     echo
     printf "${YELLOW}Enter name of device${NC} (${LCYAN}example${NC}, ${LBLUE}sdb${NC}): "
     read devname
     Quit ${devname}
     DEVPATH="${DEVICEDIR}/${devname}"
     CheckDev ${DEVPATH}
     ConfirmWrite ${filepath} ${DEVPATH}


     echo
     printf "${LCYAN}Please wait ...${NC}\n"
     echo
     pv -tpreb ${filepath} | dd bs=${BLOCKSIZE}M of=${DEVICEDIR}/${devname}
     echo
     printf "${YELLOW}Syncing buffers...${NC}\n";sync
     echo
     printf "${LGREEN}Done writing file${NC}\n"
  else
     echo "${0}: Please install 'pv' with your package manager"
  fi
  ;;
-h) echo """
    usage: idd -d [directory]
    """
    ;;

-t) FindMgr
    ;;
*) if [ -f /usr/bin/pv ];then
      clear
      printf "${LCYAN}idd - interactive dd${NC}\n"
      printf "${LPURPLE}********************************${NC}\n"
      Options
      echo
      printf "${LGREEN}Image files (img, iso) in ${HOMEDIR}${NC}\n"
      printf "${LBLUE}-------------------------------------------${NC}\n"

      FindImages
      if [ ! $? = 0 ];then
        printf "${LRED}No image files found${NC}\n"
      fi

      printf "${LBLUE}-------------------------------------------${NC}\n"

      printf "${YELLOW}Copy & paste full file path here: ${NC}"
      read filepath
      Quit ${filepath}
      CheckFile ${filepath}


      clear
      Options
      echo
      printf "${LGREEN}List of DEVICES in ${DEVICEDIR}/${NC}\n"
      printf "${LBLUE}-------------------------------------------${NC}\n"
      lsblk
      printf "${LBLUE}-------------------------------------------${NC}\n"
      echo
      printf "${YELLOW}Enter name of device${NC} (${LCYAN}example${NC}, ${LBLUE}sdb${NC}): "
      read devname
      Quit ${devname}
      DEVPATH="${DEVICEDIR}/${devname}"
      CheckDev ${DEVPATH}
      ConfirmWrite ${filepath} ${DEVPATH}

      echo
      printf "${LCYAN}Please wait ...${NC}\n"
      echo
      pv -tpreb ${filepath} | dd bs=${BLOCKSIZE}M of=${DEVICEDIR}/${devname}
      echo
      printf "${YELLOW}Syncing buffers...${NC}\n";sync
      echo
      printf "${LGREEN}Done writing file${NC}\n"
    else
        FindMgr && ${0}
    fi
    ;;
esac
