#!/usr/bin/env bash
set -e

unmount_usb() {
  if mountpoint -q /mnt/usb; then
    sudo umount /mnt/usb
    echo -e "\e[0;32mUSB drive unmounted.\e[0m"
  fi
}

mount_usb() {
  if mountpoint -q /mnt/usb; then
    echo -e "\e[0;32mUSB drive already mounted.\e[0m"
  else
    device_found=false
    for dev in sda sdb sdc sdd sde sdf sdg sdh sdi sdj sdk sdl; do
      if sudo blkid /dev/$dev | grep -iq 'TYPE="vfat"'; then
        device_found=true
        mkdir -p /mnt/usb
        sudo mount /dev/$dev /mnt/usb && { echo -e "\e[0;32mUSB drive mounted successfully on /dev/$dev.\e[0m"; break; } || echo -e "\e[0;31mFailed to mount /dev/$dev.\e[0m"
      fi
    done
    if [ "$device_found" = false ]; then
      echo -e "\e[0;31mNo USB devices found.\e[0m"
    fi
  fi
}

setup_ssh_directory() {
  export SSH_DIR=/root/.ssh
  mkdir -p $SSH_DIR
}

check_file_exists() {
  if [[ ! -f $1 ]]; then
    echo -e "\e[0;31mError: File $1 does not exist.\e[0m"
    exit 1
  fi
}

copy_keys() {
  check_file_exists "/mnt/usb/id_ed25519_agenix.pub"
  check_file_exists "/mnt/usb/id_ed25519_agenix"
  cp /mnt/usb/id_ed25519_agenix.pub $SSH_DIR
  cp /mnt/usb/id_ed25519_agenix $SSH_DIR
  chmod 600 $SSH_DIR/id_ed25519_{agenix,agenix.pub}
  echo -e "\e[0;32mKeys copied successfully.\e[0m"
}

set_keys() {
  check_file_exists "/mnt/usb/id_ed25519_github.pub"
  check_file_exists "/mnt/usb/id_ed25519_github"
  cp /mnt/usb/id_ed25519_github.pub $SSH_DIR/id_ed25519.pub
  cp /mnt/usb/id_ed25519_github $SSH_DIR/id_ed25519
  chmod 600 $SSH_DIR/id_ed25519
  chmod 644 $SSH_DIR/id_ed25519.pub
}

change_ownership() {
  chown nixos:wheel $SSH_DIR/id_ed25519{,.pub}
  chown nixos:wheel $SSH_DIR/id_ed25519_{agenix,agenix.pub}
}

trap unmount_usb EXIT

setup_ssh_directory
mount_usb
copy_keys
set_keys
change_ownership
unmount_usb