#!/bin/bash
########################################################################################
# /usr/bin/copyiso2usb
########################################################################################
# Usage: copyiso2usb [isofile]
# Installs needed packages if missing - dcfldd cdialog udisks mlocate
# Runs locate command if needed mmlocate.db doesn't exist, or is more than 24 hours
#  old, or requested.
# Uses the locate command to find iso images, presented in a selection dialog, unless
#  a file name is passed as the only argument.
# Lists usb drives in a selection dialog
# Checks that the iso will fit on the selected drive
# Confirms that the command should be run
# Copies the iso image to the usb drive.
#
# Note for maintainers 'xgettext -d copyiso2usb -L Shell /usr/bin/copyiso2usb' used
#  to create copyiso2usb.po (which was manually renamed to copyiso2usb.pot)
########################################################################################
# Author: David W. Hodgins - Mageia Qa team
# 2013-06-02 Initial writing
# 2013-06-04 Added iso as arg, fix for no partition table on usb stick, Thanks barjac.
# 2013-08-22 Converted to use gettext for translations
########################################################################################
set -u # Show usage of unset variables as an error, to help catch typos.
# Obtain translations for all static messages (message with vars translated when needed)
. /usr/bin/gettext.sh
TEXTDOMAIN=copyiso2usb
export TEXTDOMAIN
TEXTDOMAINDIR=/usr/share/locale/
export TEXTDOMAINDIR
Newline="\n"
TitleString=$(gettext "Copy iso image to usb drive")
SelectISO=$(gettext "Select iso image to be copied to a usb drive")
NoISOFound=$(eval_gettext "No iso images found.\${Newline}\${Newline}If an iso has been recently downloaded, please run /etc/cron.daily/mlocate.cron\${Newline}so that this script can find them using the locate command.")
NoUSBDrives=$(eval_gettext "No usb drives found.\${Newline}\${Newline}Please insert a usb drive, and then rerun this script.")
SelectUSBDrive=$(gettext "Select usb drive to copy the iso image to")
CopyCancelled=$(gettext "Copy cancelled. Exiting.")
Waitmlocate=$(gettext "Running /etc/cron.daily/mlocate.cron. Please wait for 'rc=0' before continuing.")
WaitCopy=$(gettext "Copying iso to usb drive. Please wait for 'rc=0' before continuing.")
FileNameTemp=$(mktemp -p /tmp "${0##*/}".XXXXXX)

ThisScript="$0"
if  [[ ! ${ThisScript:0:1} = '/' ]] ; then
    ThisScript="$(pwd)/$0"
fi

# Exit if not root.
if  [[ "$(id -u)" != "0" ]] ; then
    exec pkexec "$ThisScript" "$@"
    MustBeRoot=$(eval_gettext "Must be root to run \$ThisScript")
    ShowMsgExit "$MustBeRoot"
    exit 0
fi

#Check to see if running under pkexec, aka policykit and X
pktest=${PKEXEC_UID-"unset"}
if  [ ! "unset" =  "$pktest" ] ; then
    mydisplay=$(strings "/proc/$PPID/environ"|grep ^DISPLAY=)
    myxauthority=$(strings "/proc/$PPID/environ"|grep ^XAUTHORITY=)
    if  [ -n "$mydisplay" ] ; then
        export DISPLAY=eval "$mydisplay"
        export XAUTHORITY=eval "$myxauthority"
    fi
fi

# Set dialog command to use
#dialogcmd="/usr/bin/Xdialog" #Removed after finding errors from things like fontconfig can cause problems.
#displaytest=${DISPLAY-"unset"}
#if  [ ! -e /usr/bin/Xdialog ] || [ ! -x /usr/bin/Xdialog ] || [ "unset" = "$displaytest" ] ; then
    dialogcmd="/usr/bin/dialog"
#fi

# Function to get window size
GetDisplaySize() {
    IFS=$':, '
    read label MaxLines MaxColumns<<<"$($dialogcmd --print-maxsize --output-fd 1 2>&1)"
    unset IFS
    MenuHeight=$(( MaxLines - 5 ))
}

GetDisplaySize

# Function to list all usb drives
listusbdrives () {
    for block in $(tree -if /sys/devices|grep usb|grep /block/sd\[a-z\]\$); do
        device=${block##*/} # strip last slash and everything before it.
        printf '%s\n' "$device"
    done
}

menuitem=0
MenuText=""
DisplayMenu() {
    GetDisplaySize
    $dialogcmd --backtitle "$TitleString" --title "copyiso2usb" --menu "$DisplayText" "$MaxLines" "$MaxColumns" "$MenuHeight" $MenuText 2>$FileNameTemp
}

DisplayInfo() {
    GetDisplaySize
    $dialogcmd --backtitle "$TitleString" --title "copyiso2usb" --msgbox "$DisplayText" "$MaxLines" "$MaxColumns"
}

RunCommandX() {
    GetDisplaySize
    printf '%s\n\n' "$DisplayText"  >$FileNameTemp
    ( printf '%s\n' "$@" ; eval "$@" ; result=$? ; printf '%s\n' "rc=$result" ) >>$FileNameTemp &
    $dialogcmd --backtitle "$TitleString" \
        --title "copyiso2usb" \
        --tailbox "$FileNameTemp" "$MaxLines" "$MaxColumns"
    rm -f "$FileNameTemp"
}

RunCommand() {
    if  [[ "$dialogcmd" = "/usr/bin/dialog" ]] ; then
        printf '%s\n\n' "$DisplayText"
        eval "$@"
        result=$?
        printf '%s\n' "rc=$result"
    else
        RunCommandX "$@"
    fi
}

ShowMsgExit() {
    DisplayText="$@"
    DisplayInfo
    exit 0
}

# Function to install missing packages
InstallPackageIfFileMissing() {
# Install command, if it isn't already available
    if  [ ! -x "$1" ] ; then
        Command="$1"
        Package="$2"
        CmdNotFoundInstalling=$(eval_gettext "\$Command command not found.  Installing \$Package")
        DisplayText="$CmdNotFoundInstalling"
        RunCommand "urpmi --auto $2"
        result=$?
        if  (($result)) ; then
            FailedExiting=$(eval_gettext "urpmi failed. Return code was \$result. Exiting.")
            ShowMsgExit "$FailedExiting"
            exit 1
        fi
    fi
}

# Function to convert size to human readable format
declare -a SpaceUnits=( B KB MB GB TB PB )
declare -i SpaceUnitsIndex
declare -i SpaceInt
setspaceunits () {
    SpaceUnitsIndex=0
    (( Size *= 100 )) # add two decimal places
    while [ "$Size" -ge 102400 ] ; do
        let Size="$Size"/1024
        (( SpaceUnitsIndex++ ))
    done
    SpaceString="$Size"
    StringLength=${#SpaceString}
    let Intlength="$StringLength - 2"
    StringLeft=${SpaceString:0:$Intlength}
    StringRight=${SpaceString:$Intlength}
    SpaceString="$StringLeft.$StringRight${SpaceUnits[$SpaceUnitsIndex]}"
    if  [ "$Size" -gt 100 ] ; then
        SpaceString="$SpaceString"s
    else
        SpaceString="$SpaceString "
    fi
    printf '%s' "$SpaceString"
}


InstallPackageIfFileMissing /usr/bin/dcfldd dcfldd
InstallPackageIfFileMissing /usr/bin/dialog cdialog
InstallPackageIfFileMissing /usr/bin/udisks udisks
InstallPackageIfFileMissing /etc/cron.daily/mlocate.cron mlocate

if=""
if  [[ $# -eq 1 && -r "$1" ]] ; then
    if="$1"
fi

runmlocate="no"
if  [ -z "$if" ] ; then
    if  [ ! -s /var/lib/mlocate/mlocate.db ] ; then
        runmlocate="yes"
    else
        dbdatetext=$(stat /var/lib/mlocate/mlocate.db|grep Modify)
        dbdate="${dbdatetext#* }" # strip upto and including first space
        dbdatesecondssince1970=$(date -u -d "$dbdate" "+%s")
        currentsecondscince1970=$(date -u "+%s")
        dbageseconds=$(( $currentsecondscince1970 - $dbdatesecondssince1970 ))
        if (("$dbageseconds" > "86400")) ; then # 86400 seconds is one day
            runmlocate="yes"
        else
            mlocateAsk=$(eval_gettext "mlocate.db is \${dbageseconds} seconds old (it's used to locate iso files), would you like it updated?")        
            $dialogcmd --backtitle "$TitleString" \
                --title "copyiso2usb" \
                --yesno "$mlocateAsk" "$MaxLines" "$MaxColumns"
            Selected=$?
            if  (( ! $Selected)) ; then
                runmlocate="yes"
            fi
        fi
    fi
fi


if  [[ "$runmlocate" == "yes" ]] ; then
    DisplayText="${Waitmlocate}"
    RunCommand "/etc/cron.daily/mlocate.cron"
fi

IFS=$'\n'
isolist=($(locate -eiq .iso|grep -i iso$))
unset IFS
if  [ "${#isolist[@]}" == 0 ] ; then
    ShowMsgExit "$NoISOFound"
    exit 2
fi

for i in "${isolist[@]}" ; do
    isolistshow="${i// /_}" # Show spaces in file names as underscores
    MenuText="$MenuText $menuitem ${isolistshow}"
    (( menuitem++ ))
done

if  [[ $# -eq 1 && -r "$1" ]] ; then
    if="$1"
else
    DisplayText="$SelectISO"
    DisplayMenu
    Selected=$(<$FileNameTemp)
    rm -f "$FileNameTemp"
    if  [[ "${Selected}" == "" ]] ; then
        exit 3 # Cancel selected 
    fi
    if="${isolist[$Selected]}"
fi

shopt -s nullglob
usbdrives=($(listusbdrives))

if  [ "${#usbdrives[@]}" == 0 ] ; then
    ShowMsgExit "$NoUSBDrives"
    exit 4
fi

menuitem=0
MenuText=""

for dev in "${usbdrives[@]}" ; do
    IFS=$'\n'
    udisksinfo=($(udisks --show-info /dev/$dev |grep -e size -e label -e vendor -e model -e serial))
    unset IFS
    Size=${udisksinfo[0]##*: }
    read  -rd '' Size <<< "$Size"
    usbsize[$menuitem]="$Size"
    SizeString=$(setspaceunits)
    Label=${udisksinfo[2]##*: }
    Vendor=${udisksinfo[3]##*: }
    Model=${udisksinfo[4]##*: }
    Serial=${udisksinfo[5]##*: } 
    if  [[ "${Label}" == "" ]] ; then
        udisksinfo=($(udisks --show-info /dev/${dev}1 |grep -e label))
        Label=${udisksinfo[0]##*:}
    fi
    Mountpoint=""
    read mntdev Mountpoint rest<<<"$(grep ^/dev/$dev /proc/mounts)"
    read  -rd '' SizeString <<< "$SizeString" # Remove leading/trailing spaces
    SizeString="${SizeString// /_}" # Translate any embedded spaces to underscores
    Label=""
    read  -rd '' Label <<< "$Label"
    if  [[ "${Label}" == "" ]] ; then
        if  [ -e "/dev/${dev}1" ] ; then 
            IFS=$'\n'
            udisksinfo=($(udisks --show-info /dev/${dev}1 |grep -m 1 label))
            unset IFS
            Label=${udisksinfo[0]##*:}
        fi
    fi
    read  -rd '' Label <<< "$Label"
    Label="${Label// /_}"
    read  -rd '' Vendor <<< "$Vendor"
    Vendor="${Vendor// /_}"
    read  -rd '' Model <<< "$Model"
    Model="${Model// /_}"
    read  -rd '' Serial <<< "$Serial"
    Serial="${Serial// /_}"
    read  -rd '' Mountpoint <<< "$Mountpoint"
    Mountpoint="${Mountpoint// /_}"
    if  [[ ! "${Mountpoint}" == "" ]] ; then
         Mountpoint="---MOUNTED_at_${Mountpoint}"
    fi
    MenuText="$MenuText $menuitem ${dev}.${SizeString}.${Label}.${Vendor}.${Model}.${Serial}${Mountpoint}"
    (( menuitem++ ))
done
DisplayText="$SelectUSBDrive"
DisplayMenu
Selected=$(<$FileNameTemp) 
rm -f "$FileNameTemp"

if  [[ "${Selected}" == "" ]] ; then
    exit 5 # Cancel selected 
fi

of="${usbdrives[$Selected]}"
ofsize="${usbsize[$Selected]}"
ddcmd="dcfldd sizeprobe=if statusinterval=5 if='""$if""' of=/dev/$of bs=1M"

read label IsoSize rest<<<"$(stat $if|grep Size)"
if  (( $IsoSize > $ofsize )) ; then
    TooBig=$(eval_gettext "ISO size is \$IsoSize bytes. Selected USB stick only has room for \$ofsize bytes. Copy cancelled. Exiting.")
    ShowMsgExit "$TooBig"
    exit 7
fi

ConfirmRun=$(eval_gettext "About to run\${Newline}\${ddcmd}\${Newline}Please select yes to confirm, or no to cancel.\${Newline}NOTE: When it runs, please wait for 'rc=0' before continuing.")
$dialogcmd --backtitle "$TitleString" \
    --title "copyiso2usb" \
    --yesno "$ConfirmRun" "$MaxLines" "$MaxColumns"
Selected=$?

if  (($Selected)) ; then
    ShowMsgExit "$CopyCancelled"
    exit 7
fi

DisplayText="$WaitCopy"
RunCommand "${ddcmd} 2>&1"
exit 0