#!/bin/bash
########################################################################
#  Copyright (C) 2015 Max Mehl <mail@mehl.mx>
########################################################################
#  
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#  
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#  
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
#  
########################################################################
#  
#  This script can mount SFTP drives graphically.
#  To obtain SFTP drives, it reads out the ~/.ssh/config file, can be 
#   configured manually or directly in the program
#
#  Read my blog post for more information:
#  http://blog.mehl.mx/2014/mounting-a-sftp-storage-in-gnu-linux/
#  
########################################################################


cd "$(dirname "$(readlink -f "$0")")"
# Test if config.cfg exists and set needed variables
if [ ! -e config.cfg ]; then echo "Missing config.cfg file. Edit and rename config.cfg.sample"; exit 1; fi
source config.cfg


# Add SSH key to local keyring if not already happened
function sshadd {
  ssh-add -l >/dev/null || ssh-add
}

for ((i = 0; i < ${#PRESSH[*]}; i++))
do
  PREUSER[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $1 }')
  PREHOST[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $2 }')
  PREPORT[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $3 }')
  PREPASS[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $4 }')
  PREPATH[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $5 }')
  PRELOCAL[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $6 }')
  
  PRESSHLIST[$i]=${PREUSER[$i]}"@"${PREHOST[$i]}":"${PREPATH[$i]}
done

# Read all hosts from ~/.ssh/config 
SSHLIST=$(grep -E "Host\s[[:alnum:]]" $SSHCONFIG | awk -F' ' '{ print $2 }')

# Choose preconfigured HOST to mount
function mount {
  if ! SSH=$(zenity --list \
    --height=400 \
    --width=350 \
    --text="Please choose. Cancel to unmount drives." \
    --title="Choose SSH server" \
    --cancel-label "Unmount all / Cancel" \
    --column "Available SSH servers" \
    "<manual input>" $SSHLIST ${PRESSHLIST[*]}); then
    unmountquestion   # If you press cancel, it should ask you to unmount all drives
  fi
  
  # If you double click on an entry, it gives 'server1|server1' as result instead of 'server1'
  # This command cuts of everything after |
  SSH=$(echo $SSH | awk -F\| '{ print $1 }')
  
  ## Recognize if the input was entered manually, by .ssh/config or by the predefined SSH list in config.cfg
  
  ISPRE=$(echo ${PRESSHLIST[@]} | grep -q "$SSH" ; echo $?)
  
  if [ "$SSH" == "<manual input>" ]; then
    MAN="1"
    SSHUSER=$(zenity --entry --text "SSH user" --title "Enter SSH user")
    checkexit $?
    SSHHOST=$(zenity --entry --text "SSH host" --title "Enter SSH host")
    checkexit $?
    SSHPORT=$(zenity --entry --text "SSH port (default 22)" --title "Enter SSH port")
    checkexit $?
    SSHPASS=$(zenity --entry --hide-text --text "SSH password" --title "Enter SSH password")
    checkexit $?
    SSHPATH=$(zenity --entry --text "Desired root directory on server (default /)" --title "Enter SSH root directory")
    checkexit $?
    if [ "$SSHPORT" == "" ]; then
      SSHPORT=22
    fi
    if [ "$SSHPATH" == "" ]; then
      SSHPATH="/"
    fi
    
    SSH="m-$SSHUSER-$RANDOM"
  elif [ "$ISPRE" == "0" ]; then
    MAN="1"
    
    # Detect which predefined host was chosen and set variables accordingly
    for ((i = 0; i < ${#PRESSHLIST[*]}; i++))
    do
      if [ "$(echo ${PRESSHLIST[$i]} | grep -q "$SSH" ; echo $?)" == "0" ]; then
        SSHUSER=${PREUSER[$i]}
        SSHHOST=${PREHOST[$i]}
        SSHPORT=${PREPORT[$i]}
        SSHPASS=${PREPASS[$i]}
        SSHPATH=${PREPATH[$i]}
        SSHLOCAL=${PRELOCAL[$i]}
      fi
    done

    SSH="p-$SSHUSER-$RANDOM"
  else
    MAN="0"
  fi
  
  if [ "$SSHLOCAL" != "" ]; then
    MOUNTDIR="$SSHLOCAL"
  else
    MOUNTDIR="$LOCALMOUNTDIR"/"$SSH"
  fi
  
  # Make a local directory if not available
  if [ ! -e "$MOUNTDIR" ]; then
    mkdir -p "$MOUNTDIR"
  fi

  # Command to mount actually
  if [ "$MAN" = "0" ]; then
    sshfs -C "$SSH": "$MOUNTDIR"/ -o follow_symlinks &
  else
    echo "$SSHPASS" | sshfs -C -o password_stdin -p "$SSHPORT" "$SSHUSER"@"$SSHHOST":"$SSHPATH" "$MOUNTDIR"/ -o follow_symlinks &
  fi

  quitquestion  # one more ssh server or quit?

}

# Ask if all preconfigured SSHFS drives should be unmounted
function unmountquestion {
  zenity --question --text="Unmount all configured\nSSHFS drives now?"
  if [ "$?" = "0" ]; then
    unmount   # unmount function
  else
    exit 0
  fi
}

# Procedure to unmount all preconfigured SSHFS drives and exit program afterwards
function unmount {
  
  for FILE in "$LOCALMOUNTDIR"/*;
  do
    if [ "$(cat /etc/mtab | grep -q "$FILE"; echo $?)" == "0" ]; then
      fusermount -u "$FILE"
      echo "$FILE unmounted."
    fi
  done
  
  # Go through predefined mounts and unmount them if they're mounted
  for ((i = 0; i < ${#PRESSHLIST[*]}; i++))
    do
      PREUSER[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $1 }')
      PREHOST[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $2 }')
      PREPATH[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $5 }')
      PRELOCAL[$i]=$(echo ${PRESSH[$i]} | awk -F";" '{ print $6 }')
      SSH[$i]=${PREUSER[$i]}"@"${PREHOST[$i]}":"${PREPATH[$i]}
      if [ "$(cat /etc/mtab | grep -q "${SSH[$i]}"; echo $?)" == "0" ]; then
        fusermount -u ${PRELOCAL[$i]}
        echo "${PRELOCAL[$i]} unmounted."
      fi
    done
  
  # Remove all empty directories which have been created by the manual/predefined mount process (m-XYZ-$RANDOM)
  find "$LOCALMOUNTDIR"/ -maxdepth 1 -type d -empty -regextype sed -regex ".*/[pm]-.\+\?-[0-9]\{2,6\}" -delete
  
  exit 0
}

# Should another SSHFS storage be mounted?
function quitquestion {
  zenity --question \
  --text="Mount another SSHFS storage?" \
  --ok-label="No and exit" \
  --cancel-label="Yes"
  if [ "$?" = "0" ]; then
    exit 0
  fi
}

# Check whether one clicked "cancel" in zenity question
function checkexit {
  if [ $1 = 1 ]; then exit 1; fi
}

sshadd    # sshadd function

# Loop for endless mounts until stopped by unmount or unmountquestion
while :
do
  mount # mount function
done