#!/bin/bash

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

#set -x

set -e
set -o pipefail
set -o errtrace
set -o nounset

## https://forums.whonix.org/t/swap-file-creator-error/21128/3
## avoid printf error
#bash: printf: 3.14: invalid number
##
## Does not work:
#export LANG=C
## https://unix.stackexchange.com/questions/576701/what-is-the-difference-between-lang-c-and-lc-all-c
LC_ALL=C
export LC_ALL

pre_check_done='false'

log() {
  if [ "${VERBOSE-}" = "yes" ]; then
    printf '%s\n' "$@"
    return 0
  fi
  printf '%s\n' "$0: $@" | systemd-cat
  return 0
}

error_handler() {
   local exit_code="$?"
   trap "" ERR
   log "\
###############################################################################
## Swap File Creator ERROR
## $0
##
## BASH_COMMAND: '$BASH_COMMAND'
## exit_code: '$exit_code'
##
## For debugging, please run:
##     bash -x $0 start
## and/or:
##     bash -x $0 stop
##
## Clean the output and submit to developers.
##
## See also:
## This does not indicate a system compromise. The system can run without a swap file.
## https://www.kicksecure.com/wiki/swap-file-creator
###############################################################################\
" >&2
   exit 1
}

trap "error_handler" ERR

exit_handler() {
   local exit_code="$?"
   if [ "$exit_code" = "0" ]; then
      log "INFO: END: OK"
   else
      log "ERROR: END: with ERROR exit code: '$exit_code'" >&2
   fi
   exit "$exit_code"
}

trap exit_handler EXIT

## sets: live_status_detected
source /usr/libexec/helper-scripts/live-mode.sh

sanity_tests() {
  local cmd_list cmd

  #cryptsetup
  #mkdir
  #/usr/lib/systemd/systemd-cryptsetup
  cmd_list=(
    'safe-rm'
    'chown'
    'mkswap'
    'swapon'
    'swapoff'
    'shred'
    'cat'
    'test'
    'printf'
    'stat'
    'mountpoint'
    'grep'
    'sed'
    'blkid'
    'findmnt'
    'fallocate'
    'df'
    'free'
    'bc'
    'calculate-swap-size'
  )

  for cmd in "${cmd_list[@]}"; do
    command -v "$cmd" >/dev/null
  done
}

source_config_files() {
  shopt -s nullglob
  local i

  for i in /etc/default/swap-file-creator /etc/default/swap-file-creator.d/*.conf /usr/local/etc/default/swap-file-creator.d/*.conf; do
    if ! test -f "$i"; then
       continue
    fi
    bash -n "$i"
    source "$i"
  done
}

set_defaults() {
  [[ -v "DESC" ]] || DESC="Swap File Creator"
  [[ -v "NAME" ]] || NAME="swap-file-creator"
  [[ -v "SWAPFILE" ]] || SWAPFILE="/var/swapfile"
  [[ -v "VERBOSE" ]] || VERBOSE="no"
  [[ -v "UUID" ]] || UUID="0615ba72-85b0-4183-8d54-300bb0d2e491"
  #[[ -v "MAPPER" ]] || MAPPER="swapfile"
  #[[ -v "MAPPER_FULL" ]] || MAPPER_FULL="/dev/mapper/$MAPPER"
  [[ -v "RUN_FOLDER" ]] || RUN_FOLDER="/run/$NAME"
  ## https://gitlab.com/cryptsetup/cryptsetup/-/issues/617
  #[[ -v "cryptsetup_pre_wrapper" ]] || cryptsetup_pre_wrapper="ld-system-preload-disable"
  [[ -v "SWAPON_EXTRA" ]] || SWAPON_EXTRA=""
  [[ -v "MKSWAP_EXTRA" ]] || MKSWAP_EXTRA=""
  [[ -v "SHRED_ON_STOP" ]] || SHRED_ON_STOP="no"
  [[ -v "SHRED_OPTS" ]] || SHRED_OPTS="--verbose --iterations=1"
  [[ -v "RANDOM_DEVICE" ]] || RANDOM_DEVICE="/dev/random"
  [[ -v "ENOUGH_RAM" ]] || ENOUGH_RAM="1950"
  [[ -v "SWAP_FILE_SIZE_CUSTOM_MB" ]] || SWAP_FILE_SIZE_CUSTOM_MB=""
  [[ -v "DO_PRE_CHECK" ]] || DO_PRE_CHECK="yes"
}

do_stop() {
   local swapon_output swapoff_exitcode
   swapoff_exitcode='0'
   swapon_output="$(swapon 2>/dev/null)" || true
   if [ -f "$SWAPFILE" ]; then
      if grep --fixed-strings -- "$SWAPFILE" < <(printf '%b' "$swapon_output") &>/dev/null; then
         swapoff "$SWAPFILE" || { swapoff_exitcode="$?"; true; };
      fi

      #if [ "$swapoff_exitcode" = '0' ]; then
      #   /usr/lib/systemd/systemd-cryptsetup detach "$MAPPER"
      #fi
   fi

   ## Avoid shredding the swapfile if swapoff failed, as doing so could shred
   ## actively in-use memory and crash the system.
   if [ ! "$swapoff_exitcode" = '0' ]; then
      return 0
   fi

   if ! [ -f "$SWAPFILE" ]; then
      return 0
   fi

   if [ "$SHRED_ON_STOP" = "yes" ]; then
      log "INFO: Shredding '$SWAPFILE' before removal... This may take a while..."
      shred $SHRED_OPTS "$SWAPFILE"
      log "INFO: Done shredding '$SWAPFILE'."
   fi

   ## TODO: Why is this needed?
#    if [ "$DO_PRE_CHECK" = 'yes' ] \
#       && [ "$pre_check_done" = 'false' ] \
#       && ! luks-path-check "$SWAPFILE" &>/dev/null; then
#       log "INFO: Removing '$SWAPFILE'..."
#       safe-rm --force -- "$SWAPFILE"
#       log "INFO: Done removing '$SWAPFILE'."
#    fi

   safe-rm --force -- "$SWAPFILE"

   return 0
}

do_start() {
   local systemd_detect_virt_output hibernation_consideration dirname_swapfile fileystem_type

   log "INFO: swap-file-creator..."

   total_ram_in_mb="$(free -m | sed  -n -e '/^Mem:/s/^[^0-9]*\([0-9]*\) .*/\1/p')"

   if [ "$total_ram_in_mb" -ge "$ENOUGH_RAM" ]; then
      log "INFO: Enough RAM available (at least $ENOUGH_RAM MB). OK."
   else
      if [ "${VERBOSE-}" = "no" ]; then
         VERBOSE=true
         log "INFO: swap-file-creator..."
      fi
      log "INFO: Low RAM detected (less than $ENOUGH_RAM MB)."
      log "INFO: Consider assigning 2048 MB (2 GB) or more RAM to this machine."
      log "INFO: See: https://www.kicksecure.com/wiki/swap-file-creator"
   fi

   if ! [[ "$SWAPFILE" = /* ]]; then
      log "ERROR: 'SWAPFILE' value '$SWAPFILE' is not an absolute path. Swap file not enabled." >&2
      exit 1
   fi

   systemd_detect_virt_output=""
   if command -v systemd-detect-virt >/dev/null ; then
      if systemd_detect_virt_output="$(systemd-detect-virt 2>/dev/null)" ; then
         ## Exit code 0. Virtualization detected.
         ## Therefore disabling hibernation consideration.
         hibernation_consideration=no
      else
         ## https://forums.kicksecure.com/t/enable-and-use-zram-instead-for-swap/654/8
         #hibernation_consideration=yes
         hibernation_consideration=no
      fi
   else
      ## XXX
      hibernation_consideration=no
      systemd_detect_virt_output=""
   fi

   if [ "$DO_PRE_CHECK" = 'yes' ]; then
      if [ "$systemd_detect_virt_output" = "xen" ]; then
         log "INFO: Virtualizer 'xen' detected. Could be 'Qubes'. Skipping swap-file-creator (no swap file will be created). Doing nothing, OK."
         exit 0
      fi

      if ! luks-path-check "$SWAPFILE" &>/dev/null ; then
         log "INFO: Encrypted disk (LUKS) detected for file '$SWAPFILE': no. Not creating a swap file (expected). Doing nothing, OK. See also: https://www.kicksecure.com/wiki/swap-file-creator"
         exit 0
      else
         pre_check_done='true'
      fi
   fi

   ## User can set HIBERNATION=yes in /etc/default/swap-file-creator
   [[ -v "HIBERNATION" ]] || HIBERNATION="$hibernation_consideration"


   local swap_file_actual_size_in_bytes
   if [ -f "$SWAPFILE" ]; then
      ## Sanity test.
      test -r "$SWAPFILE"
      swap_file_actual_size_in_bytes="$(stat --format='%s' -- "$SWAPFILE")"
   else
      swap_file_actual_size_in_bytes="0"
   fi

   local folder_of_swapfile
   folder_of_swapfile="${SWAPFILE%/*}"

   if ! test -d "$folder_of_swapfile" ; then
      log "ERROR: folder_of_swapfile '$folder_of_swapfile' does not exist." >&2
      exit 1
   fi

   if ! test -w "$folder_of_swapfile" ; then
      true "INFO: File $SWAPFILE (actually folder $folder_of_swapfile) is not writeable. Maybe read-only file system? Exiting."
      exit 0
   fi

   local swap_file_actual_size_in_mb
   swap_file_actual_size_in_mb="$(( swap_file_actual_size_in_bytes / 1024 / 1024 ))"

   ## Never mind eventual small platform specific rounding errors.
   swap_file_actual_size_in_mb="$(( swap_file_actual_size_in_mb + 2 ))"

   local total_disk_space_in_mb
   dirname_swapfile="${SWAPFILE%/*}"
   total_disk_space_in_mb=$(df "$dirname_swapfile" | awk 'NR==2 {print int($2/1024)}')

   local swap_file_target_size_mb
   if [ "$SWAP_FILE_SIZE_CUSTOM_MB" = "" ]; then
      local calculate_swap_size_report
      calculate_swap_size_report="$(calculate-swap-size "$total_ram_in_mb" "$total_disk_space_in_mb" "$HIBERNATION")"

      ## For informational output.
      log "$calculate_swap_size_report"

      ## Take last word of last line.
      swap_file_target_size_mb=$(tail -1 <<< "$calculate_swap_size_report")
      swap_file_target_size_mb=$(awk '{print $NF}' <<< "$swap_file_target_size_mb")
   else
      swap_file_target_size_mb="$SWAP_FILE_SIZE_CUSTOM_MB"
   fi

   local swap_file_size_gb
   swap_file_size_gb=$(bc <<< "scale=0; $swap_file_target_size_mb / 1024")

   local free_disk_space_in_gb
   local free_disk_space_in_gb_plus
   free_disk_space_in_gb=$(df "$dirname_swapfile" | awk 'NR==2 {print int($4/1024/1024)}')
   free_disk_space_in_gb_plus=$(( free_disk_space_in_gb + 2 ))

   if [ "$swap_file_size_gb" -ge "$free_disk_space_in_gb_plus" ]; then
      log "ERROR: Only $free_disk_space_in_gb GB free disk space. Insufficient to create $swap_file_size_gb GB sized swap file. Swap file not enabled." >&2
      exit 1
   fi

   local swap_file_target_size_bytes
   swap_file_target_size_bytes=$((swap_file_target_size_mb * 1000000))

   local swap_file_target_size_gb
   swap_file_target_size_gb="$(bc <<< "scale=2; $swap_file_target_size_bytes/1024/1024/1024")"
   ## Round down to integer.
   swap_file_target_size_gb="$(printf "%.0f" "$swap_file_target_size_gb")"

   do_stop

   fileystem_type=$(findmnt --noheadings --output FSTYPE --target "$folder_of_swapfile")

   if [ "$fileystem_type" = "btrfs" ]; then
      command -v "btrfs" >/dev/null
      ## Need to use 'btrfs'. Otherwise, 'swapon "$SWAPFILE" $SWAPON_EXTRA' would fail with:
      ## > swapon: /var/swapfile: swapon failed: Invalid argument
      ## https://unix.stackexchange.com/questions/599949/swapfile-swapon-invalid-argument
      ## https://btrfs.readthedocs.io/en/latest/Swapfile.html
      btrfs filesystem mkswapfile --size "${swap_file_target_size_gb}G" "$SWAPFILE"
   else
      ## We need to check the size of $SWAPFILE, because a previous run of fallocate that
      ## got interrupted for some reason might have only created a smaller file
      ## (0 MB in worst case).
      ## fallocate is fast. Maybe no more test needed.
      #if [ "$swap_file_target_size_mb" -lt "$swap_file_actual_size_in_mb" ]; then
          ## Not exact.
          ## For example 6144 MB would only be 5.7 G in "sudo swapon --show".
          #fallocate --length "${swap_file_target_size_mb}MB" "$SWAPFILE"
          ## Hence calculate in bytes.
          fallocate --length "$swap_file_target_size_bytes" "$SWAPFILE"
      #fi
   fi

   chown --recursive root:root -- "$SWAPFILE"
   chmod --recursive 0600 -- "$SWAPFILE"

   ## AUDIT: crypto
   #/usr/lib/systemd/systemd-cryptsetup attach "$MAPPER" "$SWAPFILE" "$RANDOM_DEVICE" 'swap,cipher=aes-xts-plain64,size=512'

   #chown --recursive root:root -- "$MAPPER_FULL"
   #chmod --recursive 0600 -- "$MAPPER_FULL"

   if [ "$fileystem_type" = "btrfs" ]; then
      true "INFO: mkswap not needed for btrfs, OK."
   else
      mkswap --label "swapfile" --uuid "$UUID" "$SWAPFILE" $MKSWAP_EXTRA >/dev/null
   fi

   swapon "$SWAPFILE" $SWAPON_EXTRA

   ## /proc/swap does not show the label.
   ## output by swapon_line is not predictable, G vs M as unit.
   local swapon_line
   swapon_line="$(swapon --noheadings --raw --show=NAME,SIZE)"
   swapon_line="$(printf '%b' "$swapon_line" | grep --fixed-strings -- "$SWAPFILE" | tail -n1)" || true
   ## example swapon_line:
   #swapon_line="/var/swapfile 1024M"
   swap_file_size_with_unit=$(awk '{print $NF}' <<< "$swapon_line")

   # Check if the size is in GB or MB and convert accordingly
   if [[ $swap_file_size_with_unit == *G ]]; then
      # If size ends with 'G', remove the 'G'
      swap_file_kernel_reported_size_gb=${swap_file_size_with_unit%G}
   elif [[ $swap_file_size_with_unit == *M ]]; then
      # If size ends with 'M', remove the 'M' and divide by 1024
      swap_file_size_in_mb=${swap_file_size_with_unit%M}
      swap_file_kernel_reported_size_gb=$(bc <<< "scale=2; $swap_file_size_in_mb / 1024")
   else
      log "\
INFO: Minor note (not an error).
Swap file has been enabled; the size self-check was skipped.
Output of:
    swapon --noheadings --raw --show=NAME,SIZE
reported an unexpected unit for file '$SWAPFILE' (not M or G).
swapon_line: '$swapon_line'
swap_file_size_with_unit: '$swap_file_size_with_unit'"
      return 0
   fi

   swap_file_kernel_reported_size_gb_rounded_up="$(printf "%.0f\n" "$swap_file_kernel_reported_size_gb")"

   local swap_file_kernel_reported_size_gb_plus
   swap_file_kernel_reported_size_gb_plus=$(( swap_file_kernel_reported_size_gb_rounded_up + 1 ))

   if [ "$swap_file_kernel_reported_size_gb_plus" -ge "$swap_file_target_size_gb" ]; then
      log "INFO: File '$SWAPFILE' ($swap_file_size_gb GB) created and enabled. OK."
      return 0
   fi

   log "ERROR: swap_file_kernel_reported_size_gb_plus is smaller than expected swap_file_target_size_gb!" >&2
   log "swap_file_kernel_reported_size_gb_plus: '$swap_file_kernel_reported_size_gb_plus'" >&2
   log "swap_file_target_size_gb: '$swap_file_target_size_gb'" >&2
   exit 1
}

if [ "$live_status_detected" = "true" ]; then
  true "INFO: live mode detected. Swap file creation skipped."
  exit 0
fi

sanity_tests
source_config_files
set_defaults

if ! [ "$(id -u)" = 0 ]; then
    log "$0: ERROR: Must be run as root/sudo." >&2
    exit 1
fi

if [ $# -eq 0 ]; then
    log "ERROR: No arguments provided." >&2
    log "Usage: $0 (start|stop)" >&2
    exit 1
fi

case "$1" in
    start)
        do_start
        ;;
    stop)
        do_stop
        ;;
    *)
        log "ERROR: Invalid argument: $1" >&2
        log "Usage: $0 (start|stop)" >&2
        exit 1
        ;;
esac
