#!/bin/bash

# Script for mounting/unmounting drives (including encrypted ones)
# For encrypted drives, it works only if there's just one partition inside the
# crypt_LUKS volume
# It now also allows to power off devices so that you can safely disconnect
# them. It lists only USB devices.

# Yaroslav de la Peña Smirnov

readarray -t blkdevs <<< "$(lsblk -r -o NAME,TYPE,FSTYPE,MOUNTPOINT | tail -n +2)"

notify_and_end() {
	notify-send "💽Drive management" "$1"
	exit $2
}

get_mountable() {
  for blkdev in "${blkdevs[@]}"; do
    name="$(cut -d " " -f1 <<< "${blkdev}")"
    btype="$(cut -d " " -f2 <<< "${blkdev}")"
    fstype="$(cut -d " " -f3 <<< "${blkdev}")"
    mountpoint="$(cut -d " " -f4 <<< "${blkdev}")"
    if [ "$btype" = "part" ]; then
      if [ -z "$mountpoint" ]; then
        if [ "$fstype" = "crypto_LUKS" ]; then
          if [ "$(lsblk -r -o FSTYPE "/dev/$name" | tail -n 1)" = "crypto_LUKS" ]; then
            echo "$name"
          fi
        else
          echo "$name"
        fi
      fi
    fi
  done
}

get_unmountable() {
  for blkdev in "${blkdevs[@]}"; do
    name="$(cut -d " " -f1 <<< "${blkdev}")"
    btype="$(cut -d " " -f2 <<< "${blkdev}")"
    fstype="$(cut -d " " -f3 <<< "${blkdev}")"
    mountpoint="$(cut -d " " -f4 <<< "${blkdev}")"
    if [ "$btype" = "part" ]; then
      if [ -z "$mountpoint" ] && [ "$fstype" = "crypto_LUKS" ]; then
        mountpoint="$(lsblk -r -o MOUNTPOINT "/dev/$name" | tail -n 1)"
      fi
      if [ -n "$mountpoint" ] && [ "$mountpoint" != "/" ] && [ "$mountpoint" != "/boot" ] && [ "$mountpoint" != "[SWAP]" ]
      then
        echo "$name"
      fi
    fi
  done
}

mount_dev() {
	fstype="$(lsblk -r -o FSTYPE "/dev/$1" | tail -n 1)"
	part="/dev/$1"
	if [ "$fstype" = "crypto_LUKS" ]; then
		password="$(echo "" | rofi -dmenu -password -p "Enter passphrase for encrypted drive" --tf="$BEMENU_TF" --tb="$BEMENU_NB" --ff="$BEMENU_NB" --fb="$BEMENU_NB" --fn="$BEMENU_FN" --nb="$BEMENU_NB" --nf="$BEMENU_NF" --hf="$BEMENU_HF" --hb="$BEMENU_HB" --monitor="$BEMENU_MONITOR")"
		if [ -z "$password" ]; then
  	  	  exit 0
		fi
		printf "%s" "$password" > "$HOME/.cache/diskey"
		part="$(udisksctl unlock --key-file="$HOME/.cache/diskey" -b "/dev/$1" 2> $HOME/.cache/udiskerr | cut -d " " -f4 )"
		part="${part%?}"
		rm -f "$HOME/.cache/diskey"
		if [ $? -ne 0 ]; then
			notify_and_end "Error mounting: $(cat $HOME/.cache/udiskerr)" 1
		fi
	fi
	message="$(udisksctl mount -b "$part" 2> $HOME/.cache/udiskerr)"
	if [ $? -eq 0 ]; then
		notify_and_end "$message" 0
	else
		notify_and_end "Error mounting: $(cat $HOME/.cache/udiskerr)" 1
	fi
}

unmount_dev() {
  fstype="$(lsblk -r -o FSTYPE "/dev/$1" | sed "2q;d")"
  if [ "$fstype" = "crypto_LUKS" ]; then
    volume="/dev/mapper/$(lsblk -r -o NAME "/dev/$1" | tail -n 1)"
    if udisksctl unmount -b "$volume" 2> $HOME/.cache/udiskerr && udisksctl lock -b "/dev/$1" 2> $HOME/.cache/udiskerr; then
        notify_and_end "Device $1 has been unmounted and locked" 0
    else
		notify_and_end "Error unmounting: $(cat $HOME/.cache/udiskerr)" 1
    fi
  else
    if udisksctl unmount -b "/dev/$1" 2> $HOME/.cache/udiskerr; then
        notify_and_end "Device $1 has been unmounted" 0
    else
		notify_and_end "Error unmounting: $(cat $HOME/.cache/udiskerr)" 1
    fi
  fi
}

get_usb_storage() {
	for dev in /sys/block/*; do
		if udevadm info --query=property --path=$dev | grep -q ^ID_BUS=usb; then
			echo "$(cut -d "/" -f4 <<< ${dev})"
		fi
	done
}

poweroff_dev() {
	if udisksctl power-off -b "/dev/$1"; then
		notify_and_end "Device $1 has been powered off" 0
	else
		notify_and_end "An error occurred, could not power off device $1. Make sure that it's not in use." 1
	fi
}

action="$(printf "mount\nunmount\npoweroff\n" | rofi -dmenu -matching prefix -p "What to do?" )"

case "$action" in
  mount)
    blkdev="$(get_mountable | rofi -dmenu -p "Choose device to mount" )"
    if [ -z "$blkdev" ]; then
      exit 0
    fi
    mount_dev "$blkdev";;
  unmount)
    blkdev="$(get_unmountable | rofi -dmenu -p "Choose device to umount" )"
    if [ -z "$blkdev" ]; then
      exit 0
    fi
    unmount_dev "$blkdev";;
  poweroff)
    blkdev="$(get_usb_storage | rofi -dmenu -p "Choose device to power off" )"
    if [ -z "$blkdev" ]; then
      exit 0
    fi
    poweroff_dev "$blkdev";;
  *)
    exit 0;;
esac