#!/bin/bash
# vim: softtabstop=2 shiftwidth=2 expandtab

cleanup() {
  if [ -f "${TEMP_EFI}" ]; then
    if [ -z "${PRESERVE_EFI}" ]; then
      rm "${TEMP_EFI}"
    else
      echo "NOTICE: saving temporary EFI output ${TEMP_EFI}"
    fi
  fi
  unset TEMP_EFI PRESERVE_EFI

  [ -f "${TEMP_KCL}" ] && rm "${TEMP_KCL}"
  unset TEMP_KCL
}

usage() {
  cat <<-EOF
	USAGE: $0 [-a <arg>] [-r <arg>] [-o <out>] [-e] [-d] [bootenv|zbm-efi]

	Review or update kernel command line (KCL) associated with the ZFS boot
	environment or <bootenv> or the ZFSBootMenu EFI executable <zbm-efi>.

	When the boot environment or EFI executable is unspecified, the current
	root filesystem will be used by default. If "-" is passed, stdin will
	be read as an EFI executable.

	ARGUMENTS
	-a <arg>: Append an argument
	-r <arg>: Remove an argument
	-o <out>: Write output to <out> filesystem or executable
	-e: Open the KCL for editing in \$EDITOR
	-d: Remove entire command line
	EOF
}

check_zfs_fs() {
  zfs list -o name -H "$1" >/dev/null 2>&1
}

zerror() {
  echo "ERROR: $*" >&2
}

## BEGIN: zfsbootmenu-kcl.sh library functions
## These are duplicated here to avoid a dependency in this helper script.
## If you make improvements here, please add them back to the library too.

kcl_tokenize() {
  awk '
    BEGIN {
      strt = 1;
      quot = 0;
    }

    {
      for (i=1; i <= NF; i++) {
        # If an odd number of quotes are in this field, toggle quoting
        if ( gsub(/"/, "\"", $(i)) % 2 == 1 ) {
          quot = (quot + 1) % 2;
        }

        # Print a space if this is not the start of a line
        if (strt == 0) {
          printf " ";
        }

        printf "%s", $(i);

        if (quot == 0) {
          strt = 1;
          printf "\n";
        } else {
          strt = 0;
        }
      }
    }
  '
}

kcl_suppress() {
  local arg rem sup
  while read -r arg; do
    # Check match against all exclusions
    sup=0
    for rem in "$@"; do
      # Arguments match entirely or up to first equal
      if [[ "${arg}" == "${rem}" || "${arg%%=*}" == "${rem}" ]]; then
        sup=1
        break
      fi
    done

    # Echo argument if it was not suppressed
    [ "${sup}" -ne 1 ] && echo "${arg}"
  done
}

kcl_assemble() {
  awk '
    BEGIN{ strt = 1; }

    {
      if (strt == 0) {
        printf " ";
      }

      printf "%s", $0;
      strt = 0;
    }
  '
}

## END: zfsbootmenu-kcl.sh library functions

# Append KCL arguments to the tokenized input stream, tokenizing each in turn.
kcl_append() {
  local arg

  # Carry forward input KCL
  cat

  for arg in "$@"; do
    [ -n "${arg}" ] || continue
    kcl_tokenize <<< "${arg}"
  done
}

strip_kcl() {
  awk '
  BEGIN { first = 1; }

  /^$/ { exit; }

  {
    gsub("[[:space:]]*#.*", "");

    if (length == 0) next;

    if (first == 1) {
      first = 0;
    } else {
      printf " ";
    }

    printf "%s", $0;
  }
  '
}

delete() {
  local bootenv="${1}"

  if ! zfs list -o name -H "${bootenv}" >/dev/null 2>&1; then
    zerror "no boot environment specified"
    usage
    return 1
  fi

  if ! zfs inherit org.zfsbootmenu:commandline "${bootenv}"; then
    zerror "failed to remove KCL from ${bootenv}"
    return 1
  fi
}

raw_kcl_zfs() {
  local bootenv="${1}"

  kcl="$(zfs get -H -o value org.zfsbootmenu:commandline "${bootenv}")" || return 1
  [ "${kcl}" = "-" ] && return

  echo "${kcl}"
}

raw_kcl_efi() {
  local efi kclfile

  efi="${1}"

  if [ ! -r "${efi}" ]; then
    zerror "the EFI file cannot be read"
    return 1
  fi

  if ! kclfile="$(mktemp)"; then
    zerror "failed to create temporary file"
    return 1
  fi

  trap 'rm -f "${kclfile}"; trap - RETURN' RETURN

  if ! objout="$(objcopy --dump-section .cmdline="${kclfile}" "${efi}" 2>&1)"; then
    zerror "failed to dump EFI cmdline: ${objout}"
    return 1
  fi

  cat "${kclfile}"
}

save_kcl_zfs() {
  local kcl bootenv
  bootenv="${1}"
  kcl="$(cat)"

  if [ -n "${kcl}" ]; then
    if ! zfs set org.zfsbootmenu:commandline="${kcl}" "${bootenv}"; then
      zerror "failed to set KCL property on ${bootenv}"
      return 1
    fi
  elif ! zfs inherit org.zfsbootmenu:commandline "${bootenv}"; then
    zerror "failed to clear KCL property on ${bootenv}"
    return 1
  fi

  return 0
}

find_cmdline_gap() {
  local file
  file="${1}"

  if [ ! -r "${file}" ]; then
    zerror "unable to read object file '${file}'"
    return 1
  fi

  if ! OBJDATA="$(objdump -h -w "$1")"; then
    zerror "failed to parse object file; is objdump available?"
    return 1
  fi

  local ready idx name size vma rem offsets cmdoff
  offsets=( )
  cmdoff=
  ready=
  # shellcheck disable=SC2034
  while read -r idx name size vma rem; do
    # Object header table begins with header labeling columns
    if [ "${idx,,}" = "idx" ] && [ "${vma,,}" = "vma" ]; then
      ready="yes"
      continue
    fi

    # Ignore all lines until the header line has been encountered
    [ -n "${ready}" ] || continue

    # Make sure the index field is integral
    [ "${idx}" -eq "${idx}" ] >/dev/null 2>&1 || continue

    # Validate the VMA field, which should be hex
    vma="${vma,,}"
    # Field should not start with 0x, but tolerate it anyway
    vma="${vma#0x}"

    if ! vma="$(( "0x${vma}" ))"; then
      zerror "invalid VMA for section '${name}'"
      return 1
    fi

    if [ "${name,,}" = ".cmdline" ]; then
      cmdoff="${vma}"
    else
      offsets+=( "${vma}" )
    fi
  done <<< "${OBJDATA}"

  if [ -z "${cmdoff}" ]; then
    zerror "file '${file}' contains no .cmdline section"
    return 1
  fi

  local gap mingap
  gap=
  mingap=
  for vma in "${offsets[@]}"; do
    [ "${vma}" -gt "${cmdoff}" ] >/dev/null 2>&1 || continue;
    gap="$(( vma - cmdoff ))"
    if [ -z "${mingap}" ] || [ "${gap}" -lt "${mingap}" ]; then
      mingap="${gap}"
    fi
  done

  if [ -z "${mingap}" ]; then
    zerror "unable to determine .cmdline gap size"
    return 1
  fi

  printf "%X %X\n" "${cmdoff}" "${mingap}"
  return 0
}

save_kcl_efi() {
  local kcl efi kclfile kcloff kclgap kclsize
  efi="${1}"

  # Find offset and space available for the cmdline to replace;
  # this seems to be 4 kB with the x86_64 stub loader alignment
  if ! kclsize="$(find_cmdline_gap "${efi}")"; then
    zerror "failed to determine offset data for KCL"
    return 1
  fi

  read -r kcloff kclgap <<< "${kclsize}"

  if [ -z "${kcloff}" ] || [ -z "${kclgap}" ]; then
    zerror "offset data for KCL appears invalid; aborting"
    return 1
  fi

  if ! kclfile="$(mktemp)"; then
    zerror "failed to save temporary KCL"
    return 1
  fi

  trap 'rm -f "${kclfile}"; trap - RETURN' RETURN

  # Dracut adds a leading space; is this necessary?
  echo -n " " > "${kclfile}"
  cat > "${kclfile}"
  # Dracut also adds a null terminator
  echo -ne "\x00" >> "${kclfile}"

  if ! kclsize="$(stat -c %s "${kclfile}")"; then
    zerror "failed to determine new KCL size; is stat available?"
    return 1
  fi

  if ! [ "${kclsize}" -le "$(( "0x${kclgap}" ))" ] >/dev/null 2>&1; then
    zerror "new KCL size exceeds space available in EFI file '${efi}'; aborting"
    return 1
  fi

  if ! objout="$(objcopy --remove-section .cmdline "${efi}" 2>&1)"; then
    zerror "failed to clear existing KCL from EFI executable"
    return 1
  fi

  local objargs
  objargs=( --add-section ".cmdline=${kclfile}"
            --change-section-vma ".cmdline=0x${kcloff}" )
  if ! objout="$(objcopy "${objargs[@]}" "${efi}" 2>&1)"; then
    zerror "failed to write new KCL to EFI executable"
    return 1
  fi

  return 0
}

# Assemble and allow the user to edit an already tokenized KCL and, if it has
# changed, tokenize the edited version and save back to the input file.
#
# ARGUMENTS
# arg0: pre-existing tokenized KCL
#
# RETURNS
# 0 on successful edit
# 1 on error
edit() {
  local kclsrc kclfile

  kclsrc="${1}"

  if ! command -v "${EDITOR:=vi}" >/dev/null 2>&1; then
    zerror "define \$EDITOR to edit"
    return 1
  fi

  if ! [ -r "${kclsrc}" ]; then
    zerror "unable to read KCL"
    return 1
  fi

  if ! kclfile="$(mktemp)"; then
    zerror "failed to create temporary file"
    return 1
  fi

  trap 'rm -f "${kclfile}"; trap - RETURN' RETURN

  kcl_assemble < "${kclsrc}" > "${kclfile}"

  cat >> "${kclfile}" <<-EOF


	# KCL processing ends with the first line that contains no text.
	# Anything starting with # is considered a comment and will be ignored.
	# Multiple lines will be concatenated with spaces to form a single line.
	EOF

  if ! "${EDITOR}" "${kclfile}"; then
    zerror "failed to edit KCL"
    return 1
  fi

  if ! strip_kcl < "${kclfile}" | kcl_tokenize > "${kclsrc}"; then
    zerror "failed to save edited KCL"
    return 1
  fi

  return 0
}

getbootenv() {
  local dev mntpt fs opts

  # shellcheck disable=SC2034
  while read -r dev mntpt fs opts; do
    [ "${mntpt}" = "/" ] || continue

    if [ "${fs}" != "zfs" ]; then
      return 1
    fi

    echo "${dev}"
    return 0
  done < /proc/mounts
}

delkcl=
output=
editkcl=
appends=()
removes=()

while getopts "ha:r:o:ed" opt; do
  case "${opt}" in
    a)
      appends+=( "${OPTARG}" )
      ;;
    r)
      removes+=( "${OPTARG}" )
      ;;
    o)
      output="${OPTARG}"
      ;;
    e)
      editkcl="yes"
      ;;
    d)
      delkcl="yes"
      ;;
    h)
      usage
      exit
      ;;
    *)
      usage
      exit 1
      ;;
  esac
done

shift $((OPTIND-1))

input="${1}"
if [ -z "${input}" ]; then
  if ! input="$(getbootenv)"; then
    zerror "root does not appear to be ZFS" >&2
    exit 1
  fi
fi

if [ -z "${output}" ]; then
  output="${input}"
fi

modify=
if (( ${#removes[@]} != 0 || ${#appends[@]} != 0 )); then
  modify="yes"
elif [ "${editkcl}" = "yes" ]; then
  modify="yes"
elif [ "${delkcl}" = "yes" ]; then
  modify="yes"
fi

# Make sure to clean up temporary files on exit
unset TEMP_EFI PRESERVE_EFI TEMP_KCL
trap cleanup EXIT INT TERM

efi_mode=
kcl_reader=( "raw_kcl_zfs" "${input}" )
kcl_writer=( "save_kcl_zfs" "${output}" )

if [ "${input}" = "-" ] || [ -f "${input}" ]; then
  # When input is stdin or a file, assume it is an EFI executable
  efi_mode=yes

  if ! TEMP_EFI="$(mktemp)"; then
    zerror "unable to create temporary file for EFI executable edits"
    exit 1
  fi

  if [ "${input}" = "-" ]; then
    if [ "${editkcl}" = "yes" ]; then
      zerror "cannot edit KCL in streaming mode"
      exit 1
    fi

    if ! cat > "${TEMP_EFI}"; then
      zerror "failed to save EFI executable from stdin"
      exit 1
    fi

    input="${TEMP_EFI}"
  elif ! cp "${input}" "${TEMP_EFI}"; then
      zerror "failed to copy EFI executable '${input}' to workspace"
      exit 1
  fi

  kcl_reader=( "raw_kcl_efi" "${TEMP_EFI}" )
  kcl_writer=( "save_kcl_efi" "${TEMP_EFI}" )
elif ! check_zfs_fs "${input}"; then
  zerror "KCL source is not a ZFS filesystem or EFI executable"
  exit 1
elif [ "${output}" != "${input}" ] && ! check_zfs_fs "${output}"; then
  zerror "KCL destination is not a ZFS filesystem"
  exit 1
fi

if ! TEMP_KCL="$(mktemp)"; then
  zerror "failed to create temporary KCL file"
  exit 1
fi

# In "delete" mode, just "read" an empty string for the input kcl
if [ "${delkcl}" = "yes" ]; then
  kcl_reader=( "true" )
fi

# Read, tokenize and modify the input KCL
if ! "${kcl_reader[@]}" | kcl_tokenize | \
      kcl_suppress "${removes[@]}" | \
      kcl_append "${appends[@]}" > "${TEMP_KCL}"; then
  zerror "failed to parse input KCL"
  exit 1
fi

# If the intent is not to modify, just print the existing KCL
if [ "${modify}" != "yes" ]; then
  kcl_assemble < "${TEMP_KCL}"
  echo ""
  exit 0
fi

# Allow the user to edit if appropriate
if [ "${editkcl}" = "yes" ] && ! edit "${TEMP_KCL}"; then
  exit 1
fi

# Save the working KCL back to the output
if ! kcl_assemble < "${TEMP_KCL}" | "${kcl_writer[@]}"; then
  zerror "failed to save KCL changes"
  exit 1
fi

# In EFI mode, try to copy the temporary EFI to the output
if [ "${efi_mode}" = "yes" ]; then
  if [ "${output}" = "-" ]; then
    cat "${TEMP_EFI}"
  elif ! cp "${TEMP_EFI}" "${output}"; then
    zerror "failed to write EFI executable '${output}'; saved as '${TEMP_EFI}'"
    PRESERVE_EFI="yes"
    exit 1
  fi
fi