#!/usr/bin/bash
#
# sbupdate -- Generate and sign kernel images for UEFI Secure Boot on Arch Linux
# Copyright (C) 2016-2020 Andrey Vihrov <andrey.vihrov@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

set -eu

shopt -s extglob

readonly CONFFILE="/etc/sbupdate.conf"

# Print an error and return unsuccessfully
#  $1: error message
function error() {
  echo "$0: error: $1" >&2
  return 1
}

# Load configuration
function load_config() {
  KEY_DIR="/etc/efi-keys"
  ESP_DIR="/boot"
  OUT_DIR="EFI/Arch"
  SPLASH="/usr/share/systemd/bootctl/splash-arch.bmp"
  BACKUP=1
  EXTRA_SIGN=()
  declare -g -A CONFIGS CMDLINE INITRD OUTPUT

  shopt -s nullglob
  INITRD_PREPEND=(/boot/@(intel|amd)-ucode.img)
  shopt -u nullglob

  # shellcheck disable=SC1090
  source "${CONFFILE}"

  [[ -d "${ESP_DIR}" ]] || error "${ESP_DIR} does not exist"
  [[ -n "${CMDLINE_DEFAULT:+x}" ]] \
    || error "CMDLINE_DEFAULT is not defined or empty in ${CONFFILE}"

  local key=("${KEY_DIR}"/@(DB|db).key); readonly KEY="${key[0]}"
  local cert=("${KEY_DIR}"/@(DB|db).crt); readonly CERT="${cert[0]}"

  readonly KEY_DIR ESP_DIR OUT_DIR SPLASH BACKUP EXTRA_SIGN INITRD_PREPEND CMDLINE_DEFAULT
  readonly -A CONFIGS CMDLINE INITRD OUTPUT
}

# Parse script arguments
#  $@: arguments
function parse_args() {
  HOOK=0
  REMOVE=0

  while getopts "kr" opt; do
    case "${opt}" in
      k) HOOK=1 ;;
      r) REMOVE=1 ;;
      ?) exit 1 ;;
    esac
  done

  readonly HOOK REMOVE
}

# Find the location of the systemd EFI stub
function find_efi_stub() {
  local uname; uname="$(uname -m)"
  case "${uname}" in
    x86_64)
      readonly EFISTUB="/usr/lib/systemd/boot/efi/linuxx64.efi.stub"
      ;;
    i686)
      readonly EFISTUB="/usr/lib/systemd/boot/efi/linuxia32.efi.stub"
      ;;
    *)
      error "unsupported architecture: ${uname}"
      ;;
  esac
}

# Create a list of kernels to process
function get_kernels() {
  local force_all=0 kdir
  declare -g -a KERNELS

  if (( HOOK )); then
    # The script was run from the hook. Read standard input to determine
    # which kernels we need to update.
    while read -r target; do
      if [[ "${target}" =~ ^usr/lib/modules/.+/vmlinuz$ ]]; then
        # Regular kernel
        kdir="$(dirname "/${target}")"
        KERNELS+=("$(<"${kdir}/pkgbase")")
        [[ -f "${kdir}/kernelbase" ]] && KERNELS[-1]="$(<"${kdir}/kernelbase")"
      else
        # Another dependency; update all kernels
        force_all=1
      fi
    done
  else
    # The script was run by the user
    force_all=1
  fi

  if (( force_all )); then
    (( ! REMOVE )) || error "trying to remove all kernels"
    KERNELS=(/boot/vmlinuz-*); KERNELS=("${KERNELS[@]#/boot/vmlinuz-}")
  fi
  readonly -a KERNELS
}

# Return output file path corresponding to an image
#   $1: configuration name
function output_name() {
  echo "${ESP_DIR}/${OUTPUT[$1]:-${OUT_DIR}/$1-signed.efi}"
}

# Remove a signed kernel image
#   $1: configuration name
function remove_image() {
  local output; output="$(output_name "$1")"
  echo "Removing $(basename "${output}")"
  if (( BACKUP )); then
    mv -f "${output}" "${output}.bak"
  else
    rm "${output}"
  fi
}

# Sign a single file
#   $*: arguments to sbsign
function sign_file() {
  sbsign --key "${KEY}" --cert "${CERT}" "$@"
}

# Generate a signed kernel image
#   $1: configuration name
#   $2: kernel name
function update_image() {
  local linux="/boot/vmlinuz-$2"
  local initrd="${INITRD[$1]:-/boot/initramfs-$1.img}"
  local cmdline="${CMDLINE[$1]:-${CMDLINE_DEFAULT}}"
  local output; output="$(output_name "$1")"

  echo "Generating and signing $(basename "${output}")"

  # Create a combined binary with systemd EFI stub. For additional information see:
  #   https://github.com/systemd/systemd/blob/master/src/boot/efi/stub.c
  #   https://github.com/systemd/systemd/blob/master/test/test-efi-create-disk.sh
  #
  # Prepend initramfs files are joined with the main initramfs in one image. Refer to:
  #   https://www.kernel.org/doc/Documentation/early-userspace/buffer-format.txt
  #   https://www.kernel.org/doc/Documentation/x86/microcode.txt
  objcopy \
    --add-section .osrel="/etc/os-release"                          --change-section-vma .osrel=0x20000    \
    --add-section .cmdline=<(echo -n "${cmdline}")                  --change-section-vma .cmdline=0x30000  \
    --add-section .splash="${SPLASH}"                               --change-section-vma .splash=0x40000   \
    --add-section .linux="${linux}"                                 --change-section-vma .linux=0x2000000  \
    --add-section .initrd=<(cat "${INITRD_PREPEND[@]}" "${initrd}") --change-section-vma .initrd=0x3000000 \
    "${EFISTUB}" "${output}"
  wait $!

  # Sign the resulting output file
  sign_file --output "${output}" "${output}"
}

# Map kernel versions to image names and process changes
function process_kernels() {
  for name in "${KERNELS[@]}"; do
    for cfg in ${CONFIGS[${name}]:-${name}}; do # Note: unquoted expansion
      if (( REMOVE )); then
        remove_image "${cfg}"
      else
        update_image "${cfg}" "${name}"
      fi
    done
  done
}

# Check and sign a user-specified extra file
#   $1: file path
function check_sign_extra_file() {
  if sbverify --cert "${CERT}" "$1" >/dev/null; then
    echo "Skipping already signed file $1"
  elif (( HOOK )); then
    # Signing extra files from the hook is prohibited for security reasons
    echo "warning: failed to verify $1" >&2
  else
    echo "Signing $1"
    sign_file --output "$1" "$1"
  fi
}

# Entry point
function main() {
  load_config
  parse_args "$@"

  find_efi_stub
  get_kernels

  mkdir -p "${ESP_DIR}/${OUT_DIR}"
  process_kernels

  for f in "${EXTRA_SIGN[@]}"; do
    check_sign_extra_file "$f"
  done
}

if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then main "$@"; fi

# vim:set ts=2 sw=2 et:
