#!/bin/bash
    
##Original Code by dolphin_oracle and BitJam November 15 2016 (dolphinoracle@gmail.com)
##License is GPL 3.0
##originally developed for antiX and mx linux.

##depends on live-usb-maker, cli-shell-utils and yad

#add BitJam's soure cli-shell-utils
source /usr/local/lib/cli-shell-utils/cli-shell-utils.bash

#translations/locale stuff
TEXTDOMAINDIR=/usr/share/locale 
export TEXTDOMAIN="live-usb-maker-gui"

WORK_DIR=/run/$ME
FIFO=$WORK_DIR/fifo

CLI_PROG="live-usb-maker"

LOCK_FILE="/run/lock/$CLI_PROG"
ERR_FILE="/var/log/$CLI_PROG.error"
PROG_FILE="/var/log/$CLI_PROG.progress"
LOG_FILE="/var/log/$CLI_PROG.log"

USER_STOP=false

TITLE=$"Live-USB Maker GUI"
CHOOSE_MODE_TITLE=$"Live System Detected.  Select Live-USB creation mode"
CHOOSE_MODE=$"Select Mode"
CHOOSE_USB=$"Select USB Device"
CHOOSE_CLONE=$"Clone Mode"
CHOOSE_ISO=$"Select ISO"
CLONE_RUNNING_SYSTEM=$"Clone Running Live System"
INFO=$"Information"
BAD_SOURCE=$"Source is not an ISO file.  Please choose another file."
DESCRIPTION=$"Simple GUI for Live-USB Maker"
VERSION=$"Version"
USB_CHECK=$"No target USB device detected.  Please insert a device and try again"
LICENSE=$"License"
APPLY=$"Apply"
CLOSE=$"Close"
ABOUT=$"About"
HELP=$"Help"
ADVANCE=$"Advanced"
COMPLETE=$"LiveUSB creation successful!"
STOP_MESSAGE=$"Operation stopped by user."
MSG_LUM_LOCK_EXISTS=$"A %s process (using PID %s) is already running"
MSG_LUM_LOCK_EXISTS_2=$"Please close that process before starting a new one"
MSG_OTHER_LOCK_ERROR=$"Failed to obtain lock on %s"
PROGRESS_CLOSED=$"Progress window closed, but job not done. Do you want to keep going?"
YES=$"Yes"
NO=$"No"
ENCRYPT=$"Encrypt Live-USB"
dd_mode=$"dd mode"
copy_mode=$"Update mode"
full_featured=$"Normal mode - full featured live USB"
dd_title="making liveUSB in dd mode"


PRETEND=$1
echo "Pretend Mode = " $PRETEND
if [ "$PRETEND" = "-p" ]; then
    echo "Pretend Mode Enabled"
else
    unset PRETEND
fi 

dd_default=$1
if [ "$dd_default" = "--dd" ]; then
    dd_switch="^"
    echo "dd default enabled " $dd_switch
else
    unset dd_default
fi

##If on mx-viewer is present, use mx-viewer, else wise use generic xdg-open 
if [ -e "/usr/bin/mx-viewer" ]; then
    BROWSER="mx-viewer"
else
    BROWSER="xdg-open"
fi


##-----------------------------------------------------

main() {
    echo "creating lock ..."
    trap clean_up EXIT
    gui_flock $LOCK_FILE || flock_error
    mkdir -p $WORK_DIR

    mkfifo $FIFO
    exec 10<> $FIFO
    truncate -s0 $PROG_FILE


    ##MODE 1 is clone, no encyrption available
    ##MODE 2 is standard with no encryption available
    ##MODE 3 is clone with encryption available
    ##MODE 4 is standard iso selection with encryption available
    ##We start with mode 2, standard no encyrption, for backward compatibility
    MODE=2
    status=0
    ##start gui and run actions
    ##user that launched the app
    launch_user=$(ps -aux |grep -v grep| grep -m 1 live-usb-maker-gui |awk '{print $1}')
    echo "launch user is " $launch_user

    #offer clone mode if running live and dd_switch is blank
    if [ "$dd_defaLC_ALL=en_US.utf8ult" = "--dd" ]; then
        echo "dd mode selected, don't check clone mode"
    else
        check_clone_mode
    fi
    
    ##check that encrypt option is available
    live-usb-maker --help | grep -q -- --encrypt 
    if [ $? = 0 ]; then 
        add_encrypt_checkbox
    fi

    ##check that gui progress option is available
    live-usb-maker --help | grep -q -- --gui-progress
    if [ $? = 0 ]; then
        PROGRESS_OPTION="--gui-progress yad --center --window-icon=drive-removable-media-usb --no-escape --title=\"$TITLE\" --width=300 --on-top --progress --no-buttons --auto-close --undecorated "
    else
        PROGRESS_OPTION=""
    fi
    echo Progress Option is $PROGRESS_OPTION 

    while [ "$status" = "0" ]; do
        display_selections
        actions
        status=$actionstatus
    
    done
}

#define functions

check_clone_mode() {
    LIVE_CHECK=$(df -T / |tail -n1 |awk '{print $2}')
    echo "live check = " $LIVE_CHECK
    if [ "$LIVE_CHECK" = "overlay" ] || [ "$LIVE_CHECK" = "aufs" ]; then

        MODE_SELECTION=$(yad --window-icon=drive-removable-media-usb --form --num-output --button="$APPLY"!emblem-default!:0 --button="$CLOSE"!gtk-close!:1 --title="$TITLE" \
            --field="<b>$CHOOSE_MODE_TITLE</b>":LBL ""\
            --field="":CB "$CLONE_RUNNING_SYSTEM!$CHOOSE_ISO"\
            --field="":LBL "")

        case $? in
            0)     MODE=$(echo $MODE_SELECTION|cut -d "|" -f2)
                   echo "Mode is " $MODE ;;

            1)     exit 0    ;;

	  252)     exit 0    ;;
        esac
    fi
}


display_selections() {

    build_usb_list

    # get user selections with above for defaults
    #display selection dialog

    case $MODE in

        1) selections=$(yad --window-icon=drive-removable-media-usb --form --center --button="$ABOUT"!help-about!:8 --button="$HELP"!help-contents!:4 --button="$APPLY"!emblem-default!:0 --button="$CLOSE"!gtk-close!:1 --width=300 --title="$TITLE" \
            --field="<b>Clone Mode</b>":LBL ""\
            --field="$CHOOSE_USB":CB "$usb_list"\
            --field="$CHOOSE_CLONE":CB "clone"\
            --field="":LBL "") ;;

        2) selections=$(yad --window-icon=drive-removable-media-usb --form --center --button="$ABOUT"!help-about!:8 --button="$HELP"!help-contents!:4 --button="$APPLY"!emblem-default!:0 --button="$CLOSE"!gtk-close!:1 --width=300 --title="$TITLE" \
            --field="<b>$TITLE</b>":LBL ""\
            --field="$CHOOSE_USB":CB "$usb_list"\
            --field="$CHOOSE_ISO":FL "$source" --file-filter="iso|*.iso" \
            --field=$"Options":CB "1: $full_featured!$dd_switch 2: $dd_mode!"\
            --field="":LBL "") ;;
            
        3) selections=$(yad --window-icon=drive-removable-media-usb --form --center --button="$ABOUT"!help-about!:8 --button="$HELP"!help-contents!:4 --button="$APPLY"!emblem-default!:0 --button="$CLOSE"!gtk-close!:1 --width=300 --title="$TITLE" \
            --field="<b>Clone Mode</b>":LBL ""\
            --field="$CHOOSE_USB":CB "$usb_list"\
            --field="$CHOOSE_CLONE":CB "clone"\
            --field=$"Options":CB "1: $full_featured!4: $ENCRYPT"\
            --field="":LBL "") ;;

        4) selections=$(yad --window-icon=drive-removable-media-usb --form --center --button="$ABOUT"!help-about!:8 --button="$HELP"!help-contents!:4 --button="$APPLY"!emblem-default!:0 --button="$CLOSE"!gtk-close!:1 --width=300 --title="$TITLE" \
            --field="<b>$TITLE</b>":LBL ""\
            --field="$CHOOSE_USB":CB "$usb_list"\
            --field="$CHOOSE_ISO":FL "$source" --file-filter="iso|*.iso" \
            --field=$"Options":CB "1: $full_featured!$dd_switch 2: $dd_mode!3: $copy_mode!4: $ENCRYPT"\
            --field="":LBL "") ;;
    esac
}


actions() {
    status2=$?

    case $status2 in

        0) adjust_settings ;;
        1) quit            ;;
        8) display_about   ;;
        4) display_help    ;;
        5) display_license ;;
	252)     quit    ;;
    esac
}

adjust_settings() {

    echo $selections

    device_and_source

    echo "usb_device is " $usb_device
    echo "source is " $source
    echo "OPTIONS is " $OPTION
    
    check_unmountable
    
    #check to make sure source is a iso

    if [[ $MODE = 2 || $MODE = 4 ]]; then
        if [ "$(file "$source" |grep -i -q iso && echo $?)" = "0" ]; then
            echo "is a iso file"
            is_a_iso=true
        else
            echo "not an iso file"
            is_a_iso=false
            yad --form --title=$TITLE --button=gtk-ok --text="$BAD_SOURCE" --buttons-layout=center
        fi
    fi

    if [ "$OPTION" = "dd" ]; then
	echo "OPTION IS DD"
        if [ $is_a_iso = true ]; then
        check_size
        make_dd_usb
        error_notification
        else
        actionstatus=0
        fi
    else
        case $MODE in
    
            1) source="clone"
               make_live_usb
               error_notification
                ;;
                
            3) source="clone"
               make_live_usb
               error_notification
                ;;
                
            2)  if [ $is_a_iso = true ]; then
                    check_size
                    make_live_usb
                    error_notification
                else
                    actionstatus=0
                fi ;;
                
            4)  if [ $is_a_iso = true ]; then
                    check_size
                    make_live_usb
                    error_notification
                else
                    actionstatus=0
                fi ;;
        esac 
    fi
}


display_about() {
    device_and_source
    info=$(yad --selectable-labels --window-icon=drive-removable-media-usb --align=center --form --button="$LICENSE":5 --button=OK:0 --width=300 --title=$"About Live-USB Maker GUI"\
        --field="<b>$TITLE</b>":LBL ""\
        --field=" ":LBL " "\
        --field="$VERSION 16.11":LBL ""\
        --field=" ":LBL " "\
        --field="<b>$DESCRIPTION</b>":LBL ""\
        --field=" ":LBL " "\
        --field="http://mepis.antix.org":BTN "su $launch_user -c '$BROWSER http://antix.mepis.org'"\
        --field="http://mxlinux.org":BTN "su $launch_user -c '$BROWSER https://mxlinux.org'"\
        --field=" ":LBL " "\
        --field="copyright (c) antiX":LBL "")

    case $? in

        0) actionstatus=0  ;;
        5) display_license ;;
    esac
}    
    
display_help() {
    device_and_source
    su $launch_user -c "$BROWSER https://mxlinux.org/wiki/help-files/help-live-usb-maker &"
    actionstatus=0
}

display_license() {
    su $launch_user -c "$BROWSER https://mxlinux.org/wiki/licenses/license-live-usb-maker &"
    display_about
}

quit() {
    echo quit
    actionstatus=1
}

build_usb_list() {
    #Check if running live
    local line VENDOR NAME SIZE MODEL
    Live_device=$(get_drive $(get_live_dev))
    echo  "Live Device is " $Live_device
    unset usb_list
    #build list for usb_device selection box
    while read line; do
        unset VENDOR NAME SIZE MODEL
        echo "$line"
        eval "$line"
        is_usb_or_removable $NAME
        test=$?
        echo "test is " "$test"
        if [ "$NAME" != "$Live_device" ]; then
            if [ $test = 0 ]; then
                echo $NAME $VENDOR $MODEL $SIZE
                if [ "$line" != "" ]; then
                    usb_list="$NAME   $VENDOR $MODEL $SIZE!""$usb_list"
                fi
            fi
        fi
    done<<LsBlk
$(lsblk --nodeps --pairs -no name,size,model,vendor -I 3,8,22,179,259)
LsBlk

    if [ "$usb_device" != "" ]; then
        usb_list="$usb_device"!"$usb_list"
    fi

    if [ "$usb_list" = "" ]; then
        yad --form --title="$TITLE" --button=gtk-ok --text="$USB_CHECK" --buttons-layout=center
        exit 0
    fi

    echo "usb_list is " $usb_list
    echo ""
}

error() {
    echo "Error: $*" >&2
    exit 3
}

clean_up() 
{
    test -e $FIFO     && rm -f $FIFO
    test -d $WORK_DIR && rmdir $WORK_DIR
    #reset dirty_bytes if we changed it and we have an unexpected exit
    [ "$ORIG_DIRTY_BYTES" ] && sysctl vm.dirty_bytes=$ORIG_DIRTY_BYTES >> $LOG_FILE
    [ "$ORIG_DIRTY_RATIO" ] && sysctl vm.dirty_ratio=$ORIG_DIRTY_RATIO >> $LOG_FILE
    unflock $LOCK_FILE

}

error_notification()
{

if [ $USER_STOP = true ]; then
    yad --window-icon=drive-removable-media-usb --center --width=300 --align=center --form --title="$TITLE" --button=gtk-ok --buttons-layout=center --field="<b>$STOP_MESSAGE</b>":LBL ""
    actionstatus=1
fi

if [ $USER_STOP = false ]; then    
    if test -f $ERR_FILE; then 
        echo "cli code had a fatal error"
        MESSAGE="$(cat $ERR_FILE)"
        yad --window-icon=drive-removable-media-usb --center --width=300 --form --title="$TITLE" --button=gtk-ok --text="$MESSAGE" --buttons-layout=center
        actionstatus=0
    else
        echo "cli code ran without error"
        yad --window-icon=drive-removable-media-usb --center --width=300 --align=center --form --title="$TITLE" --button=gtk-ok --buttons-layout=center --field="<b>$COMPLETE</b>":LBL ""
        actionstatus=1
    fi
fi
}

make_live_usb()
{
local lum_pid yad_pid tail_pid 

    while true; do

        yad --window-icon=drive-removable-media-usb --center --title="$TITLE" --text-info --button="$CLOSE"!gtk-close!:1 --width=880 --height=350 --tail --center <&10 &
        yad_pid=$!

        tail -n20  -f --pid=$yad_pid $PROG_FILE >&10 &
	
        if [ -z "$lum_pid" ]; then 
            $CLI_PROG gui $PRETEND $OPTION --msdos --force=flock --from="$source" -t /dev/$usb_device $PROGRESS_OPTION 2>/dev/null &
            lum_pid=$!
            if [ "$PROGRESS_OPTION" = "" ]; then
            fake_progress
            fi
	    echo "LUM PID IS " $lum_pid
        fi

        while true; do
            sleep 0.1
            test -d /proc/$yad_pid && test -d /proc/$lum_pid && continue
            break
        done

        # Suspend lum
        test -d /proc/$lum_pid && kill -STOP $lum_pid

        sleep .1

        test -d /proc/$yad_pid || echo "yad exited"
        test -d /proc/$lum_pid || echo "lum exited"

        if test -d /proc/$lum_pid; then
           kill_pids $yad_prog_pid
           yad --window-icon=drive-removable-media-usb --width=300 --align=center --form --center --title="$TITLE" --button="$NO"!gtk-close!:1 --button="$YES"!gtk-ok!:0 --buttons-layout=center --field="$PROGRESS_CLOSED":LBL ""
       	   local ret=$?
           kill -CONT $lum_pid
           case $ret in

                0) if [ "$PROGRESS_OPTION" = "" ]; then
                   fake_progress
                   else
                   echo continueing
                   fi
                   ;;

                1) USER_STOP=true && break ;;

           esac
           continue
        else
           break
        fi
    done
    #echo -n "Press <Enter> to continue ... "
    #read x
    echo "killing yad and lum as needed  ..."
    # I don't know why I no longer get the Terminate message
    kill_pids $yad_pid $lum_pid $yad_prog_pid
}

kill_pids() {
    local pid
    for pid; do
        test -z $pid && continue
        disown  $pid 2>/dev/null
        pkill --parent $pid
        kill    $pid 2>/dev/null
    done
}

flock_error()
{
local flck_pid
flck_pid=$(flock_pid)
echo $flck_pid
if [ ${#flck_pid} -gt 0 ]; then
    msg=$(printf "$MSG_LUM_LOCK_EXISTS \n $MSG_LUM_LOCK_EXISTS_2" live-usb-maker $flck_pid)
else
    msg=$(printf "$MSG_OTHER_LOCK_ERROR" $LOCK_FILE)
fi
yad_error "$msg"

}

yad_error()
{

local error_message="$1"
yad --window-icon=drive-removable-media-usb --width=300 --align=center --center --form --title="$TITLE" --button=gtk-ok --buttons-layout=center --field="<b>$error_message</b>":LBL ""
    exit 0
}

device_and_source()
{
echo $selections
usb_device=$(echo $selections|cut -d "|" -f2|cut -d " " -f1)
source=$(echo $selections|cut -d "|" -f3)
source=$(readlink -f "$source")
OPTION=$(echo $selections|cut -d "|" -f4 |cut -d ":" -f1| tr -d '[:space:]')

# figure options

case $OPTION in

       1) OPTION="";;
       2) OPTION="dd";;
        3) OPTION="--update";;
        4) OPTION="--encrypt";;
esac

echo "OPTION IN DEVICE_AND_SOURCE IS " $OPTION

}

fake_progress()
{
tail -n20  -f --pid=$yad_pid $PROG_FILE  |yad --window-icon=drive-removable-media-usb --on-top --center --undecorated --no-escape --title="$TITLE" --width=300 --progress --progress-text="$TITLE" --pulsate --no-buttons &
yad_prog_pid=$!
}

dd_progress()
{
tail -n1 -f --pid=$dd_pid $PROG_FILE |yad --window-icon=drive-removable-media-usb --on-top --center --undecorated --no-escape --title="$TITLE" --width=300 --progress --percentage="0" --progress-text="$dd_title" --button=gtk-cancel --auto-kill &
yad_prog_pid=$!
}

add_encrypt_checkbox()
{
echo encypt test
echo MODE = $MODE

##MODE 1 is clone, no encyrption available
##MODE 2 is standard with no encryption available
##MODE 3 is clone with encryption available
##MODE 4 is standard iso selection with encryption available

case $MODE in

    1) MODE=3;;
    2) MODE=4;;

esac
}

dd_or_liveUSB()
{
if [ "$OPTION" != "dd" ]; then
    make_dd_usb
else
    make_live_usb
fi
}

make_dd_usb()
{
    local trip dd_pid check_yad check_dd start_io target_io iso_io current_io percent_complete
    iso_io=$(du_size $source)
    echo "iso_io is: " $iso_io
    #convert to sectors
    iso_io=$(($iso_io * 1024 / 512 * 1024 ))
    echo "adjusted iso_io size is: " $iso_io
    start_io=$(cat /sys/block/"$usb_device"/stat |awk '{print $7}')
    echo "start+io is: " $start_io
    target_io=$(($start_io + iso_io))
    echo "target_io is : " $target_io    

    #adjust vm.dirty.bytes for progress meter and speed bump
    ORIG_DIRTY_RATIO=$(sysctl -n vm.dirty_ratio)
    ORIG_DIRTY_BYTES=$(sysctl -n vm.dirty_bytes)
    usb_dirty_bytes_gui=20000000
    sysctl vm.dirty_bytes=$usb_dirty_bytes_gui >> $LOG_FILE
    #dd conv=fdatasync bs=1M if="$source" of=/dev/$usb_device &
    dd bs=1M if="$source" of=/dev/$usb_device &
    dd_pid=$!
    echo "dd pid is " $dd_pid
    dd_progress
    echo "yad pid is " $yad_prog_pid
    trip=0
    while [ $trip = 0 ]; do
        #echo "dd pid is " $dd_pid
        #echo "yad pid is " $yad_prog_pid
        ps --pid $dd_pid 2>/dev/null
        check_dd=$?
        echo "check_dd is " $check_dd
        ps --pid $yad_prog_pid 2>/dev/null
        check_yad=$?
        echo "check_yad is " $check_yad
        if [ $check_dd = 0 ]; then
            if [ $check_yad != 0 ]; then
            USER_STOP=true
            trip=1
            else
            trip=0
            current_io=$(cat /sys/block/"$usb_device"/stat |awk '{print $7}')
            echo "current_io is : " $current_io
            percent_complete=$(awk "BEGIN {print ($current_io - $start_io)/$iso_io * 100}")
            echo "percent complete is " $percent_complete
            echo $percent_complete > $PROG_FILE
            sleep 1
            fi
         else
            trip=1
         fi
    done        
    kill_pids $dd_pid $yad_prog_pid 
    sync
    #put vm.dirty.bytes back the way we found theme
    sysctl vm.dirty_bytes=$ORIG_DIRTY_BYTES >> $LOG_FILE
    sysctl vm.dirty_ratio=$ORIG_DIRTY_RATIO >> $LOG_FILE
    unset ORIG_DIRTY_BYTES ORIG_DIRTY_RATIO
      
}

check_size()
{
local total_size iso_size less_esp
total_size=$(parted --script /dev/$usb_device unit MiB print 2>/dev/null | sed -rn "s/^Disk.*: ([0-9]+)MiB$/\1/p")
iso_size=$(du_size $source)

echo "total size is " $total_size
echo "iso size is " $iso_size

less_esp=$(($total_size - 50))

echo "less_esp is " $less_esp

if [ "$less_esp" -gt "$iso_size" ]; then
    echo "size check ok"
else
    msg=$(printf $"USB device is too small for this operation")
    yad_error "$msg"
fi
}

check_unmountable()
{
local ret
umount_all $usb_device
ret=#?
if [ $ret = 1 ]; then
    msg=$(printf "$msg" $usb_device)
    yad_error "$msg"
fi
}



#launch gui

main "$@"

