#!/bin/bash

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

set -e

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

error_handler() {
   local exit_code="$?"
   trap "" ERR
   echo "\
###############################################################################
## Swap File Creator ERROR
## $0
##
## BASH_COMMAND: $BASH_COMMAND
## exit_code: $exit_code
##
## For debugging, please run:
## bash -x $0
##
## Clean the output and submit to developers.
###############################################################################\
" >&2
   exit 1
}

trap "error_handler" ERR

sanity_tests() {
   command -v cryptsetup >/dev/null 2>&1
   command -v mkdir >/dev/null 2>&1
   command -v rm >/dev/null 2>&1
   command -v chown >/dev/null 2>&1
   command -v mkswap >/dev/null 2>&1
   command -v swapon >/dev/null 2>&1
   command -v swapoff >/dev/null 2>&1
   command -v shred >/dev/null 2>&1
   command -v cat >/dev/null 2>&1
   command -v test >/dev/null 2>&1
   command -v echo >/dev/null 2>&1
   command -v stat >/dev/null 2>&1
   command -v /usr/lib/systemd/systemd-cryptsetup >/dev/null 2>&1
}

source_config_file() {
  if [ -f "/etc/default/swap-file-creator" ]; then
    bash -n "/etc/default/swap-file-creator"
    source "/etc/default/swap-file-creator"
  fi
}

set_defaults() {
  [[ -v "DESC" ]] || DESC="Encrypted Swap File Creator"
  [[ -v "NAME" ]] || NAME="swap-file-creator"
  [[ -v "SWAPFILE" ]] || SWAPFILE="/var/swapfile"
  [[ -v "VERBOSE" ]] || VERBOSE="yes"
  [[ -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=""
}

do_stop() {
   if test -e "$MAPPER_FULL" ; then
      swapoff "$MAPPER_FULL" >/dev/null 2>&1 || true
   fi

   /usr/lib/systemd/systemd-cryptsetup detach "$MAPPER" 2>/dev/null

   if [ "$SHRED_ON_STOP" = "yes" ]; then
      if [ -f "$SWAPFILE" ]; then
         echo "INFO: Shredding '$SWAPFILE'... This may take a while..."
         shred $SHRED_OPTS "$SWAPFILE"
         echo "INFO: Done shred '$SWAPFILE'."
      fi
   fi

   return 0
}

do_start() {
   echo "INFO: swap-file-creator..."

   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
   fi

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

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

   local swap_file_actual_size_in_bytes
   if [ -f "$SWAPFILE" ]; then
      swap_file_actual_size_in_bytes="$(stat --format='%s' "$SWAPFILE")"
   else
      swap_file_actual_size_in_bytes="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 ))"

   if [ "$total_ram_in_mb" -ge "$ENOUGH_RAM" ]; then
      echo "INFO: Enough RAM available. More than $ENOUGH_RAM MB RAM. OK."
   else
      echo "INFO: Less than $ENOUGH_RAM MB RAM available."
      echo "INFO: Ideally you should assign 2048 MB (2 GB) or more RAM to this machine."
      echo "INFO: See: https://www.kicksecure.com/wiki/swap-file-creator"
   fi

   local dirname_swapfile
   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.
      echo "$calculate_swap_size_report"

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

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

   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 [ "$systemd_detect_virt_output" = "xen" ]; then
      echo "INFO: xen detected. Could be a Qubes HVM. Skipping swap-file-creator. Doing nothing, ok."
      exit 0
   fi

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

   local swap_file_target_size_bytes=$((swap_file_target_size_mb * 1000000))

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

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

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

   do_stop

   ## 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"

   mkswap --label "swapfile" --uuid "$UUID" "$MAPPER_FULL" $MKSWAP_EXTRA >/dev/null
   swapon "$MAPPER_FULL" $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 --raw --show=LABEL,SIZE | grep '\bswapfile\b')
   swap_file_size_with_unit=$(echo "$swapon_line" | awk '{print $NF}')

   # 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
      echo "\
WARNING: Minor issue.
Unit test if swap file with correct size has been created has been skipped.
Output of:
    swapon --raw --show=LABEL,SIZE | grep '\bswapfile\b'
showed unknown unit, not M or G.
swap_file_size_with_unit: '$swap_file_size_with_unit'"
   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
      echo "INFO: Encrypted '$SWAPFILE' ($swap_file_size_gb GB) ready. OK."
      return 0
   fi

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

if grep --no-messages --quiet 'boot=live' /proc/cmdline; then
   true "INFO: grub-live mode detected. Swap file creation skipped."
   exit 0
elif grep --no-messages --quiet 'root=live' /proc/cmdline; then
   true "INFO: ISO live mode detected. Swap file creation skipped."
   exit 0
fi

sanity_tests
source_config_file
set_defaults

if [ $# -eq 0 ]; then
    echo "ERROR: No arguments provided."
    echo "Usage: $0 [start|stop]"
    exit 1
fi

case "$1" in
    start)
        do_start
        ;;
    stop)
        do_stop
        ;;
    *)
        echo "ERROR: Invalid argument: $1"
        echo "Usage: $0 [start|stop]"
        exit 1
        ;;
esac
