#!/usr/bin/bash

############################################
#          Claire Robinson (MrsB)          #
#             for Mageia QA                #
############################################
#            Thanks also go to..           #
#        Alfred Kretschmer (alfred)        #
#        Barry Jackson (barjac)            #
############################################

############################################
# If dcfldd is installed you will get a    #
# progress indicator when dumping to USB.  #
############################################

XDG_CONFIG_HOME="${XDG_CONFIG_HOME:-$HOME/.config}"
conf_file="$XDG_CONFIG_HOME/Mageia/dorsync.conf"

generate_config_file () {
  mkdir -p "$XDG_CONFIG_HOME/Mageia"
  cat << EOF > "$conf_file"
#############################################
# Configuration file for the dorsync script #
#############################################
location="" # Full path to the local folder where ISOs will be downloaded.
release=""  # Name of the directory on the rsync server.
user=""     # Username used by rsync.
password="" # Password used by rsync.
EOF
}

if [ -e "$conf_file" ]; then
  source "$conf_file"
else
  generate_config_file
  echo "Missing file $conf_file has been created."
  echo "Please edit this file and then rerun dorsync."
  exit
fi

red='\e[0;31m'
green='\e[0;32m'
bold='\e[01m'
endcolour='\e[0m'

# Check required packages are installed
deps=(coreutils udisks2 sed grep rsync)
for dep in "${deps[@]}"; do
  if ! rpm --quiet -q "$dep"; then
    echo "You need to install $dep to run this program"
    baddep=1
  fi
done
[[ ${#baddep} -gt 0 ]] && exit 1

if [ ! -d "$location" ]; then
  echo "Directory $location does not exist."
  echo "Please create this directory or fix your 'location' parameter in dorsync.conf."
  exit
fi

cd "$location"
rsync_url="rsync://$user@bcd.mageia.org/isos/$release/"

select_isos () {
  isos=`RSYNC_PASSWORD="$password" rsync --list-only --no-recursive "$rsync_url" | tr -s ' ' | cut -f5 -d" " | grep -vE "^\.|torrent"`
  rm -f dorsync.skip
  echo -e "\nFor each of the following ISOs, press Y to download/update it, or any other key to skip it:\n"
  for iso in $isos; do
    read -r -n1 -e -p "  $iso : " yn
    if [[ $yn = [Yy] ]]; then
      echo "+$iso" >> dorsync.skip
    else
      echo "$iso" >> dorsync.skip
    fi
  done
  echo -e "\nYour selection has been saved in dorsync.skip.\n"
}

# Function to do the actual rsync and run checks afterwards
dorsync () {
if [ -e dorsync.skip ]; then
  echo -e "\ndorsync.skip file found:\n"
  cat dorsync.skip | while read -r iso
  do
    if [[ `echo "$iso" | cut -c 1` = '+' ]]; then
      iso=`echo "$iso" | cut -c 2-`
      echo -e "  $iso $green SELECTED $endcolour"
    else
      echo -e "  $iso $red SKIPPED $endcolour"
    fi
  done
  echo -e "\nDo you want to continue with this selection?"
  read -r -n1 -e -p "Press Y to continue, or any other key to modify this list: " yn
  if [[ $yn != [Yy] ]]; then
    select_isos
  fi
else
  echo -e "\nNo dorsync.skip file found. Do you really want to download/update all ISOs?"
  read -r -n1 -e -p "Press Y to continue with all ISOs, or any other key to choose which ISOs to select: " yn
  if [[ $yn != [Yy] ]]; then
    select_isos
  fi
fi

rsyncargs=("-avHP")
# Exclude torrent files. We only want ISOs.
rsyncargs+=("--exclude=*.torrent")

if [ -e dorsync.skip ]; then
  rsyncargs+=("--exclude-from=dorsync.skip")
fi

echo "Starting rsync.."
RSYNC_PASSWORD="$password" rsync "${rsyncargs[@]}" "$rsync_url" .
OUT=$?
echo ""
if [ $OUT -eq 0 ];then
   echo "Completed rsync"
else
   echo "Rsync failed. Exiting."
   exit
fi

echo ""
du -h | sort -k 2 | sed "s|\./||" | sed -e "s|\(\s\)\.|\1Total\n|"
echo ""
echo "Starting checks, please ensure the dates are correct"
echo ""
# Loop through folders in $location and read info into arrays
# Hopefully able to use this later to dump to USB stick
# Perform checks too and highlight green OK red Failed
i=0
OIFS="$IFS"
IFS=$'\n'

# Sorted list with classic isos before live isos
for line in $(find . -type f \( ! -name "*Live*" \) -and \( -name "Mageia*.iso" \) -printf '%P\n' | sort && find . -type f -name "Mageia*Live*.iso" -printf '%P\n' | sort)
do
  directory[i]=`dirname "$line"`
  iso[i]=`basename "$line"`
  echo "PATH: ${directory[i]}"
  echo "ISO:  ${iso[i]}"
  cd "${directory[i]}"
  # Check for DATE.txt and display it or warn if missing
  if [ -e DATE.txt ]
  then
    echo "DATE: `cat DATE.txt`"
    missingdate[i]=false
  else
    echo -e "DATE: ${red}Missing${endcolour}"
    missingdate[i]=true
  fi
  # Check for md5 file and then md5sum it or warn if missing
  if [ -e "${iso[i]}.md5" ]
  then
    md5sum -c --quiet --status ./*.md5 2>/dev/null
    OUT=$?
    if [ $OUT -eq 0 ];then
      echo -e "MD5:  ${green}OK${endcolour}     \c"
      md5ok[i]=true
      missingmd5[i]=false
    else
      echo -e "MD5:  ${red}Failed${endcolour}     \c"
      md5ok[i]=false
      missingmd5[i]=false
    fi
  else
    echo -e "MD5:  ${red}Missing${endcolour}     \c"
    missingmd5[i]=true
  fi

  # Check for sha3 file and then sha3-512sum it or warn if missing
  if [[ -e "${iso[i]}.sha3" && -x "/usr/bin/sha3-512sum" ]]
  then
    sha3-512sum -c ./*.sha3 1>/dev/null 2>&1
    OUT=$?
    if [ $OUT -eq 0 ];then
      echo -e "SHA3: ${green}OK${endcolour}     \c"
      sha3ok[i]=true
      missingsha3[i]=false
    else
      echo -e "SHA3: ${red}Failed${endcolour}     \c"
      sha3ok[i]=false
      missingsha3[i]=false
    fi
  elif [ -e "${iso[i]}.sha3" ]
  then
    echo -e "SHA3: ${red}sha3-512sum missing${endcolour}     \c"
      sha3ok[i]=false
      missingsha3[i]=false
  else
    echo -e "SHA3: ${red}Missing${endcolour}     \c"
    sha3ok[i]=false
    missingsha3[i]=true
  fi

  # Check for sha512 file and then sha512sum it or warn if missing
  if [ -e "${iso[i]}.sha512" ]
  then
    sha512sum -c --quiet --status ./*.sha512 2>/dev/null
    OUT=$?
    if [ $OUT -eq 0 ];then
      echo -e "SHA512: ${green}OK${endcolour}"
      sha512ok[i]=true
      missingsha512[i]=false
    else
      echo -e "SHA512: ${red}Failed${endcolour}"
      sha512ok[i]=false
      missingsha512[i]=false
    fi
  else
    echo -e "SHA512: ${red}Missing${endcolour}"
    sha512ok[i]=false
    missingsha512[i]=true
  fi

  echo ""
  cd "$location"
  (( i++ ))
done
IFS="$OIFS"
# Check if anything failed and provide a summary
total=${#directory[@]}
warnings=0
startwarning="--------------------------------------------------\n${bold}Warning:${endcolour}\n"
for (( i=0; i<=$(( total -1 )); i++ ))
do
  # Warn of missing DATE.txt files
  if [ "${missingdate[i]}" = true ]
  then
    if [ $warnings = 0 ]
    then
      echo -e "$startwarning"
    fi
    echo "${directory[i]}/DATE.txt missing"
    (( warnings++ ))
  fi
  # Warn of missing md5 files
  if [ "${missingmd5[i]}" = true ]
  then
    if [ $warnings = 0 ]
    then
      echo -e "$startwarning"
    fi
    echo "${directory[i]}/${iso[i]}.md5 missing"
    (( warnings++ ))
  else
    # Warn of bad md5sum if present
    if [ "${md5ok[i]}" = false ]
    then
      if [ $warnings = 0 ]
      then
        echo -e "$startwarning"
      fi
      echo "${iso[i]} md5sum failed"
      (( warnings++ ))
    fi
  fi
  # Warn of missing sha3 files
  if [ "${missingsha3[i]}" = true ]
  then
    if [ $warnings = 0 ]
    then
      echo -e "$startwarning"
    fi
    echo "${directory[i]}/${iso[i]}.sha3 missing"
    (( warnings++ ))
  else
    # Warn of bad sha3sum if present
    if [ "${sha3ok[i]}" = false ]
    then
      if [ $warnings = 0 ]
      then
        echo -e "$startwarning"
      fi
      echo "${iso[i]} sha3-512sum failed"
      (( warnings++ ))
    fi
  fi
  # Warn of missing sha512 files
  if [ "${missingsha512[i]}" = true ]
  then
    if [ $warnings = 0 ]
    then
      echo -e "$startwarning"
    fi
    echo "${directory[i]}/${iso[i]}.sha512 missing"
    (( warnings++ ))
  else
    # Warn of bad sha512sum if present
    if [ "${sha512ok[i]}" = false ]
    then
      if [ $warnings = 0 ]
      then
        echo -e "$startwarning"
      fi
      echo "${iso[i]} sha512sum failed"
      (( warnings++ ))
    fi
  fi
done
echo ""
# If anything failed, provide a count
if [ $warnings != 0 ]
then
  if [ $warnings -gt 1 ]
  then
    message="There are $warnings potential problems."
  else
    message="There is 1 potential problem."
  fi
  echo -e "${bold}$message${endcolour}"
  echo "--------------------------------------------------"
else
  echo "------------------------------------"
  echo ""
  echo -e "${bold}All tests passed OK. Yippee!${endcolour}"
  echo "Please check the dates are correct."
  echo ""
  echo "------------------------------------"
fi

}


# Function to list USB sticks
liststicks () {
  i=0
  for usb in $(lsblk -ndpro NAME,TRAN,TYPE | grep "usb disk$" | cut -f1 -d" ")
  do
    showinfo=`udisksctl info -b "$usb"`
    usb_size=`echo "$showinfo" | grep "  Size:" | cut -f2 -d: | sed 's/^[ ]*//'`
    # Ignore empty drives.
    if [ $usb_size = 0 ]; then
      continue
    fi
    usbdev[i]="$usb"
    found[i]=`date`
    sizebytes[i]=$usb_size
    sizegb[i]=`echo "scale=1;${sizebytes[i]}/1024/1024/1024" | bc`

    showinfo=`udevadm info -q property -n "$usb"`
    vendor[i]=`echo "$showinfo" | grep "ID_VENDOR=" | cut -f2 -d=`
    model[i]=`echo "$showinfo" | grep "ID_MODEL=" | cut -f2 -d=`
    (( i++ ))
  done
}


# Minimal function to find isos if rsync is skipped
findisos () {
  i=0
  OIFS="$IFS"
  IFS=$'\n'
  # Sorted list with classic isos before live isos
  for line in $(find . -type f \( ! -name "*Live*" \) -and \( -name "Mageia*.iso" \) -printf '%P\n' | sort && find . -type f -name "Mageia*Live*.iso" -printf '%P\n' | sort)
  do
    directory[i]=`dirname "$line"`
    iso[i]=`basename "$line"`
    (( i++ ))
  done
  IFS="$OIFS"
}


# Function to wait for usb stick to be inserted
dodoh () {
  while true
  do
    echo "No USB drives found. Please insert one now."
    read -r -n1 -e -p "Press Q to quit or any other key to continue: " num
    case $num in
      [Qq]* ) exit;;
      * ) liststicks
	  echo ""
	  if ! [ ${#found[@]} = 0 ]
	  then
	    doh="Yay"
	    break
	  fi
	  ;;
    esac
  done
}


# Function to actually do the dump onto usb
dothedeed () {
  echo ""
  echo "Please be patient this can take a while"
  echo ""
  echo "This operation requires root privileges"
  echo "Please enter your root password.."
  echo ""
  export rootwhere="${usbdev[$((usbchoice-1))]}"
  export rootiso="$fullisopath"
  if [ -e /usr/bin/dcfldd ]
  then
    su -c 'dcfldd sizeprobe=if statusinterval=5 if="$rootiso" of="$rootwhere" bs=1M && sync'
  else
    su -c 'dd if="$rootiso" of="$rootwhere" bs=1M && sync'
  fi
  OUT=$?
  echo ""
  if [ $OUT -eq 0 ];then
    echo "Completed OK."
  else
    echo "Failed. Exiting."
    exit
  fi
}

# Function to choose usb stick & iso and check it'll fit
dodump () {
  echo ""
  didrsync=$1
  if [ "$didrsync" == 1 ]
  then
    # Ask to dump to USB or exit
    echo ""
    while true
    do
      read -r -n1 -e -p "Do you want to dump one onto USB? [y/n]" yn
      case $yn in
	[Yy] ) liststicks
		break;;
	[Nn] ) exit;;
	   * ) echo "Please choose y or n";;
      esac
    done
    echo ""
  else
    liststicks
  fi

  # If no USBs found D'oh! loop until there is one..
  total=${#found[@]}
  if [ "$total" = 0 ]
  then
    doh="doh"
    while [ "$doh" = "doh" ]
    do
      dodoh
    done
    total=${#found[@]}
  fi

  # Now there is a USB. Print the info
  echo ""
  echo "USB drives found:"
  echo ""
  echo -e "    Device\tSize\tWhen Found\t\t\tModel"

  for (( i=0; i<=$(( total -1 )); i++ ))
  do
    echo -e "$((i+1))) ${usbdev[i]}\t${sizegb[i]}Gb\t${found[i]}\t${vendor[i]} ${model[i]}"
  done

  # Choose which USB to use
  echo ""
  while true
  do
  read -r -e -p "Please choose which USB to use or Q to quit and press <enter>: " usbchoice
   if [[ $usbchoice = [Qq] ]]; then
   exit
   elif [[ "$usbchoice" =~ ^[0-9]+$ ]] && (( $usbchoice <= $total )) && (( $usbchoice > 0 )); then
   break
   else echo "Invalid selection"
   fi
  done

  # List the ISOs
  echo ""
  echo "List of ISOs:"

  total=${#iso[@]}
  for (( i=0; i<=$(( total -1 )); i++ ))
  do
    echo "$((i+1))) ${iso[i]}"
  done
  echo ""

  # Choose which ISO to use
  while true
  do
   read -r -e -p "Please choose which ISO to use or Q to quit and press <enter>: " isochoice
   if [[ $isochoice = [Qq] ]]; then
   exit
   elif [[ "$isochoice" =~ ^[0-9]+$ ]] && (( $isochoice <= $total )) && (( $isochoice > 0 )); then
   break
   else echo "Invalid selection"
   fi
  done

  # Will it fit?
  fullisopath=`readlink -f "${directory[$((isochoice-1))]}/${iso[$((isochoice-1))]}"`
  isosize=`du -b "$fullisopath" | cut -f1`
  usbsize=${sizebytes[$((usbchoice-1))]}
  if [ "$usbsize" -lt "$isosize" ]
  then
    echo ""
    echo "Sorry. It won't fit! Exiting."
    exit
  fi

  # Warn
  usbinfo="${usbdev[$((usbchoice-1))]} (${sizegb[$((usbchoice-1))]}Gb ${vendor[$((usbchoice-1))]} ${model[$((usbchoice-1))]} found at ${found[$((usbchoice-1))]})"
  echo ""
  echo -e "${bold}About to dump $fullisopath${endcolour}"
  echo -e "${bold}onto $usbinfo${endcolour}"
  echo ""
  echo "This will destroy any data already on the USB device."

  # Confirm
  while true
  do
    read -r -n1 -e -p "Press Y to confirm or Q to quit: " yq
    case $yq in
      [Qq] ) exit;;
      [Yy] ) dothedeed
	     break;;
	 * ) echo -e "Invalid selection\n";;
    esac
  done
}


# Offer to skip rsync and just dump to usb
echo ""
while true
do
  read -r -n1 -e -p "Do you want to Rsync (R), skip Rsync and just dump to usb (D) or quit (Q): " rdq
case $rdq in
    [Rr] )  dorsync; dodump 1; exit;;
    [Dd] )  findisos; dodump 0; exit;;
    [Qq] )  exit;;
       * )  echo " $rdq Invalid selection"; echo "";;
  esac
done

