#!/bin/bash
# Perform mount and umount of log disks (either ext4 or luks format) and deploy
# usb disk (ext4) for main-ai.
# Inputs arguments: ACTION, ID_FS_TYPE, ID_FS_LABEL (empty for LUKS disk), DEVNAME
# e.g., add crypto_LUKS  MKNSSDRE1TB_MK1603071001AB6E3 /dev/sda1 8 1
# e.g., remove crypto_LUKS  MKNSSDRE1TB_MK1603071001AB6E3 /dev/sda1 8 1

# This function should only be called through udev rule

# Enable unofficial bash strict mode.
set -euo pipefail

# Constants
readonly key_dir="/tmp/"
readonly map_point="/dev/mapper/"
readonly mount_point="/media/zoox/"
readonly log_file="/var/log/zooxops/log_drive_mount"

function main() {
  # Looks for any unencrypted drives that are added or removed and attempts to
  # mount/unmount them
  # :param action: A change of drive state for the system (e.g. added, removed,
  # changed)
  # :param fs_type: The type of filesystem (FS_TYPE=...) of the device (encrypted
  # ext4 drives do not show up as ext4)
  # :param fs_label: The device label associated with a drive (e.g. zooxdata-263)
  local action; action="$1"; shift;
  local fs_type; fs_type="$1"; shift;
  local fs_label; fs_label="$1"; shift;
  local devname; devname="$1"; shift;

  # Create log file if doesn't exist
  if [[ ! -f "$log_file" ]]; then
    if [[ ! -d /var/log/zooxops ]]; then
      mkdir --parents /var/log/zooxops
    fi
    touch "$log_file"
  fi
  # Set a lock for the device name
  # This can be safely called simultaneously on different devices
  if [[ ! -d /var/lock/dev ]]; then
    mkdir --parents /var/lock/dev
  fi
  set_lockfile_exit_trap "/var/lock/$devname"

  echo "Udev Action Caught: $action $fs_type $fs_label $devname" >> "$log_file"

  if [[ "$fs_type" != "ext4" ]]; then
    echo "File system types other than 'ext4' are not supported" >> "$log_file"
    usage 122
  fi

  if [[ "$fs_label" != zooxusb* ]] && [[ "$fs_label" != zooxdata-* ]] &&
     [[ "$fs_label" != zooxdeploy* ]]; then
    echo "Device is not a Zoox logging or deploy drive" >> "$log_file"
    usage 122
  fi

  case "$action" in
    add)
      mount_ext4 $fs_label $devname
      ;;

    remove)
      umount_ext4 $fs_label
      ;;

    *)
      echo "unexpected action: $action" >> "$log_file"
      usage 133
  esac
}

trap_function() {
  # Trap function for the lockfile exit trap
  # :param lock_file: The lock file
  rc="$?"
  local lock_file="$1"; shift;

  case "$rc" in
    0)
      echo "Cleanly exited" >> "$log_file"
      ;;

    100)
      echo "Unencrypted drive, mounting successful: exiting gracefully" >> "$log_file"
      ;;

    101|111)
        echo "Logging Drive Mounting failed with status ${rc}" >> "$log_file"
      ;;

    122|133)
      echo "Unhandled action requested, exiting gracefully" >> "$log_file"
      ;;

    *)
      echo "Logging Drive encountered unexpected error with status ${rc}" >> "$log_file"
  esac

  rm --force ${lock_file}
}

# Sets a lockfile exit trap.
set_lockfile_exit_trap() {
  #Sets a lockfile exit trap
  local lock_file="$1"; shift;
  # Attempt to write the lock file by setting 'noclobber' then trying to
  # write the file (or overwrite if it is there).  If this fails, exit.
  # Otherwise
  if ! (set -o noclobber; : > "${lock_file}") ; then
    echo "$(date --utc): ${lock_file} exists; another instance is running." >> "$log_file"
    exit 0
  fi
  trap "trap_function $lock_file" EXIT
}

function usage() {
  # Exits function and displays correct usage information
  # :param code: exit code to return
  local code; code="$1"; shift;

  echo "usage: device_action fs_type fs_label device_id" >> "$log_file"
  exit "$code"
}

function umount_ext4() {
  # Unmounts a drive that has been removed. If there is an encrypted counterpart
  # drive, also unmount that.
  # :param label: The device label associated with a drive (e.g. zooxdata-263)
  local label; label="$1"; shift;

  echo "Unmount Ext4 $label" >> $log_file

  # Verify metadata partition is mounted
  if [[ -d "$mount_point/$label" ]]; then
    # Verify encrypted partition is mapped in memory
    if [[ -L "$map_point/$label" ]]; then
      umount_enc "$label"
    fi

    umount "$mount_point/$label"
    rmdir "$mount_point/$label"
  fi
}

function umount_enc() {
  # Unmounts an encrypted drive when its unencrypted counterpart is removed.
  # :param label: The device label associated with a drive (e.g. zooxdata-263)
  # the encrypted label is derived by adding the "_e" extension
  local label; label="$1"; shift;
  local enc_label; enc_label="$label"; enc_label+="_e";

  echo "Unmount Enc $label" >> $log_file

  # Verify encrypted partition is mounted
  if [[ -d "$mount_point/$enc_label" ]]; then
    echo "Unmounting Encrypted Disk $enc_label" >> $log_file
    umount "$mount_point/$enc_label"
    rmdir "$mount_point/$enc_label"

    # Verify that encrypted partition is mapped in memory
    if [[ -L "$map_point/$label" ]]; then
      cryptsetup luksClose "$map_point/$label"
    fi
  fi
}

function mount_ext4() {
  # Mounts a drive that has been added. If there is an encrypted counterpart
  # drive, also attempt to mount that
  # :param label: The device label associated with a drive (e.g. zooxdata-263)
  # :param devname: The device id associated with a drive (e.g. sda1)
  local label; label="$1"; shift;
  local devname; devname="$1"; shift;

  echo "Mount Ext4 $label $devname" >> $log_file

  # Make sure a directory exists where we are mounting
  if [[ ! -d "$mount_point/$label" ]]; then
    mkdir --parents "$mount_point/$label"
  fi

  mount "$devname" "$mount_point/$label"

  # Give correct permissions to USB drives
  if [[ "$label" =~ zooxusb* ]]; then
    chown --recursive zoox:zoox "$mount_point/$label"
  fi

  # Identify if the mounted drive is part of a set (multiple partitions) and is the first partition
  if [[ "$devname" =~ sd.1 ]]; then
    mount_enc $(echo "$devname" | sed 's/.$//') "$label"
  fi
}

function mount_enc() {
  # Mounts an encrypted drive when its unencrypted counterpart is added. Checks
  # to see if it can download the associated key, then if so decrypts the drive.
  # :param devbase: The device portion of the device id (e.g. sda instead of
  # sda1)
  # :param label: The device label associated with a drive (e.g. zooxdata-263)
  # devname is used to refer to the encrypted partition
  # devmeta is used to refer to the unencrypted partition
  # enc_label is derived by adding the "_e" extension to the label
  local devbase; devbase="$1"; shift;
  local label; label="$1"; shift;
  local devname; devname="$devbase"; devname+="2";
  local devmeta; devmeta="$devbase"; devmeta+="1";
  local enc_label; enc_label=$label; enc_label+="_e";
  local uuid

  echo "Mount Enc $label $devbase" >> $log_file

  # Verify metadata partition is mounted
  if [[ ! -d "$mount_point/$label" ]]; then
    echo "Metadata partition has dissapeared before encrypted partition could be mounted" >> "$log_file"
    exit 101
  fi

  if [[ ! -s "$mount_point/$label/uuid" ]]; then
    echo "Not an encrypted drive" >> "$log_file"
    exit 100
  fi
  uuid=$(cat "$mount_point/$label/uuid")

  #Verify key matching uuid doesn't already exist
  if [[ ! -s "$key_dir/$uuid.key" ]]; then
    # Query HWSD for key matching uuid
    echo "Attempt to Curl Key: $uuid from HWSD" >> $log_file
    #Make sure key file doesn't exist
    rm -f "$key_dir/$uuid.key"
    curl --silent "http://hwsd.zooxlabs.com:6060/api/run/v1/run?test_name=$uuid" | grep resultValue | head -1 | xargs | sed 's/resultValue: //' | base64 --decode > "$key_dir/$uuid.key"
  fi

  # Verify a key matching uuid was found
  if [[ -s "$key_dir/$uuid.key" ]]; then
    echo "Attempt to Decrypt with key: $uuid" >> $log_file
    cryptsetup --key-file "$key_dir/$uuid.key" luksOpen "$devname" "$label"

    # If directory for mount  doesn't exist create it
    if [[ ! -d "$mount_point/$enc_label" ]]; then
      mkdir --parents "$mount_point/$enc_label"
    fi

    mount "$map_point/$label" "$mount_point/$enc_label"

    # Give correct permissions to USB drives
    if [[ "$label" =~ zooxusb* ]]; then
      chown --recursive zoox:zoox "$mount_point/$enc_label"
    fi
  else
    echo "No Key Found for $uuid" >> "$log_file"
    usage 111
  fi

}

main $@
