#!/bin/bash

## Copyright (c) 2025 - 2025 ENCRYPTED SUPPORT LLC <adrelanos@whonix.org>
## See the file COPYING for copying conditions.

set -o errexit
set -o nounset
set -o errtrace
set -o pipefail

# shellcheck source=../libexec/helper-scripts/get_colors.sh
source /usr/libexec/helper-scripts/get_colors.sh
# shellcheck source=../libexec/helper-scripts/accountctl.sh
source /usr/libexec/helper-scripts/accountctl.sh

get_config_section() {
   local file_name section found_section line
   file_name="$1"
   section="$2"
   found_section='n'
   while read -r -- line; do
      if [ "${found_section}" = 'y' ]; then
         if [[ "${line}" =~ ^\[.*\]$ ]]; then
            break
         fi
         printf '%s\n' "${line}"
      fi
      if [ "${line}" = "[${section}]" ]; then
         found_section='y'
      fi
   done < "${file_name}"

   true "INFO: Extracted config section '$section' from file '$file_name'."
}

remove_from_config_section() {
   local file_name section remove_entity do_literal_remove filter_stage line \
     lines output line_was_removed
   file_name="$1"
   section="$2"
   remove_entity="$3"
   do_literal_remove="$4"
   filter_stage='1'
   line_was_removed=""

   mapfile -t lines < "$file_name"

   output=()
   for line in "${lines[@]}"; do
      if [ "${filter_stage}" = '1' ]; then
         output+=( "${line}" )
         if [ "${line}" = "[${section}]" ]; then
            filter_stage='2'
         fi
      elif [ "${filter_stage}" = '2' ]; then
         if [ "${do_literal_remove}" = 'true' ]; then
            ## 'remove_entity' is treated as a string literal here.
            if [[ "${line}" =~ ^"${remove_entity}"$ ]]; then
               line_was_removed="true"
               continue
            fi
         else
            ## 'remove_entity' is treated as a regex here.
            if [[ "${line}" =~ ^${remove_entity}$ ]]; then
               line_was_removed="true"
               continue
            fi
         fi
         output+=( "${line}" )
         if [[ "${line}" =~ ^\[.*\]$ ]]; then
            filter_stage='3'
         fi
      else
         output+=( "${line}" )
      fi
   done

   printf '%s\n' "${output[@]}" | sponge -- "$file_name"

   if [ "$line_was_removed" = "true" ]; then
      printf '%s\n' "INFO: Removed lines matching '$remove_entity' from section '$section' of file '$file_name'." >&2
      return 0
   fi
   return 1
}

enable_autologin() {
   local user
   user="$1"

   mkdir --parents -- /etc/lightdm/lightdm.conf.d
   printf '%s\n' "[Seat:*]
autologin-user=${user}" | sponge -- '/etc/lightdm/lightdm.conf.d/40_autologin.conf'
   printf '%s\n' "INFO: Wrote file '/etc/lightdm/lightdm.conf.d/40_autologin.conf' to enable autologin for account '$user'." >&2
   printf '%s\n' "$delimiter" >&2
   cat -- '/etc/lightdm/lightdm.conf.d/40_autologin.conf' >&2
   printf '%s\n' "$delimiter" >&2

   mkdir --parents -- /etc/sddm.conf.d
   printf '%s\n' "[Autologin]
User=${user}" | sponge -- '/etc/sddm.conf.d/z-autologin.conf'
   printf '%s\n' "INFO: Wrote file '/etc/sddm.conf.d/z-autologin.conf' to enable autologin for account '$user'." >&2
   printf '%s\n' "$delimiter" >&2
   cat -- '/etc/sddm.conf.d/z-autologin.conf' >&2
   printf '%s\n' "$delimiter" >&2

   mkdir --parents -- /etc/greetd/config.toml.d
   printf '%s\n' "[initial_session]
command = \"/usr/libexec/desktop-config-dist/start-lxqt-session\"
user = \"${user}\"" | sponge -- '/etc/greetd/config.toml.d/40_autologin.conf'
   printf '%s\n' "INFO: Wrote file '/etc/greetd/config.toml.d/40_autologin.conf' to enable autologin for account '$user'." >&2
   printf '%s\n' "$delimiter" >&2
   cat -- '/etc/greetd/config.toml.d/40_autologin.conf'
   printf '%s\n' "$delimiter" >&2
   printf '%s\n' "INFO: Rebuilding greetd config."
   build-config-file /etc/greetd/config.toml.d /etc/greetd/config.toml
}

disable_autologin() {
   local user lightdm_config_file_list file_name file_contents pattern match_count
   user="$1"

   ## TODO: All of these routines may misbehave if the spacing around the =
   ## sign in the respective configuration lines is different than what the
   ## script assumes, or if one of the lines has an otherwise unexpected
   ## value.

   readarray -t lightdm_config_file_list < <(find /etc/lightdm)
   for file_name in "${lightdm_config_file_list[@]}"; do
      if ! [ -f "${file_name}" ]; then
         continue
      fi

      pattern="^autologin-user=[[:space:]]*${user}$"
      match_count=$(grep --extended-regexp --count -- "$pattern" "${file_name}") || true

      sed -i -E "/$pattern/d" -- "${file_name}"

      if [ "$match_count" -gt 0 ]; then
         printf '%s\n' "INFO: Removed $match_count matching line(s) from file '${file_name}'." >&2
      else
         true "INFO: No matching lines found in file '${file_name}'."
      fi

      #sed -i "/^autologin-user=\s*${user}$/d" -- "${file_name}"
      true "INFO: Processed autologin configuration removal check for account '$user' from file '$file_name'."
      file_contents="$(cat -- "${file_name}")"
      if [ -z "${file_contents}" ] \
         || [ "${file_contents}" = '[Seat:*]' ]; then
         printf '%s\n' "INFO: File '$file_name' contains no meaningful configuration data, removing." >&2
         safe-rm --verbose -- "${file_name}"
      fi
   done

   for file_name in /etc/sddm.conf /etc/sddm.conf.d/*; do
      if ! [ -f "${file_name}" ]; then
         continue
      fi
      remove_from_config_section "${file_name}" 'Autologin' "User=${user}" 'true' || true
      true "INFO: Processed autologin configuration removal check for account '$user' from file '$file_name'."
      file_contents="$(cat -- "${file_name}")"
      if [ -z "${file_contents}" ] \
         || [ "${file_contents}" = '[Autologin]' ]; then
         printf '%s\n' "INFO: File '$file_name' contains no meaningful configuration data, removing." >&2
         safe-rm --verbose -- "${file_name}"
      fi
   done

   for file_name in /etc/greetd/config.toml.d/*; do
      if ! [ -f "${file_name}" ]; then
         continue
      fi
      if remove_from_config_section "${file_name}" 'initial_session' "user = \"${user}\"" 'true'; then
         remove_from_config_section "${file_name}" 'initial_session' "command = \".*\"" 'false' || true
      fi
      true "INFO: Processed autologin configuration removal check for account '$user' from file '$file_name'."
      file_contents="$(cat -- "${file_name}")"
      if [ -z "${file_contents}" ] \
         || [ "${file_contents}" = '[initial_session]' ]; then
         printf '%s\n' "INFO: File '$file_name' contains no meaningful configuration data, removing." >&2
         safe-rm --verbose -- "${file_name}"
      fi
   done
   printf '%s\n' "INFO: Rebuilding greetd config."
   build-config-file /etc/greetd/config.toml.d /etc/greetd/config.toml
}

enable_sysmaint_autologin() {
   mkdir --parents -- /etc/user-sysmaint-split.conf.d
   printf '%s\n' 'sysmaint-autologin=yes' | sponge -a -- /etc/user-sysmaint-split.conf.d/40_autologin.conf
   printf '%s\n' "INFO: Wrote autologin configuration for account 'sysmaint' to file '/etc/user-sysmaint-split.conf.d/40_autologin.conf'." >&2
   printf '%s\n' "$delimiter" >&2
   cat -- /etc/user-sysmaint-split.conf.d/40_autologin.conf >&2
   printf '%s\n' "$delimiter" >&2
}

disable_sysmaint_autologin() {
   local file_name file_contents
   for file_name in /etc/user-sysmaint-split.conf.d/*; do
      if ! [ -f "${file_name}" ]; then
         continue
      fi
      sed -i '/^sysmaint-autologin=/d' -- "${file_name}"
      ## TODO: Report only actually removed instances similar to disable_autologin function.
      printf '%s\n' "INFO: Removed autologin configuration for account 'sysmaint' from file '${file_name}'." >&2
      file_contents="$(cat -- "${file_name}")"
      if [ -z "${file_contents}" ]; then
         printf '%s\n' "INFO: File '$file_name' contains no meaningful configuration data, removing." >&2
         safe-rm --verbose -- "${file_name}"
      fi
   done
}

validate_username() {
   local user found_user user_entry

   user="${1:-}"
   if [ -z "${user}" ]; then
      printf '%s\n' "ERROR: No username provided. Please specify a username." >&2
      exit 1
   fi

   if ! id "${user}" &>/dev/null ; then
      printf '%s\n' "ERROR: Account '$user' does not exist. Please check the username and try again." >&2
      exit 1
   fi

   found_user='n'
   for user_entry in "${user_list[@]}"; do
      if [ "${user}" = "${user_entry}" ]; then
         found_user='y'
         break
      fi
   done

   if [ "${found_user}" = 'n' ]; then
      printf '%s\n' "ERROR: Account '$user' is not a normal user account. Please check the username and try again." >&2
      exit 1
   fi
}

filter_out_sysmaint_maybe() {
   ## This function expects to have data piped to in on standard input.
   if [ "${special_sysmaint_handling}" = 'yes' ]; then
      grep --invert-match -- 'sysmaint'
   else
      cat
   fi
}

check_autologin() {
   local user
   user="${1:-}"
   ## Don't run check_root, this might need to be used as a normal user by
   ## systemcheck

   ## NOTE: Exit code is important here because pwchange uses
   ##       'autologinchange -c "$user"' and relies on the exit code.
   # shellcheck disable=SC2076
   if [[ " ${autologin_users[*]} " =~ " ${user} " ]]; then
      exit 0
   else
      exit 1
   fi
}

cli_enable_autologin() {
   local user

   check_root
   create_dm_config_dirs

   user="${1:-}"
   validate_username "${user}"

   if [ "${user}" = 'sysmaint' ] && [ "${special_sysmaint_handling}" = 'yes' ]; then
      ## Account 'sysmaint' requires special handling.
      if [ "${sysmaint_autologin}" = 'yes' ]; then
         printf '%s\n' "INFO: Autologin already enabled for account 'sysmaint', exiting." >&2
         exit 0
      fi
      enable_sysmaint_autologin
      exit 0
   fi

   # shellcheck disable=SC2076
   if [[ " ${autologin_users[*]} " =~ " ${user} " ]]; then
      printf '%s\n' "INFO: Autologin already enabled for account '$user', exiting." >&2
      exit 0
   fi

   enable_autologin "${user}"
}

cli_disable_autologin() {
   local user

   check_root
   create_dm_config_dirs

   user="${1:-}"
   validate_username "${user}"

   if [ "${user}" = 'sysmaint' ] && [ "${special_sysmaint_handling}" = 'yes' ]; then
      ## Account 'sysmaint' requires special handling.
      if [ "${sysmaint_autologin}" = 'no' ]; then
         printf '%s\n' "INFO: Autologin already disabled for account 'sysmaint', exiting." >&2
         exit 0
       fi
       disable_sysmaint_autologin
       exit 0
   fi

   # shellcheck disable=SC2076
   if ! [[ " ${autologin_users[*]} " =~ " ${user} " ]]; then
      printf '%s\n' "INFO: Autologin already disabled for account '$user', exiting." >&2
      exit 0
   fi

   disable_autologin "${user}"
}

warn_on_empty_user_password() {
   if [ -z "$(get_clean_pass "$1")" ]; then
      printf '%s\n' "${red}WARNING:${nocolor} Account '$1' has no password set." >&2
      printf '%s\n' "Users can log into this account knowing only the username." >&2
      printf '%s\n' "You can use the 'pwchange' utility to change this." >&2
      printf '%s\n' "See https://www.kicksecure.com/wiki/Login for more information." >&2
   fi
}

autologinchange() {
   local user user_entry autologin_user enable_yn disable_yn

   check_root
   create_dm_config_dirs

   printf '%s\n' "[${green}INFO${nocolor}]: This tool enables or disables GUI autologin only." >&2
   printf '%s\n' "Users present on the system:" >&2
   for user_entry in "${user_list[@]}"; do
      printf '%s\n' "   ${user_entry}" >&2
   done

   read -r -p "Enter the username you would like to toggle autologin for: " -- user

   validate_username "${user}"

   # shellcheck disable=SC2076
   if [[ " ${autologin_users[*]} " =~ " ${user} " ]]; then
      printf '%s\n' "Account '$user' is currently configured to automatically log in." >&2
      read -r -p "Would you like to disable autologin? [Y/N] " -- disable_yn

      if [ "${disable_yn,,}" = 'y' ]; then
         disable_autologin "${user}"
         printf '%s\n' "SUCCESS: Autologin for account '$user' disabled." >&2
         warn_on_empty_user_password "$user"
      else
         printf '%s\n' "CANCELLED disabling autologin for account '$user'." >&2
      fi
      exit 0
   elif [ "${user}" = 'sysmaint' ] && [ "${special_sysmaint_handling}" = 'yes' ]; then
      if [ "${sysmaint_autologin}" = 'yes' ]; then
         printf '%s\n' "Account 'sysmaint' is currently configured to automatically log in." >&2
         read -r -p "Would you like to disable autologin for sysmaint session? [Y/N] " -- disable_yn

         if [ "${disable_yn,,}" = 'y' ]; then
            disable_sysmaint_autologin
            printf '%s\n' "SUCCESS: Autologin for account 'sysmaint' disabled." >&2
            warn_on_empty_user_password 'sysmaint'
         else
            printf '%s\n' "CANCELLED disabling autologin for account 'sysmaint'." >&2
         fi
         exit 0
      else
         printf '%s\n' "Account 'sysmaint' is currently NOT configured to automatically log in." >&2
         read -r -p "Would you like to enable autologin for sysmaint session? [Y/N] " -- enable_yn

         if [ "${enable_yn,,}" = 'y' ]; then
            ## Disable sysmaint autologin first, then re-enable it, to make
            ## sure other configuration lying around doesn't override the
            ## user's choice.
            disable_sysmaint_autologin
            enable_sysmaint_autologin
            printf '%s\n' "SUCCESS: Autologin for account 'sysmaint' enabled." >&2
         else
            printf '%s\n' "CANCELLED enabling autologin for account 'sysmaint'." >&2
         fi
         exit 0
      fi
   else
      printf '%s\n' "Account '$user' is currently NOT configured to automatically log in." >&2
      read -r -p "Are you sure you want to enable autologin for account '$user'? [Y/N] " -- enable_yn

      if [ "${enable_yn,,}" = 'y' ]; then
         if [ "${#autologin_users[@]}" != '0' ]; then
            for autologin_user in "${autologin_users[@]}"; do
               disable_autologin "${autologin_user}"
            done
         fi
         enable_autologin "${user}"
         printf '%s\n' "SUCCESS: Autologin for account '$user' enabled." >&2
      else
         printf '%s\n' "CANCELLED enabling autologin for account '$user'." >&2
      fi
      exit 0
   fi
}

check_root() {
   if [ "$(id -u)" != "0" ]; then
      printf '%s\n' "ERROR: This must be run as root (sudo)!" >&2
      exit 1
   fi
}

create_dm_config_dirs() {
   mkdir --parents -- /etc/lightdm
   mkdir --parents -- /etc/sddm.conf.d
}

if [ -f '/usr/share/qubes/marker-vm' ]; then
   printf '%s\n' "INFO: GUI autologin is not applicable to Qubes OS." >&2
   exit 0
fi

delimiter="####################"
autologin_users_lightdm=()
autologin_users_sddm=()
autologin_users=()

special_sysmaint_handling='no'
sysmaint_autologin='no'
if [ -x /usr/libexec/user-sysmaint-split/sysmaint-boot ]; then
   special_sysmaint_handling='yes'
   if [ "$(/usr/libexec/user-sysmaint-split/sysmaint-boot query-sysmaint-autologin 2>/dev/null)" = 'yes' ]; then
      sysmaint_autologin='yes'
   fi
fi

autologinchange_mode=''
while true; do
   case "${1:-}" in
      '-c')
         shift
         autologinchange_mode='check_autologin'
         ;;
      '-e')
         shift
         autologinchange_mode='cli_enable_autologin'
         ;;
      '-d')
         shift
         autologinchange_mode='cli_disable_autologin'
         ;;
      '-r')
         ## -r = raw, don't use special handling on sysmaint account even if
         ## user-sysmaint-split is installed
         shift
         special_sysmaint_handling='no'
         ;;
      '--')
         shift
         break
         ;;
      *)
         break
         ;;
   esac
done
if [ -z "${autologinchange_mode}" ]; then
   autologinchange_mode='autologinchange'
fi

## Ensure that lightdm doesn't have a multi-seat configuration, we can't
## safely manage those
## TODO: consider cases with leading spaces
## TODO: ignore comments
if [ -d /etc/lightdm ]; then
   if grep --quiet --recursive --ignore-case -- '^\[Seat:[^*]' /etc/lightdm; then
      printf '%s\n' "ERROR: Multi-seat lightdm configuration detected, cannot proceed!" >&2
      exit 1
   fi
fi

autologin_users_lightdm=()
if [ -d /etc/lightdm ]; then
   readarray -t autologin_users_lightdm < <(
      grep --recursive --ignore-case -- '^autologin-user=' /etc/lightdm \
         | filter_out_sysmaint_maybe \
         | awk -F'=' '{ print $NF }' || true
   ) || true
fi
true "INFO: Collected list of users with lightdm autologin enabled."

readarray -t autologin_users_sddm < <(
   for file_name in /etc/sddm.conf /etc/sddm.conf.d/*; do
      if ! [ -f "${file_name}" ]; then
         continue
      fi
      get_config_section "${file_name}" 'Autologin'  \
         | grep -- '^User=' \
         | filter_out_sysmaint_maybe \
         | awk -F'=' '{ print $NF }' || true
   done
) || true
true "INFO: Collected list of users with sddm autologin enabled."

readarray -t autologin_users_greetd < <(
   for file_name in /etc/greetd/config.toml.d/*; do
      if ! [ -f "${file_name}" ]; then
         continue
      fi
      get_config_section "${file_name}" 'initial_session' \
         | grep -- '^user\s*=\s*"[^"]*"$' \
         | filter_out_sysmaint_maybe \
         | sed 's/^.*"\([^"]*\)"$/\1/' || true
   done
) || true

readarray -t autologin_users < <(
  printf '%s\n' \
     "${autologin_users_lightdm[@]}" \
     "${autologin_users_sddm[@]}" \
     "${autologin_users_greetd[@]}" \
     | sed '/^$/d' \
     | sort -u
)

if ! user_list_str="$(/usr/libexec/helper-scripts/get-user-list)"; then
    printf '%s\n' "ERROR: Failed to get user list!" >&2
    exit 1
fi
readarray -t user_list <<< "$user_list_str"

"${autologinchange_mode}" "$@"
