#!/bin/bash

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

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  set -o errexit
  set -o nounset
  set -o errtrace
  set -o pipefail
fi

source /usr/libexec/helper-scripts/wc-test.sh
source /usr/libexec/helper-scripts/live-mode.sh

overlay_tmpfs_repo='/var/lib/grub-live/overlay_tmpfs_repo'
target_overlay_mount_list=()
fs_type_whitelist=(
  'nfs'
  'proc'
  'tmpfs'
  'autofs'
  'binfmt_misc'
  'cgroup2'
  'configfs'
  'devpts'
  'devtmpfs'
  'fusectl'
  'hugetlbfs'
  'mqueue'
  'ramfs'
  'securityfs'
  'sysfs'
  'tracefs'
  'bpf'
  'portal'
  'debugfs'
  'overlay'
)

if ! [ -v kernel_cmdline ]; then
  if [ -f '/proc/cmdline' ]; then
    kernel_cmdline="$(cat -- /proc/cmdline)"
  elif [ -f '/proc/1/cmdline' ]; then
    kernel_cmdline="$(cat -- /proc/1/cmdline)"
  else
    kernel_cmdline=''
  fi
fi
if ! [ -v proc_mount_contents ]; then
  proc_mount_contents="$(cat -- /proc/self/mounts)"
fi
if ! [ -v lsblk_output ]; then
  lsblk_output="$(lsblk --raw --output=MOUNTPOINTS,RM | tail -n+2)"
fi

cmd_wrapper() {
  printf "%s\n" "$0: INFO: executing: $*"
  if "$@"; then
    printf "%s\n" "$0: INFO: Success."
    return 0
  else
    printf "%s\n" "$0: INFO: Non-zero exit code."
    return 1
  fi
}

check_in_live_mode() {
  ## We can't use systemd's ConditionKernelCommandLine here, since we want to
  ## run if any of several kernel parameters are set.
  ## ConditionKernelCommandLine only allows us to run if ALL listed kernel
  ## parameters are specified, which isn't expected to happen ever since some
  ## of the parameters are specific to Dracut and others are specific to
  ## initramfs-tools.
  ##
  ## We don't run live-hardener under ISO Live mode, it can cause installation
  ## issues and doesn't add any useful amount of security (aside from things
  ## like locking down the ability to modify UEFI variables, which is the very
  ## behavior that causes installation issues).

  printf '%s\n' "live_status_detected_live_mode_environment_machine=${live_status_detected_live_mode_environment_machine}" 1>&2

  if [ "${live_status_detected_live_mode_environment_machine}" = 'false' ] \
    || [[ "${live_status_detected_live_mode_environment_machine}" =~ ^iso-live ]]; then
    printf "%s\n" "$0: INFO: iso-live mode detected, exiting, ok."
    exit 0
  fi
}

get_mount_list_to_overlay() {
  local lsblk_raw_list lsblk_path_list \
    lsblk_removable_list lsblk_raw_str lsblk_raw_path_str \
    lsblk_raw_path_list lsblk_raw_path proc_mount_path_list \
    proc_mount_attr_list lsblk_idx lsblk_path lsblk_removable proc_mount_idx \
    fs_attr_list fs_attr_item is_fs_readonly is_removable is_fs_whitelisted \
    fs_whitelist_idx proc_mount_type proc_mount_path proc_mount_type_list

  if [ -z "${lsblk_output}" ]; then
    printf "%s\n" "$0: ERROR: empty output from 'lsblk --raw --output=MOUNTPOINTS,RM'!" 1>&2
    exit 1
  elif [ -z "${proc_mount_contents}" ]; then
    printf "%s\n" "$0: ERROR: empty output from 'cat -- /proc/self/mounts'!" 1>&2
    exit 1
  fi

  ## lsblk --raw output is a pain to parse unfortunately...
  readarray -t lsblk_raw_list <<< "${lsblk_output}"
  if (( ${#lsblk_raw_list[@]} <= 1 )) \
    || [ -z "${lsblk_raw_list[0]:-}" ]; then
    printf "%s\n" "$0: ERROR: 'lsblk_raw_list' array is empty!" 1>&2
    exit 1
  fi
  lsblk_path_list=()
  lsblk_removable_list=()
  for (( lsblk_idx = 0; lsblk_idx < ${#lsblk_raw_list[@]}; lsblk_idx++ )); do
    lsblk_raw_str="${lsblk_raw_list[lsblk_idx]}"

    if [ "${lsblk_raw_str:0:1}" = ' ' ]; then
      ## Current drive has no mountpoint, skip it
      continue
    fi

    lsblk_raw_path_str="$(cut -d' ' -f1 <<< "${lsblk_raw_str}")"
    ## Path field from lsblk contains escaped newlines, so this will convert
    ## them into real newlines (and fix any other escapes that happen to be
    ## in the string).
    printf -v lsblk_raw_path_str "%b" "${lsblk_raw_path_str}"
    readarray -t lsblk_raw_path_list <<< "${lsblk_raw_path_str}"
    lsblk_removable="$(cut -d' ' -f2 <<< "${lsblk_raw_str}")"

    for lsblk_raw_path in "${lsblk_raw_path_list[@]}"; do
      lsblk_path_list+=( "${lsblk_raw_path}" )
      lsblk_removable_list+=( "${lsblk_removable}" )
    done
  done

  ## TODO: /proc/self/mounts contents may have spaces and other encoded characters
  ## https://unix.stackexchange.com/questions/317476/shell-code-to-check-if-a-device-or-file-with-spaces-in-the-path-is-mounted
  readarray -t proc_mount_path_list < <(
    awk '{ print $2 }' <<< "${proc_mount_contents}"
  )
  readarray -t proc_mount_type_list < <(
    awk '{ print $3 }' <<< "${proc_mount_contents}"
  )
  readarray -t proc_mount_attr_list < <(
    awk '{ print $4 }' <<< "${proc_mount_contents}"
  )

  if (( ${#lsblk_path_list[@]} <= 1 )) \
    && [ -z "${lsblk_path_list[0]:-}" ]; then
    printf "%s\n" "$0: ERROR: 'lsblk_path_list' array is empty!" 1>&2
    exit 1
  elif (( ${#proc_mount_path_list[@]} <= 1 )) \
    && [ -z "${proc_mount_path_list[0]:-}" ]; then
    printf "%s\n" "$0: ERROR: 'proc_mount_path_list' array is empty!" 1>&2
    exit 1
  fi

  for (( proc_mount_idx = 0; proc_mount_idx < ${#proc_mount_path_list[@]};
    proc_mount_idx++ )); do
    proc_mount_type="${proc_mount_type_list[proc_mount_idx]}"
    is_fs_whitelisted='false'
    for (( fs_whitelist_idx = 0; \
      fs_whitelist_idx < ${#fs_type_whitelist[@]}; fs_whitelist_idx++ )); do
      if [ "${fs_type_whitelist[fs_whitelist_idx]}" = "${proc_mount_type}" ];
        then is_fs_whitelisted='true'
        break
      fi
    done
    if [ "${is_fs_whitelisted}" = 'true' ]; then
      continue
    fi

    proc_mount_path="${proc_mount_path_list[proc_mount_idx]}"
    is_removable='false'
    if [[ "${proc_mount_path}" =~ ^/media/ ]] \
      || [[ "${proc_mount_path}" =~ ^/mnt/ ]] \
      || [ "${proc_mount_path}" = '/media' ] \
      || [ "${proc_mount_path}" = '/mnt' ]; then
      for (( lsblk_idx = 0; lsblk_idx < ${#lsblk_path_list[@]}; \
        lsblk_idx++ )); do
        lsblk_path="${lsblk_path_list[lsblk_idx]}"
        if [ "${lsblk_path}" = "${proc_mount_path}" ]; then
          lsblk_removable="${lsblk_removable_list[lsblk_idx]}"
          if [ "${lsblk_removable}" = '1' ]; then
            is_removable='true'
            break
          fi
        fi
      done
      if [ "${is_removable}" = 'true' ]; then
        continue
      fi
    fi

    is_fs_readonly='true'
    IFS=',' read -r -a fs_attr_list \
      <<< "${proc_mount_attr_list[proc_mount_idx]}"
    for fs_attr_item in "${fs_attr_list[@]}"; do
      if [ "${fs_attr_item}" = 'rw' ]; then
        is_fs_readonly='false'
        break
      fi
    done

    if [ "${is_fs_readonly}" = 'false' ]; then
      target_overlay_mount_list+=( "${proc_mount_path}" );
    fi
  done
}

mount_overlays() {
  local mount_idx target_overlay_mount tmpfs_dir tmpfs_upper_dir \
    tmpfs_work_dir

  if [ "${#target_overlay_mount_list[@]}" = '0' ]; then
    printf "%s\n" "$0: INFO: If there are no directories to overlay, skip everything else, ok."
    return
  fi

  cmd_wrapper safe-rm -r -f -- "${overlay_tmpfs_repo}"
  cmd_wrapper mkdir --parents -- "${overlay_tmpfs_repo}"

  for (( mount_idx = 0; mount_idx < ${#target_overlay_mount_list[@]}; \
    mount_idx++ )); do
    target_overlay_mount="${target_overlay_mount_list[mount_idx]}"

    if ! cmd_wrapper mount -o remount,ro -- "${target_overlay_mount}"; then
      continue
    fi

    tmpfs_dir="${overlay_tmpfs_repo}/${mount_idx}"
    tmpfs_upper_dir="${tmpfs_dir}/upper"
    tmpfs_work_dir="${tmpfs_dir}/work"
    cmd_wrapper mkdir --parents -- "${tmpfs_dir}"
    cmd_wrapper mount -t tmpfs tmpfs -- "${tmpfs_dir}"
    cmd_wrapper mkdir --parents -- "${tmpfs_upper_dir}"
    cmd_wrapper mkdir --parents -- "${tmpfs_work_dir}"

    if ! cmd_wrapper \
      mount \
      -t overlay overlay \
      -o lowerdir="${target_overlay_mount}" \
      -o upperdir="${tmpfs_upper_dir}" \
      -o workdir="${tmpfs_work_dir}" \
      -- \
      "${target_overlay_mount}"; then
      ## Some filesystems don't support being the lower directory of an
      ## overlayfs (for instance, fat32). With these directories, we simply
      ## leave them read-only and clean up the directories that would have
      ## been used to place a writable overlay on them.
      cmd_wrapper umount -- "${tmpfs_dir}"
      cmd_wrapper safe-rm -r -f -- "${tmpfs_dir}"
    fi
  done
}

check_in_live_mode
get_mount_list_to_overlay
if [ "${LIVE_HARDENER_TEST:-}" = 'true' ]; then
  printf '%s\n' "${target_overlay_mount_list[@]}"
else
  mount_overlays
fi

## TODO: breaks unit test
#printf "%s\n" "$0: INFO: Success."
