#!/bin/bash

# configure reads TITLE and DESCRIPTION
# shellcheck disable=SC2034
TITLE="Network storage"
# shellcheck disable=SC2034
DESCRIPTION="Mount shared folders from your network in /mnt"

# What's mounted, one "kind|source|mount point" per line, e.g.
# smb|//pi-star.local/files|/mnt/pi-star-files. It exists once you've run
# this, even if you mounted nothing, so this isn't "to do" without a server.
STORAGE_CONF="$HOME/.config/coredesktop/storage"
# Passwords for SMB shares, readable only by root, which mounts them
SMB_CREDENTIALS="/etc/coredesktop/smb"
BOOKMARKS="$HOME/.config/gtk-3.0/bookmarks"
source "$COREDESKTOP/configure.d/lib/mdns.sh"

is_installed() {
  command -v avahi-browse
}

is_configured() {
  [ -f "$STORAGE_CONF" ]
}

# A name usable in a path: "Pi-Star Files" -> pi-star-files
slug() {
  local s
  s="$(tr '[:upper:]' '[:lower:]' <<< "$1" | sed -E 's/[^a-z0-9._-]+/-/g; s/^-+//; s/-+$//')"
  echo "${s:-share}"
}

# /mnt/<host>[-<name>], with -2, -3... if that's taken
mount_point() {
  local base n=2 point
  base="/mnt/$(slug "${1%.local}")"
  [ -n "$2" ] && base="$base-$(slug "$2")"
  point="$base"
  while [ -e "$point" ] || grep -qsF "|$point" "$STORAGE_CONF"; do
    point="$base-$n"
    n=$((n + 1))
  done
  echo "$point"
}

# --- SMB: Windows, Samba (e.g. a NAS) and Mac file sharing --------------------

# The disk shares on an SMB server, one per line, leaving out admin shares
# like C$. smbclient reads the password from PASSWD, not the command line.
smb_shares() {
  local host="$1" user="$2" password="$3" auth=(-N)
  [ -n "$user" ] && auth=(-U "$user")
  PASSWD="$password" timeout 20 smbclient -L "//$host" "${auth[@]}" -g 2> /dev/null \
    | awk -F'|' '$1 == "Disk" && $2 !~ /\$$/ { print $2 }'
}

# The /etc/fstab line: root mounts the share the first time something opens
# the folder (x-systemd.automount), as you, and unmounts it after 10 idle
# minutes. nofail and the timeout keep a missing server from holding up boot
# or hanging whatever opened the folder. Spaces are \040 in fstab.
smb_fstab_line() {
  local source="${1// /\\040}" point="$2" credentials="$3" auth
  if [ -n "$credentials" ]; then auth="credentials=$credentials"; else auth="guest"; fi
  echo "$source $point cifs $auth,uid=$(id -u),gid=$(id -g),iocharset=utf8,nofail,_netdev,x-systemd.automount,x-systemd.idle-timeout=600,x-systemd.mount-timeout=20 0 0"
}

add_smb() {
  local host="$1" share="$2" user password shares point source credentials="" choice i
  echo "Sign in to $host. Leave the user name empty for a share anyone can open."
  read -rp "User name: " user
  password=""
  if [ -n "$user" ]; then
    read -rsp "Password: " password
    echo
  fi

  if [ -z "$share" ]; then
    mapfile -t shares < <(smb_shares "$host" "$user" "$password")
    if [ ${#shares[@]} -eq 0 ]; then
      echo "$host didn't list any shares for ${user:-guests} (wrong password, or it doesn't list them)."
      read -rp "The share's name, if you know it (Enter to stop): " share
      [ -n "$share" ] || return 1
    else
      for i in "${!shares[@]}"; do
        printf '  %2d) %s\n' $((i + 1)) "${shares[$i]}"
      done
      read -rp "Which share? " choice
      if ! [[ "$choice" =~ ^[0-9]+$ ]] || (( choice < 1 || choice > ${#shares[@]} )); then
        echo "Choose a number from the list."
        return 1
      fi
      share="${shares[$((choice - 1))]}"
    fi
  fi

  source="//$host/$share"
  point="$(mount_point "$host" "$share")"
  if [ -n "$user" ]; then
    credentials="$SMB_CREDENTIALS/$(basename "$point")"
    printf 'username=%s\npassword=%s\n' "$user" "$password" \
      | sudo install -D -m 600 -o root -g root /dev/stdin "$credentials" || return 1
  fi
  sudo install -d "$point" || return 1
  smb_fstab_line "$source" "$point" "$credentials" | sudo tee -a /etc/fstab > /dev/null || return 1
  sudo systemctl daemon-reload || return 1
  sudo systemctl restart "$(systemd-escape -p --suffix=automount "$point")" || return 1
  finish_add smb "$source" "$point" "$share on ${host%.local}"
}

remove_smb() {
  local point="$1" tmp
  sudo systemctl stop "$(systemd-escape -p --suffix=automount "$point")" \
    "$(systemd-escape -p --suffix=mount "$point")" 2> /dev/null
  tmp="$(mktemp)" || return 1
  awk -v point="$point" '!($2 == point && $3 == "cifs")' /etc/fstab > "$tmp"
  sudo cp /etc/fstab /etc/fstab.bak && sudo install -m 644 "$tmp" /etc/fstab || { rm -f "$tmp"; return 1; }
  rm -f "$tmp"
  sudo systemctl daemon-reload
  sudo rm -f "$SMB_CREDENTIALS/$(basename "$point")"
  sudo rmdir "$point" 2> /dev/null
  return 0
}

# --- SFTP: any machine you can ssh to --------------------------------------------

sftp_unit_name() {
  echo "coredesktop-sftp-$(basename "$1").service"
}

# systemd would expand $ and % in the unit's ExecStart
sftp_unit() {
  local source="$1" point="$2"
  cat << EOF
# Written by CoreDesktop (configure storage)
[Unit]
Description=Mount $source at $point
After=network-online.target

[Service]
ExecStart=%h/.local/share/coredesktop/system/sftp-mount "$source" "$point"
ExecStop=/usr/bin/fusermount3 -u "$point"
Restart=on-failure
RestartSec=30

[Install]
WantedBy=default.target
EOF
}

add_sftp() {
  local host="$1" user="${2:-$USER}" path="$3" source point name unit
  if ! command -v sshfs > /dev/null; then
    echo "Installing sshfs..."
    sudo pacman -S --needed --noconfirm sshfs || return 1
  fi
  read -rp "User name on $host [$user]: " name
  user="${name:-$user}"
  if [ -z "$path" ]; then
    read -rp "Folder on $host (Enter for your home folder): " path
  fi
  if [[ "$user$host$path" == *[\$%\"\\]* ]]; then
    echo "The user name and folder can't contain \$, %, \" or \\."
    return 1
  fi
  source="$user@$host:$path"

  # Accepts the server's key the first time; the mount can't ask
  echo "Checking that $host accepts one of your SSH keys..."
  if [ -S "$HOME/.bitwarden-ssh-agent.sock" ]; then
    export SSH_AUTH_SOCK="$HOME/.bitwarden-ssh-agent.sock"
  fi
  ssh -o ConnectTimeout=10 -o PasswordAuthentication=no -o KbdInteractiveAuthentication=no \
    "$user@$host" true || {
    echo "$host didn't accept any of your SSH keys, and the mount can't ask for a password."
    echo "Add your key to $host first (e.g. ssh-copy-id $user@$host), then try again."
    return 1
  }

  point="$(mount_point "$host" "$(basename "${path:-}")")"
  unit="$(sftp_unit_name "$point")"
  sudo install -d -o "$USER" -g "$(id -gn)" "$point" || return 1
  mkdir -p "$HOME/.config/systemd/user"
  sftp_unit "$source" "$point" > "$HOME/.config/systemd/user/$unit" || return 1
  systemctl --user daemon-reload
  systemctl --user enable --now "$unit" || return 1
  finish_add sftp "$source" "$point" "${host%.local}${path:+ $(basename "$path")}"
}

remove_sftp() {
  local point="$1" unit
  unit="$(sftp_unit_name "$point")"
  systemctl --user disable --now "$unit" 2> /dev/null
  rm -f "$HOME/.config/systemd/user/$unit"
  systemctl --user daemon-reload
  sudo rmdir "$point" 2> /dev/null
  return 0
}

# --- Both ------------------------------------------------------------------------

# Record it, add it to the Files sidebar and check it opens
finish_add() {
  local kind="$1" source="$2" point="$3" label="$4"
  echo "$kind|$source|$point" >> "$STORAGE_CONF"
  mkdir -p "$(dirname "$BOOKMARKS")"
  grep -qs "^file://$point " "$BOOKMARKS" || echo "file://$point $label" >> "$BOOKMARKS"
  if timeout 30 ls "$point" > /dev/null 2>&1 && mountpoint -q "$point"; then
    echo "$source is at $point (\"$label\" in the Files sidebar)."
  else
    echo "$source was added at $point, but didn't open yet. If it doesn't, see:"
    if [ "$kind" = smb ]; then
      echo "  journalctl -b -u '$(systemd-escape -p --suffix=mount "$point")'"
    else
      echo "  journalctl --user -u $(sftp_unit_name "$point")"
    fi
  fi
}

remove_storage() {
  local kind="$1" point="$3" tmp
  "remove_$kind" "$point" || return 1
  tmp="$(mktemp)" || return 1
  grep -vF "|$point" "$STORAGE_CONF" > "$tmp"
  mv -f "$tmp" "$STORAGE_CONF"
  if [ -f "$BOOKMARKS" ]; then
    grep -v "^file://$point " "$BOOKMARKS" > "$tmp"
    mv -f "$tmp" "$BOOKMARKS"
  fi
  echo "$point is no longer mounted."
}

# The file servers on the network: "kind|host|name" per line
find_storage() {
  local own name type host
  own="$(mdns_own_host)"
  mdns_services | while IFS='|' read -r name type host _; do
    [[ "${host,,}" == "${own,,}" ]] && continue
    case "$type" in
      _smb._tcp) echo "smb|$host|$name" ;;
      _sftp-ssh._tcp) echo "sftp|$host|$name" ;;
    esac
  done | sort -u
}

find_and_add() {
  local found kind host name choice i
  echo "Looking for shared folders on your network..."
  mapfile -t found < <(find_storage)
  if [ ${#found[@]} -eq 0 ]; then
    echo "Nothing found. A Mac shares folders over SMB once File Sharing is on (System Settings >"
    echo "General > Sharing). Or enter the address instead."
    return 0
  fi
  for i in "${!found[@]}"; do
    IFS='|' read -r kind host name <<< "${found[$i]}"
    printf '  %2d) %-26s %s\n' $((i + 1)) "$host" "$( [ "$kind" = smb ] && echo "SMB shared folders" || echo "SFTP (over SSH)")"
  done
  read -rp "Which? " choice
  if ! [[ "$choice" =~ ^[0-9]+$ ]] || (( choice < 1 || choice > ${#found[@]} )); then
    echo "Choose a number from the list."
    return 1
  fi
  IFS='|' read -r kind host _ <<< "${found[$((choice - 1))]}"
  "add_$kind" "$host"
}

# smb://host/share or sftp://[user@]host[/path]
enter_and_add() {
  local address rest host user="" path
  read -rp "Address, e.g. smb://nas.local/files or sftp://me@server.local/srv/data: " address
  rest="${address#*://}"
  case "$address" in
    smb://*/?*) add_smb "${rest%%/*}" "${rest#*/}" ;;
    sftp://*)
      [[ "$rest" == *@* ]] && { user="${rest%%@*}"; rest="${rest#*@}"; }
      host="${rest%%/*}"
      path=""
      [[ "$rest" == */* ]] && path="/${rest#*/}"
      add_sftp "$host" "$user" "$path" ;;
    "") return 0 ;;
    *) echo "Start it with smb:// (and include the share) or sftp://."; return 1 ;;
  esac
}

choose_and_remove() {
  local entries choice i
  mapfile -t entries < <(grep -v '^$' "$STORAGE_CONF")
  for i in "${!entries[@]}"; do
    IFS='|' read -r _ source point <<< "${entries[$i]}"
    printf '  %2d) %-30s %s\n' $((i + 1)) "$point" "$source"
  done
  read -rp "Which? " choice
  if ! [[ "$choice" =~ ^[0-9]+$ ]] || (( choice < 1 || choice > ${#entries[@]} )); then
    echo "Choose a number from the list."
    return 1
  fi
  IFS='|' read -r -a entry <<< "${entries[$((choice - 1))]}"
  remove_storage "${entry[@]}"
}

run_configure() {
  local choice source point
  mkdir -p "$(dirname "$STORAGE_CONF")" && touch "$STORAGE_CONF" || return 1
  echo "This mounts shared folders from other machines on your network in /mnt: SMB shares"
  echo "(Windows, a NAS, a Mac with File Sharing on) and, over SSH, any machine you can ssh to."

  while true; do
    echo
    if [ -s "$STORAGE_CONF" ]; then
      while IFS='|' read -r _ source point; do
        echo "  $point  <- $source"
      done < "$STORAGE_CONF"
      echo
    fi
    echo "  1) Find shared folders on your network"
    echo "  2) Enter an address"
    [ -s "$STORAGE_CONF" ] && echo "  3) Remove one"
    echo "  q) Done"
    echo
    read -rp "Choose: " choice || return 0
    case "$choice" in
      1) find_and_add ;;
      2) enter_and_add ;;
      3) [ -s "$STORAGE_CONF" ] && choose_and_remove ;;
      *) return 0 ;;
    esac
  done
}