#!/bin/sh
#
# SPDX-FileCopyrightText: 2017 - 2020 EvaDogStar <evastar@protonmail.com>
# SPDX-FileCopyrightText: 2024 - 2025 Benjamin Grande M. S. <ben.grande.b@gmail.com>
#
# SPDX-License-Identifier: MIT
#
# Take screenshot in Qubes GuiVM and copy/move to qube.
#
# Dialog tools: kdialog, zenity
# Shot tools: spectacle, xfce4-screenshooter, maim, scrot

set -eu

msg(){
  printf '%s\n' "$*" >&2
}

err_dialog(){
  msg="${1}"
  msg "[ERROR] ${msg}"
  case "${dialog_cmd}" in
    zenity) zenity "${dialog_title}" --error --text -- "${msg}";;
    kdialog) kdialog "${dialog_title}" --error "${msg}";;
    *) msg "Unsupported dialog command"; exit 1;;
  esac
  exit 1
}

unsupported_type_per_tool(){
  msg "${screenshot_cmd}: Unsupported screenshot type: ${screenshot_type}"
  exit 1
}

remove_screenshot(){
  rm -f -- "${screenshot_file}" ||
    err_dialog "Failed to remove file: '${screenshot_file}'"
}

take_screenshot() {
  screenshot_type="${1}"

  shift "${#}"
  set -- "${screenshot_file}"
  arg=""
  case "${screenshot_cmd}" in
    spectacle)
      set -- -o "${@}"
      case "${screenshot_type}" in
        region) arg="-r";;
        window) arg="-a";;
        fullscreen) arg="-f";;
        *) unsupported_type_per_tool;;
      esac
      ;;
    xfce4-screenshooter)
      set -- -s "${@}"
      case "${screenshot_type}" in
        region) arg="-r";;
        window) arg="-w";;
        fullscreen) arg="-f";;
        *) unsupported_type_per_tool;;
      esac
      ;;
    scrot)
      set -- -b -F "${@}"
      case "${screenshot_type}" in
        region|window) arg="-s";;
        fullscreen) arg="";;
        *) unsupported_type_per_tool;;
      esac
      ;;
    maim)
      set -- -o -u "${@}"
      case "${screenshot_type}" in
        region|window) arg="-s";;
        fullscreen) arg="";;
        *) unsupported_type_per_tool;;
      esac
      ;;
    *) msg "Unsupported screenshot tool"; exit 1;;
  esac
  if test -n "${arg}"; then
    set -- "${arg}" "${@}"
  fi
  "${screenshot_cmd}" "${@}"
}

print_help(){
  # editorconfig-checker-disable
  printf '%s\n' "Usage: ${0##*/} [OPTIONS]
 -h, --help                  print this help message and exit
Capture mode:
 -r, --region                capture a region
 -w, --window                capture a window
 -f, --fullscreen            capture everything
File outcome:
 -d, --qube NAME             qube to save screenshot
 --qube-file-manager         open file manager in qube
 --move                      move file instead of copy
Development mode:
 -D, --dialog-cmd NAME      dialog tool: kdialog, zenity
 -S, --screenshot-cmd NAME  screenshot tool: maim, scrot, spectacle,
                             xfce4-screenshooter
Note:
  maim and scrot:
    They do not have a separate option for region or window, therefore,
    selecting either of them will have the same effect, which is, capture a
    window by clicking on it, capture a region by dragging the mouse.
  xfce4-screenshooter:
    Window option can only capture the active window.
  spectacle:
    After screenshot is captured and edited, click on 'Save' and then close
    the window to continue the operation."
  # editorconfig-checker-enable
  exit 1
}

## Expand directory only in the qube.
qube_pictures_dir="\$(xdg-user-dir PICTURES)"
guivm_pictures_dir="$(xdg-user-dir PICTURES)"
mkdir -p -- "${guivm_pictures_dir}" || exit 1

dialog_title="--title=${0##*/}"
current_date="$(date +"%Y-%m-%d-%H%M%S")"
screenshot_basename="${current_date}.png"
screenshot_file="${guivm_pictures_dir%*/}/${screenshot_basename}"
qube_screenshot_file="${qube_pictures_dir}/${screenshot_basename}"
screenshot_type_text=""
screenshot_action_text=""
screenshot_action_supplied=""
qube=""
exit_required=0
file_manager=0
file_move=0
screenshot_cmd=""
screenshot_cmd_wanted=""
dialog_cmd=""
dialog_cmd_wanted=""

while test "$#" -gt 0; do
  key="${1}"
  case "${key}" in
    -h|--help)
      print_help
      ;;
    -r|--region)
      screenshot_type_text="Region"
      ;;
    -w|--window)
      screenshot_type_text="Window"
      ;;
    -f|--fullscreen)
      screenshot_type_text="Fullscreen"
      ;;
    -d|--qube)
      shift
      qube="${1}"
      ;;
    --qube-file-manager)
      file_manager=1
      screenshot_action_supplied="1"
      ;;
    --move)
      file_move=1
      screenshot_action_supplied="1"
      ;;
    -S|--screenshot-cmd)
      shift
      screenshot_cmd_wanted="${1}"
      ;;
    -D|--dialog-cmd)
      shift
      dialog_cmd_wanted="${1}"
      ;;
    *)
      msg "Unknown option: ${key}"
      exit 1
    ;;
  esac
  shift
done

if test -n "${dialog_cmd_wanted}"; then
  if ! command -v "${dialog_cmd_wanted}" >/dev/null; then
    msg="wanted dialog program not found: ${dialog_cmd_wanted}"
    msg "[ERROR] ${msg}"
    exit 1
  fi
  case "${dialog_cmd_wanted}" in
    kdialog|zenity);;
    *)
      msg="wanted dialog program unsupported: ${dialog_cmd_wanted}"
      msg "[ERROR] ${msg}"
      exit 1
      ;;
  esac
  dialog_cmd="${dialog_cmd_wanted}"
else
  if command -v kdialog >/dev/null; then
    dialog_cmd="kdialog"
  elif command -v zenity >/dev/null; then
    dialog_cmd="zenity"
  fi
  if test -z "${dialog_cmd}"; then
    msg "[ERROR] dialog programs not found: zenity kdialog"
    exit 1
  fi
fi

if test -n "${screenshot_cmd_wanted}"; then
  if ! command -v "${screenshot_cmd_wanted}" >/dev/null; then
    msg="wanted screenshot program not found: ${screenshot_cmd_wanted}"
    msg "[ERROR] ${msg}"
    case "${dialog_cmd}" in
      zenity) zenity "${dialog_title}" --info --text -- "${msg}";;
      kdialog) kdialog "${dialog_title}" --msgbox "${msg}";;
      *) msg "Unsupported dialog command"; exit 1;;
    esac
    exit 1
  fi
  case "${screenshot_cmd_wanted}" in
    maim|scrot|spectacle|xfce4-screenshooter);;
    *)
      msg="wanted screenshot program unsupported: ${screenshot_cmd_wanted}"
      msg "[ERROR] ${msg}"
      exit 1
      ;;
  esac
  screenshot_cmd="${screenshot_cmd_wanted}"
else
  if command -v maim >/dev/null; then
    screenshot_cmd="maim"
  elif command -v scrot >/dev/null; then
    screenshot_cmd="scrot"
  elif command -v spectacle >/dev/null; then
    screenshot_cmd="spectacle"
  elif command -v xfce4-screenshooter >/dev/null; then
    screenshot_cmd="xfce4-screenshooter"
  fi
  if test -z "${screenshot_cmd}"; then
    msg="screenshot programs not found"
    msg="${msg}: spectacle xfce4-screenshooter scrot maim"
    msg "[ERROR] ${msg}"
    case "${dialog_cmd}" in
      zenity) zenity "${dialog_title}" --info --text -- "${msg}";;
      kdialog) kdialog "${dialog_title}" --msgbox "${msg}";;
      *) msg "Unsupported dialog command"; exit 1;;
    esac
    exit 1
  fi
fi

if test -z "${screenshot_type_text}"; then
  # shellcheck disable=SC2086
  dialog_msg="Select capture mode:"
  case "${dialog_cmd}" in
    zenity)
      screenshot_type_text="$(zenity "${dialog_title}" --list \
        --text "${dialog_msg}" \
        --radiolist \
        --column "Pick" --column "Mode" -- \
        FALSE "Region" \
        FALSE "Window" \
        FALSE "Fullscreen" \
      )"
      ;;
    kdialog)
      screenshot_type_text="$(kdialog "${dialog_title}" \
        --radiolist "${dialog_msg}" -- \
        "Region" "Region" off \
        "Window" "Window" off \
        "Fullscreen" "Fullscreen" off \
      )"
      ;;
    *) msg "Unsupported dialog command"; exit 1;;
  esac
fi

case "${screenshot_type_text}" in
  "Region") take_screenshot region;;
  "Window") take_screenshot window;;
  "Fullscreen") take_screenshot fullscreen;;
  *) msg "[ERROR] mode not selected"; exit 1;;
esac

if ! test -f "${screenshot_file}"; then
  msg="Screenshot was not saved in GuiVM"
  msg "[ERROR] ${msg}"
  case "${dialog_cmd}" in
    zenity) zenity "${dialog_title}" --warning --text -- "${msg}";;
    kdialog) kdialog "${dialog_title}" --sorry "${msg}";;
    *) msg "Unsupported dialog command"; exit 1;;
  esac
  exit 1
fi

if test "${screenshot_action_supplied}" != "1"; then
  dialog_msg="${screenshot_type_text} capture saved: ${screenshot_file}."
  dialog_msg="${dialog_msg}\nWhat do you want to do with it?"
  case "${dialog_cmd}" in
    zenity)
      screenshot_action_text="$(zenity "${dialog_title}" \
        --list \
        --width=280 --height=210 \
        --text "${dialog_msg}" \
        --separator="\n" \
        --checklist --column "Pick" --column "Resolution" -- \
        FALSE "Exit" \
        FALSE "Open file manager in qube" \
        FALSE "Move file"
      )"
      ;;
    kdialog)
      screenshot_action_text="$(kdialog "${dialog_title}" \
        --checklist "${dialog_msg}" \
        --separate-output -- \
        "Exit" "Exit" off \
        "Open file manager in qube" "Open file manager in qube" off \
        "Move file" "Move file" off
      )"
      ;;
    *) msg "Unsupported dialog command"; exit 1;;
  esac

  if test -z "${screenshot_action_text}"; then
    exit 0
  fi

  IFSOLD="${IFS}"
  IFS="|"
  screenshot_action_text="$(printf '%s\n' "${screenshot_action_text}" | \
    tr "\n" "|")"
  for val in ${screenshot_action_text}; do
    case "${val}" in
      "Exit") exit_required=1;;
      "Open file manager in qube") file_manager=1;;
      "Move file") file_move=1;;
      *) exit 1;;
    esac
  done
  IFS="${IFSOLD}"
fi

if test "${exit_required}" = "1"; then
  exit 0
fi

qube_list="$(qvm-ls --no-spinner --raw-data --fields=NAME,CLASS | \
              awk -F "|" '$1 !~ /(^dvm-|-dvm$)/ &&
                          $2 !~ /^(AdminVM|TemplateVM)$/{print $1}')"

if test -z "${qube}"; then
  dialog_msg="Select destination qube"
  case "${dialog_cmd}" in
    zenity)
      qube_list="$(printf '%s\n' "${qube_list}" | sed -e "s/^/FALSE /")"
      # shellcheck disable=SC2086
      qube="$(zenity "${dialog_title}" --list  --width=200 --height=390 \
              --text "${dialog_msg}" \
              --radiolist --column "Pick" --column "qube" -- ${qube_list})"
      ;;
    kdialog)
      qube_list="$(printf '%s\n' "${qube_list}" | \
        sed -e "s/\(.*\)/\1 \1 off/")"
      # shellcheck disable=SC2086
      qube="$(kdialog "${dialog_title}" --radiolist "${dialog_msg}" \
              -- ${qube_list})"
      ;;
    *) msg "Unsupported dialog command"; exit 1;;
  esac
  if test -z "${qube}"; then
    msg="qube was not selected"
    msg "[ERROR] ${msg}"
    case "${dialog_cmd}" in
      zenity) zenity "${dialog_title}" --error --text -- "${msg}";;
      kdialog) kdialog "${dialog_title}" --error "${msg}";;
      *) msg "Unsupported dialog command"; exit 1;;
    esac
    exit 1
  fi
fi

if ! qvm-check -- "${qube}" >/dev/null 2>&1; then
  msg="qube doesn't exist: ${qube}"
  msg "[ERROR] ${msg}"
  case "${dialog_cmd}" in
    zenity) zenity "${dialog_title}" --error --text -- "${msg}";;
    kdialog) kdialog "${dialog_title}" --error "${msg}";;
    *) msg "Unsupported dialog command"; exit 1;;
  esac
  exit 1
fi

qube_qrexec="$(python3 -c "import qubesadmin
dom = qubesadmin.Qubes().domains[\"${qube}\"]
print(dom.features.check_with_template('qrexec'))")"

## Unix system.
if test "${qube_qrexec}" = "1" &&
  qvm-run --no-gui --pass-io -- "${qube}" true >/dev/null 2>&1
then
  qvm-run --no-gui -- "${qube}" "mkdir -p -- \"${qube_pictures_dir}\"" ||
    err_dialog "Failed to create directory: ${qube}: '${qube_pictures_dir}'"
  qvm-run --no-gui --pass-io -- "${qube}" \
    "cat -- > \"${qube_screenshot_file}\"" < "${screenshot_file}" ||
    err_dialog "Failed to copy screenshot to qube: '${qube}'"

  if test "${file_move}" = "1"; then
    remove_screenshot
  fi

  if test "${file_manager}" = "1"; then
    qube_gui="$(qvm-features -- "${qube}" gui || true)"
    qube_gui_emulated="$(qvm-features -- "${qube}" gui-emulated || true)"
    if test -z "${qube_gui}" && test -z "${qube_gui_emulated}"; then
      err_dialog "Refusing to open file manager in qube with disabled GUI"
    fi
    file_manager_cmd="xdg-open \"${qube_pictures_dir}\""
    qvm-run -- "${qube}" "${file_manager_cmd}" ||
      err_dialog "Failed to open file manager: ${qube}: '${file_manager_cmd}'"
  fi
  exit
fi

## Any other system that supports Qrexec.
if test "${qube_qrexec}" = "1" &&
  qvm-copy-to-vm "${qube}" "${screenshot_file}" >/dev/null 2>&1
then
  if test "${file_move}" = "1"; then
    remove_screenshot
  fi
  exit
fi

err_dialog "Failed to copy screenshot to qube: '${qube}'"