#!/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

## sets the following variables:
## live_status_detected_live_mode_environment_pretty
## live_status_detected_live_mode_environment_machine
## live_status_word_pretty
## live_status_detected
## live_status_maybe_iso_live_message
source /usr/libexec/helper-scripts/live-mode.sh

## The location where tmpfs directories are created. These are used as the upper
## dirs for overlayfs mounts later.
overlay_tmpfs_repo='/run/grub-live/overlay_tmpfs_repo'

## Directories that should be remounted read-only and (if possible) overlaid
## with a tmpfs.
target_overlay_mount_list=()

## Mountpoints on the system that have other mountpoints beneath them. We
## usually (but not always) want to avoid mounting overlays on these mounts,
## because OverlayFS will hide all the submounts.
dir_submount_list=()

## 'true' and 'false' values, one for each mount in target_overlay_mount_list.
## If a mount has a 'true' value in this list, it indicates that the mount
## should only be remounted read-only, but not overlaid.
skip_dir_overlay_bool_list=()

## Lists of mount paths, filesystem types, and mount attributes from
## /proc/self/mounts. The values in these lists are sorted, see
## populate_proc_mount_lists for details of the sort order.
proc_mount_path_list=()
proc_mount_type_list=()
proc_mount_attr_list=()

## Filesystems we do not want to remount read-only and overlay. Network
## filesystems and special filesystems belong here.
fs_type_whitelist=(
  'nfs'
  'proc'
  'tmpfs'
  'autofs'
  'binfmt_misc'
  'cgroup2'
  'configfs'
  'devpts'
  'devtmpfs'
  'fusectl'
  'hugetlbfs'
  'mqueue'
  'ramfs'
  'securityfs'
  'sysfs'
  'tracefs'
  'bpf'
  'portal'
  'debugfs'
  'overlay'
)

## Filesystems we want to remount read-only but not overlay. Filesystems that
## OverlayFS considers "weird" should be listed here, see
## https://lkml.kernel.org/linux-ext4/CAOQ4uxgPXBazE-g2v=T_vOvnr_f0ZHyKYZ4wvn7A3ePatZrhnQ@mail.gmail.com/T/
fs_type_nooverlay_list=(
  'cifs'
  'efivarfs'
  'iso9660'
  'jfs'
  'vfat'
  'overlay'
)

## Mounts we want to overlay even if it requires hiding submounts.
allow_hiding_submounts_list=(
  ## /boot will usually have /boot/efi as a submount. It is very unlikely
  ## anyone will need to modify /boot/efi in live mode, so we're fine with
  ## hiding it.
  '/boot'
)

## Mounts we don't even want to remount read-only, if they have submounts.
disallow_readonly_with_submounts_list=(
  ## /var is always expected to be writable. If it has no submounts, we can
  ## make it read-only and overlay it, but if it does have submounts, we can't
  ## overlay it without hiding those submounts. We can't remount it read-only
  ## either in such a situation.
  '/var'
)

## Mounts we don't want to remount read-only under any circumstance.
disallow_readonly_list=(
  ## /run/rootfsbase has a read/write tmpfs overlay on it already when booted
  ## in live mode. It *also* has several other hidden mounts on it, so the
  ## check for the 'overlay' filesystem isn't enough to prevent the script
  ## from trying to remount it read-only, so we need to ignore it explicitly.
  '/run/rootfsbase'
)

## The ASCII control codes SOH and STX. These are used for adjusting string
## sort order.
ascii_soh="$(printf '%b' '\001')"
ascii_stx="$(printf '%b' '\002')"

## The contents of /proc/self/mounts. This variable may be set by the
## regression test script, so we only populate it with real data if we need
## to.
if ! [ -v proc_mount_contents ]; then
  proc_mount_contents="$(cat -- /proc/self/mounts)"
fi
## A string containing a list of mountpoints and whether each one is located
## on a removable disk or not. lsblk --raw output provides info about one disk
## on each line. The two fields we ask lsblk to return are separated by a
## space; the first field is the path(s) the disk is mounted to, the second
## field is 1 if the disk is removable and 0 otherwise. The mount data uses
## octal escape codes for spaces, newlines, etc.
if ! [ -v lsblk_output ]; then
  lsblk_output="$(lsblk --raw --output=MOUNTPOINTS,RM | tail -n+2)"
fi

## Wrapper to run a command and report success or failure.
cmd_wrapper() {
  printf "%s\n" "$0: INFO: executing: $*"
  if "$@"; then
    printf "%s\n" "$0: INFO: Success."
    return 0
  else
    printf "%s\n" "$0: ERROR: Non-zero exit code."
    return 1
  fi
}

## Checks if the system is booted in live mode or not.
check_in_live_mode() {
  ## Better to use the unified live mode testing code than to rely on
  ## systemd's ConditionKernelCommandLine here.
  ##
  ## 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
}

## Populates proc_mount_path_list, proc_mount_type_list, and
## proc_mount_attr_list.
populate_proc_mount_lists() {
  local proc_mount_annotated_str line proc_mount_path

  ## In order for submount detection to work, we need to have the proc_mount
  ## lists specify their items such that all paths under one directory are
  ## grouped together, and all sibling directories are sorted alphabetically.
  ## This allows us to determine if a path contains submounts or not, by
  ## detecting if the path after it starts with the original path name plus a
  ## trailing slash. I.e., given the following list of mounts:
  ## - /etc
  ## - /var
  ## - /var/log
  ## - /vara
  ## - /home
  ## We can determine that /var contains submounts, because the path after it
  ## starts with '/var/'.
  ##
  ## This sort order is equivalent to an alphabetic sort, except for the
  ## forward slash is given the absolute highest priority. The easiest way to
  ## adjust the sort order this way is to replace the forward slash with some
  ## character that sorts higher than anything else. The ASCII control code
  ## STX is a good candidate for this, because it is extremely unlikely to
  ## occur in a real path, and it is the character with the second-highest
  ## sorting order in the C locale.
  ##
  ## We want to sort the entire contents of proc_mount_contents by the second
  ## field, so in order to accomplish this, we use a Decorate-Sort-Undecorate
  ## routine. The decorator is the second field from proc_mount_contents, with
  ## all slashes replaced with the SOH control code. (If we wanted to cope
  ## with the possibility of a SOH character in a path, we could replace SOH
  ## characters with forward slashes, but this is probably unnecessary.)
  ##
  ## We need the separator between the decorator and the rest of the mount
  ## information to sort higher even than the slash. SOH is a good choice for
  ## this. (Even if SOH characters appear in the path, this won't cause a
  ## problem, because only the first SOH in each decorated line acts as a
  ## separator.)

  ## Decorate the lines from proc_mount_contents.
  proc_mount_annotated_str=''
  while read -r line; do
    ## 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
    proc_mount_path="$(cut -d' ' -f2 <<< "${line}")"
    proc_mount_path="${proc_mount_path//\//"${ascii_stx}"}"
    proc_mount_annotated_str+="${proc_mount_path}${ascii_soh}${line}"$'\n'
  done <<< "${proc_mount_contents}"

  ## Sort and undecorate the lines.
  proc_mount_annotated_str="$(LC_ALL='C' sort \
    <<< "${proc_mount_annotated_str}")"
  proc_mount_contents="$(cut -d"${ascii_soh}" -f2- \
    <<< "${proc_mount_annotated_str}")"

  ## Split the sorted string into the proc_mount lists.
  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 (( ${#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
}

## Finds all mounts on the system that should be remounted read-only and
## overlaid.
get_mount_list_to_harden() {
  local lsblk_raw_list lsblk_path_list overlay_mount_list_str \
    lsblk_removable_list lsblk_raw_str lsblk_raw_path_str \
    lsblk_raw_path_list lsblk_raw_path lsblk_idx lsblk_path lsblk_removable \
    proc_mount_idx fs_attr_list fs_attr_item skip_fs_remount_ro is_removable \
    is_fs_whitelisted fs_whitelist_idx proc_mount_type proc_mount_path \
    disallow_readonly_dir submount_dir mount_idx target_overlay_mount \
    skip_dir_overlay allow_hide_submounts submount_dir allow_hide_submounts \
    fs_nooverlay_item allow_hide_submounts_dir target_dir_idx target_dir

  ## Ensure lsblk_output and proc_mount_contents are set
  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

  ## Expand the raw output from lsblk into two correlated lists, one
  ## specifying mount paths, the other specifying whether each mount path is
  ## located on a removable disk.
  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 in "${!lsblk_raw_list[@]}"; 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

  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
  fi

  ## Build the list of mounts to harden from the list of mounts on the system.
  for proc_mount_idx in "${!proc_mount_path_list[@]}"; do
    ## Skip mounts for whitelisted filesystems.
    proc_mount_type="${proc_mount_type_list[proc_mount_idx]}"
    is_fs_whitelisted='false'
    for fs_whitelist_idx in "${!fs_type_whitelist[@]}"; 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

    ## Skip mounts for removable media under (or on) /media and /mnt.
    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 in "${!lsblk_path_list[@]}"; 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

    ## Skip mounts that are already read-only.
    skip_fs_remount_ro='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
        skip_fs_remount_ro='false'
        break
      fi
    done
    if [ "${skip_fs_remount_ro}" = 'true' ]; then
      continue
    fi

    ## Skip mounts that are listed in disallow_readonly_list.
    for disallow_readonly_dir in "${disallow_readonly_list[@]}"; do
      if [ "${disallow_readonly_dir}" = "${proc_mount_path}" ]; then
        skip_fs_remount_ro='true'
        break
      fi
    done
    if [ "${skip_fs_remount_ro}" = 'true' ]; then
      continue
    fi

    ## Skip mounts that are listed in disallow_readonly_with_submounts_list,
    ## if those mounts contain submounts.
    for submount_dir in "${dir_submount_list[@]}"; do
      if [ "${proc_mount_path}" = "${submount_dir}" ]; then
        for disallow_readonly_dir in "${disallow_readonly_with_submounts_list[@]}"; do
          if [ "${proc_mount_path}" = "${disallow_readonly_dir}" ]; then
            skip_fs_remount_ro='true'
            break
          fi
        done
        if [ "${skip_fs_remount_ro}" = 'true' ]; then
          break
        fi
      fi
    done
    if [ "${skip_fs_remount_ro}" = 'true' ]; then
      continue
    fi

    ## If we got this far, this is a mount we want to harden. Add it to the
    ## list and mark it as needing remounted read-only and overlaid. (We will
    ## mark mounts that need to be remounted read-only but not overlaid
    # a bit later.)
    overlay_mount_list_str+="${proc_mount_path}"$'\n'
    skip_dir_overlay_bool_list+=( 'false' )
  done

  ## Trim a trailing newline from overlay_mount_list_str to keep it from
  ## becoming an additional array element.
  overlay_mount_list_str="$(head -c-1 <<< "${overlay_mount_list_str}")"

  ## Sort target_overlay_mount_list so submounts can be easily found within
  ## it, and split it into individual paths. The same sort order as specified
  ## in populate_proc_mount_lists is used, though it is much simpler to
  ## implement here.
  overlay_mount_list_str="${overlay_mount_list_str//\//"${ascii_stx}"}"
  overlay_mount_list_str="$(LC_ALL='C' sort <<< "${overlay_mount_list_str}")"
  overlay_mount_list_str="${overlay_mount_list_str//"${ascii_stx}"/\/}"
  readarray -t target_overlay_mount_list <<< "${overlay_mount_list_str}"

  ## Find and mark dirs that should not be overlaid.
  for mount_idx in "${!target_overlay_mount_list[@]}"; do
    target_overlay_mount="${target_overlay_mount_list[mount_idx]}"
    skip_dir_overlay="false"
    allow_hide_submounts='false'

    ## If a dir has submounts, and it does not have an exception in
    ## allow_hiding_submounts_list, prevent it from being overlaid.
    for submount_dir in "${dir_submount_list[@]}"; do
      if [ "${target_overlay_mount}" = "${submount_dir}" ]; then
        for allow_hide_submounts_dir in "${allow_hiding_submounts_list[@]}"; do
          if [ "${target_overlay_mount}" = "${allow_hide_submounts_dir}" ]; then
            allow_hide_submounts='true'
            break
          fi
        done
        if [ "${allow_hide_submounts}" = 'true' ]; then
          break
        fi

        skip_dir_overlay='true'
        break
      fi
    done

    ## If a mount's filesystem is marked non-overlay-able, prevent it from
    ## being overlaid.
    if [ "${skip_dir_overlay}" = 'false' ]; then
      for proc_mount_idx in "${!proc_mount_path_list[@]}"; do
        if [ "${proc_mount_path_list[proc_mount_idx]}" \
          != "${target_overlay_mount}" ]; then
          continue
        fi
        for fs_nooverlay_item in "${fs_type_nooverlay_list[@]}"; do
          if [ "${proc_mount_type_list[proc_mount_idx]}" \
            = "${fs_nooverlay_item}" ]; then
            skip_dir_overlay='true'
            break
          fi
        done
      done
    fi

    ## If a directory is not going to be overlaid, mark it in the list.
    if [ "${skip_dir_overlay}" = 'true' ]; then
      skip_dir_overlay_bool_list[mount_idx]='true'
    fi

    ## Avoid mounting overlays on mounts that are going to end up hidden.
    if [ "${allow_hide_submounts}" = 'true' ]; then
      ## Mark all submounts so that they are not overlaid.
      for target_dir_idx in "${!target_overlay_mount_list[@]}"; do
        target_dir="${target_overlay_mount_list[target_dir_idx]}"
        if [ "${target_dir}" != "${target_overlay_mount}" ] \
          && [[ "${target_dir}" == "${target_overlay_mount%/}/"* ]]; then
          skip_dir_overlay_bool_list[target_dir_idx]='true'
        fi
      done
    fi
  done
}

## Remounts filesystems read-only and applies tmpfs overlays.
harden_mounts() {
  local mount_idx target_overlay_mount tmpfs_dir tmpfs_upper_dir \
    tmpfs_work_dir skip_dir_overlay overlay_idx

  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

  ## Reset the tmpfs repo.
  cmd_wrapper safe-rm -r -f -- "${overlay_tmpfs_repo}"
  cmd_wrapper mkdir --parents -- "${overlay_tmpfs_repo}"

  ## Remount targeted directories read-only.
  for mount_idx in "${!target_overlay_mount_list[@]}"; do
    target_overlay_mount="${target_overlay_mount_list[mount_idx]}"

    if ! cmd_wrapper mount -o remount,ro -- "${target_overlay_mount}"; then
      skip_dir_overlay_bool_list[mount_idx]='true'
    fi
  done

  ## Mount writable overlays over those directories that are safe to overlay.
  overlay_idx='0'
  for mount_idx in "${!target_overlay_mount_list[@]}"; do
    target_overlay_mount="${target_overlay_mount_list[mount_idx]}"
    skip_dir_overlay="${skip_dir_overlay_bool_list[mount_idx]}";

    if [ "${skip_dir_overlay}" = 'true' ]; then
      continue
    fi

    (( overlay_idx++ )) || true
    tmpfs_dir="${overlay_tmpfs_repo}/${overlay_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
      ## If the overlay mount fails, simply leave it remounted read-only and
      ## clean up the tmpfs that would have been overlaid on it.
      cmd_wrapper umount -- "${tmpfs_dir}"
      cmd_wrapper safe-rm -r -f -- "${tmpfs_dir}"
    fi
  done
}

## Finds mounts in proc_mount_path_list that contain submounts, and stores
## those paths in dir_submount_list.
find_submount_dirs() {
  local prev_item counter

  if (( ${#proc_mount_path_list[@]} <= 1 )); then
    return
  fi

  prev_item="${proc_mount_path_list[0]}"
  ## counter = 1, not 0, since we're comparing each item in the list with the
  ## item before it.
  for (( counter = 1; counter < ${#proc_mount_path_list[@]}; counter++ )); do
    if ! [[ "${proc_mount_path_list[counter]}" == /* ]]; then
      prev_item="${proc_mount_path_list[counter]}"
      continue
    fi
    if ! [[ "${proc_mount_path_list[counter]}" == "${prev_item%/}/"* ]]; then
      prev_item="${proc_mount_path_list[counter]}"
      continue
    fi
    dir_submount_list+=( "${prev_item}" )
    prev_item="${proc_mount_path_list[counter]}"
  done
}

check_in_live_mode
populate_proc_mount_lists
find_submount_dirs
get_mount_list_to_harden
if [ "${LIVE_HARDENER_TEST:-}" = 'true' ]; then
  printf '%s\n' "${target_overlay_mount_list[@]}"
  printf '%s\n' "${skip_dir_overlay_bool_list[@]}"
else
  harden_mounts
fi

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