#!/bin/bash
set -euo pipefail

# The regex path arguments for --keep-paths are treated as regular arguments
opts=$(getopt --name "$(basename "${0}")" --options 'hF:U:KMBSPWDR' \
       --longoptions 'help,ignition-file:,ignition-url:,keep-paths,keep-machine-id,backup,stop,preview-delete,preview-keep,delete-backup,restore-backup' -- "${@}")
eval set -- "${opts}"

KEEPMACHINEID=
IGNITIONFILE=
IGNITIONURL=
HASKEEPPATHS=
BACKUP=
STOP=
PREVIEWDELETE=
PREVIEWKEEP=
DELETEBACKUP=
RESTOREBACKUP=

while true; do
  case "$1" in
  -h|--help)
    echo "Usage: $(basename "${0}") [--ignition-file FILE] [--ignition-url URL] [--keep-machine-id] [--keep-paths REGEX...]"
    echo "  Resets Flatcar Container Linux through a (selective) cleanup of the root filesystem during the next boot."
    echo "  Paths for data to retain can be specified as regular expressions."
    echo "  Ignition will run again, and a local or remote Ignition configuration source can also be set up."
    echo "  A full or selective discard of the root filesystem allows to reconfigure the system while avoiding config drift."
    echo "  The /etc/machine-id file will be deleted but it is possible to keep the machine ID by letting it be part of the kernel cmdline."
    echo "  When paths to keep are specified, only needed paths should be used and not those set up by the old Ignition config"
    echo "  or side effects of it, to really discard the old configuration state. When a path specified is a folder, the contents are"
    echo "  preserved as well because MYPATH/.* is automatically appended as additonal regular expression for paths to keep."
    echo "  To delete the contents of a folder but keep the folder itself, specify it as equivalent regular expression in the form of"
    echo "  '^/etc/mypath', '/etc/mypath$', '/etc/mypat[h]', '/etc/(mypath)', or '(/etc/mypath)'. The used regular expression language"
    echo "  is that of egrep."
    echo "  Meaningful examples are:"
    echo "  - '/etc/ssh/ssh_host_.*' to preserve SSH host keys"
    echo "  - '/var/log' to preserve system logs"
    echo "  - '/var/lib/docker' '/var/lib/containerd' to preserve container state and images"
    echo "  The rootfs does not include the files from /etc that are provided by the overlay mount unless they were copied up."
    echo "  Therefore, you won't see them in the preview or backup."
    echo "Options:"
    echo "  -F, --ignition-file <FILE>	Writes the given Ignition config JSON file to /oem/config.ign"
    echo "  -U, --ignition-url <URL>	Writes the given Ignition config JSON URL as kernel cmdline parameter to /oem/grub.cfg"
    echo "  -K, --keep-paths <REGEX>...	Writes the given regular expressions for paths to keep as combined OS reset info to /selective-os-reset"
    echo "  -M, --keep-machine-id		Writes the current machine ID as kernel cmdline parameter to /oem/grub.cfg to preserve it"
    echo "  -B, --backup			Copies the files that will be deleted to /flatcar-backup/"
    echo "Actions (exclusive, no OS reset will be staged):"
    echo "  -S, --stop			Stops the staged OS reset and nothing will happen on the next boot"
    echo "  -P, --preview-delete		Prints the files that will be deleted"
    echo "  -W, --preview-keep		Prints the files that will be kept"
    echo "  -D, --delete-backup		Deletes the backup under /flatcar-backup/"
    echo "  -R, --restore-backup		Copies the files from /flatcar-backup/ to their original location on the rootfs"
    echo
    echo "Example for selectively resetting the OS with retriggering Ignition while keeping SSH host keys, logs, and machine ID:"
    echo "  sudo $(basename "${0}") --keep-machine-id --keep-paths '/etc/ssh/ssh_host_.*' /var/log"
    echo "  sudo systemctl reboot"
    exit 1
    ;;
  -F|--ignition-file)
    shift
    if [[ -n "${IGNITIONURL}" ]]; then
        echo "Error: Can't specify both Ignition URL and Ignition file at the same time" > /dev/stderr ; exit 1
    fi
    IGNITIONFILE="$1"
    ;;
  -U|--ignition-url)
    shift
    if [[ -n "${IGNITIONFILE}" ]]; then
        echo "Error: Can't specify both Ignition URL and Ignition file at the same time" > /dev/stderr ; exit 1
    fi
    IGNITIONURL="$1"
    ;;
  -K|--keep-paths)
    HASKEEPPATHS=1
    ;;
  -M|--keep-machine-id)
    KEEPMACHINEID=1
    ;;
  -B|--backup)
    BACKUP=1
    ;;
  -S|--stop)
    STOP=1
    ;;
  -P|--preview-delete)
    PREVIEWDELETE=1
    ;;
  -W|--preview-keep)
    PREVIEWKEEP=1
    ;;
  -D|--delete-backup)
    DELETEBACKUP=1
    ;;
  -R|--restore-backup)
    RESTOREBACKUP=1
    ;;
  --)
    shift
    break;;
  esac
  shift
done

KEEP=("$@")
if [ "${KEEP[*]}" != "" ] && [ "${HASKEEPPATHS}" != 1 ]; then
  echo "Error: Found unused arguments: ${KEEP[*]}" > /dev/stderr ; exit 1
fi
if [ "${KEEP[*]}" = "" ]; then
  if [ "${HASKEEPPATHS}" = 1 ]; then
    echo "Error: No paths to keep specified for --keep-paths argument" > /dev/stderr ; exit 1
  fi
fi
if [ "${HASKEEPPATHS}" = 1 ]; then
  for ENTRY in "${KEEP[@]}"; do
    if [[ "${ENTRY}" = './'* ]] || [[ "${ENTRY}" =~ ^[^/\\\(\{\[\$^].*$ ]]; then
      echo "Error: Invalid path to keep, must be an absolute path or a regex for an absolute path: ${ENTRY}" > /dev/stderr ; exit 1
    fi
  done
fi

if [ -e "/selective-os-reset" ]; then
  echo "INFO: An OS reset was staged already from a previous run."
  echo
fi

checkargs="${STOP}${PREVIEWDELETE}${PREVIEWKEEP}${DELETEBACKUP}${RESTOREBACKUP}"
if [ "${checkargs}" != "" ] && [ "${checkargs}" != "1" ]; then
  echo "Error: Only one exclusive action allowed" > /dev/stderr
  exit 1
fi

[ "$EUID" = "0" ] || { echo "Need to be root: sudo $0 $opts" > /dev/stderr ; exit 1 ; }

UKI_ADDON_DIR="/boot/EFI/Linux/acl.efi.extra.d"
UKI_ADDON_TEMPLATE="/boot/acl/uki-addons/firstboot.addon.efi"

# Detect UKI (systemd-boot) mode by the presence of the addon directory.
is_uki_mode() {
  [[ -d "${UKI_ADDON_DIR}" ]]
}

function generate_regex() {
  local ENTRY=
  echo -n '('
  for ENTRY in "${KEEP[@]}"; do
    # If it ends with / we cut it away as it's optional and also won't match the paths find prints for directories
    ENTRY="${ENTRY%/}"
    # If this here starts with / and doesn't end with $|)|*|]|? we will generate an additional regex entry to keep not only the path but also its contents
    if [[ "${ENTRY}" = /* ]] && [[ "${ENTRY}" != *'$' ]] && [[ "${ENTRY}" != *')' ]] && [[ "${ENTRY}" != *'*' ]] && [[ "${ENTRY}" != *']' ]] && [[ "${ENTRY}" != *'?' ]]; then
      echo -n "${ENTRY}/.*|"
    fi
    echo -n "${ENTRY}|"
  done
  echo '/flatcar-backup|/flatcar-backup/.*|/selective-os-reset)'
  # If nothing should be kept but we need to have at least one entry,
  # therefore, use the flag file itself as entry which will be removed anyway
}

function walkroot() {
  local action="$1"
  local extraarg="${2-}"
  while IFS= read -r -d '' entry; do
    "${action}" "${entry}"
  done < <(unshare -m sh -c "umount /etc && find / -xdev -regextype egrep ${extraarg} -regex '$(generate_regex)' -print0")
  # Don't use -depth to make sure we process directories first.
  # Do the print0 as last action, after filtering.
  true # Do not carry any last condition evaluation over as return code
}

# Handle exclusive actions

if [ "${STOP}" = 1 ]; then
  echo "Removing /selective-os-reset and /boot/flatcar/first_boot"
  rm -f "/selective-os-reset" "/boot/flatcar/first_boot"
  rm -f "${UKI_ADDON_DIR}/firstboot.addon.efi"
  exit 0
elif [ "${PREVIEWDELETE}" = 1 ]; then
  # For find -not means that we look at all files that are not matched by the keep regex
  walkroot echo -not
  echo "Note that it is ok to delete the /bin or /lib symlinks and any other OS files/directories like /etc or /.etc-work as they will be recreated."
  exit 0
elif [ "${PREVIEWKEEP}" = 1 ]; then
  walkroot echo
  exit 0
elif [ "${DELETEBACKUP}" = 1 ]; then
  echo "Removing /flatcar-backup/ directory"
  rm -rf "/flatcar-backup/"
  exit 0
elif [ "${RESTOREBACKUP}" = 1 ]; then
  if [ ! -d "/flatcar-backup/" ]; then
    echo "Error: The directory /flatcar-backup/ does not exist" > /dev/stderr
    exit 1
  fi
  echo "Restoring rootfs files from /flatcar-backup/"
  # TODO: our rsync does not support --acls
  unshare -m sh -c "umount /etc && rsync -x -a --sparse --inplace -v /flatcar-backup/ /"
  echo "You should reboot now"
  exit 0
fi

### Default action is to stage a reset ###

function backup_cp() {
  local entry="$1"
  local newpath="/flatcar-backup/${entry}"
  if [ "${entry}" != "/etc" ] && mountpoint -q "${entry}"; then
    return # Don't copy a mountpoint folder like /proc because we can't restore it well (also skips /)
  fi
  if [ "${entry}" = "/.etc-work" ] || [[ "${entry}" = "/.etc-work/"* ]]; then
    return # No need to store the overlay work dir either
  fi
  if [ ! -d "$(dirname "${newpath}")" ]; then
    local tocreate=""
    while IFS= read -r -d '' pathpart; do
      if [ "${pathpart}" = "" ]; then
        continue
      fi
      tocreate="${tocreate}/${pathpart}"
      if [ ! -d "/flatcar-backup${tocreate}" ]; then
        # Use rsync to create the directory with the right permissions but don't copy its contents
        # TODO: --acls
        unshare -m sh -c "umount /etc && rsync -x -a --exclude='*' '${tocreate}' '/flatcar-backup${tocreate}'"
      fi
    done < <(dirname -z "${entry}" | tr '/' '\0')
  fi
  # TODO: --acls
  unshare -m sh -c "umount /etc && if [ -d '${entry}' ] && [ ! -L '${entry}' ]; then rsync -x -a --exclude='*' '${entry}/' '${newpath}/'; else cp -a '${entry}' '${newpath}'; fi"
}

if [ "${BACKUP}" = 1 ]; then
  echo "Removing existing /flatcar-backup/ directory"
  rm -rf "/flatcar-backup/"
  echo "Copying files that will be deleted to /flatcar-backup/"
  mkdir /flatcar-backup
  walkroot backup_cp -not
else
  echo "WARNING: Running without --backup can cause data loss if the keep paths don't work as expected."
  echo "Also check whether your regex works as wanted with --preview-delete and --preview-keep."
  echo
fi

if [ "${KEEPMACHINEID}" = 1 ]; then
  MACHINEID=$(cat /etc/machine-id)
  touch /oem/grub.cfg
  sed -i "s/systemd\.machine_id=[a-f0-9]*//g" /oem/grub.cfg
  echo "set linux_append=\"\$linux_append systemd.machine_id=${MACHINEID}\"" >> /oem/grub.cfg
  echo "Wrote machine ID as kernel cmdline parameter to /oem/grub.cfg"
else
  if [ -e /oem/grub.cfg ]; then
    sed -i "s/systemd\.machine_id=[a-f0-9]*//g" /oem/grub.cfg
    echo "Removed any hardcoded systemd.machine_id kernel cmdline parameter in /oem/grub.cfg"
  fi
fi
if [ "${IGNITIONFILE}" != "" ]; then
  if [ -e /oem/grub.cfg ]; then
    sed -i "s/ignition.config.url=[^ \"']*//g" /oem/grub.cfg
    echo "Removed any ignition.config.url kernel cmdline parameter in /oem/grub.cfg"
  fi
  cp "${IGNITIONFILE}" /oem/config.ign
  echo "Wrote Ignition file /oem/config.ign"
fi
if [ "${IGNITIONURL}" != "" ]; then
  if [ -e /oem/config.ign ]; then
    rm /oem/config.ign
    echo "Removed Ignition file /oem/config.ign"
  fi
  touch /oem/grub.cfg
  sed -i "s/ignition.config.url=[^ \"']*//g" /oem/grub.cfg
  echo "set linux_append=\"\$linux_append ignition.config.url=${IGNITIONURL}\"" >> /oem/grub.cfg
  echo "Wrote Ignition URL as kernel cmdline parameter to /oem/grub.cfg"
fi
# Throw away rests of previous modifications that are now no-ops
if [ -e /oem/grub.cfg ]; then
  # shellcheck disable=SC2016 # We want to literally match $linux_append
  sed -i '/set linux_append="\$linux_append *"/d' /oem/grub.cfg
fi

generate_regex > /selective-os-reset

touch /boot/flatcar/first_boot

if is_uki_mode; then
  # In UKI (systemd-boot) mode there is no grub.cfg, so the first_boot
  # marker file alone cannot re-trigger Ignition.  Restore the firstboot
  # addon from the template that was saved at image build time.
  # systemd-boot auto-discovers it and appends flatcar.first_boot=detected
  # to the kernel cmdline on the next boot.
  if [[ ! -f "${UKI_ADDON_TEMPLATE}" ]]; then
    echo "Error: Firstboot addon template not found at ${UKI_ADDON_TEMPLATE}" > /dev/stderr
    exit 1
  fi
  cp "${UKI_ADDON_TEMPLATE}" "${UKI_ADDON_DIR}/firstboot.addon.efi"
  echo "Restored UKI firstboot addon from template"
fi

echo "Prepared /selective-os-reset and /boot/flatcar/first_boot"
echo "Staged OS reset, you can reboot now"
