#!/usr/bin/env bash
# shellcheck disable=SC2034

set -eu -o pipefail
shopt -qs failglob

# import the partition helper functions
# shellcheck source=partyplanner
. "${0%/*}/partyplanner"

OUTPUT_FMT="raw"
BUILDER_ARCH="$(uname -m)"
OVF_TEMPLATE=""

GRUB_SET_PRIVATE_VAR="no"
XFS_DATA_PARTITION="no"
UEFI_SECURE_BOOT="no"

for opt in "$@"; do
   optarg="$(expr "${opt}" : '[^=]*=\(.*\)')"
   case "${opt}" in
      --package-dir=*) PACKAGE_DIR="${optarg}" ;;
      --output-dir=*) OUTPUT_DIR="${optarg}" ;;
      --output-fmt=*) OUTPUT_FMT="${optarg}" ;;
      --os-image-size-gib=*) OS_IMAGE_SIZE_GIB="${optarg}" ;;
      --data-image-size-gib=*) DATA_IMAGE_SIZE_GIB="${optarg}" ;;
      --os-image-publish-size-gib=*) OS_IMAGE_PUBLISH_SIZE_GIB="${optarg}" ;;
      --data-image-publish-size-gib=*) DATA_IMAGE_PUBLISH_SIZE_GIB="${optarg}" ;;
      --partition-plan=*) PARTITION_PLAN="${optarg}" ;;
      --ovf-template=*) OVF_TEMPLATE="${optarg}" ;;
      --with-grub-set-private-var=*) GRUB_SET_PRIVATE_VAR="${optarg}" ;;
      --xfs-data-partition=*) XFS_DATA_PARTITION="${optarg}" ;;
      --with-uefi-secure-boot=*) UEFI_SECURE_BOOT="${optarg}" ;;
   esac
done

case "${OUTPUT_FMT}" in
   raw|qcow2|vmdk) ;;
   *)
      echo "unexpected image output format '${OUTPUT_FMT}'" >&2
      exit 1
      ;;
esac

case "${PARTITION_PLAN}" in
  split|unified) ;;
  *)
    echo "unexpected partition plan '${PARTITION_PLAN}'" >&2
    exit 1
    ;;
esac

# Fail fast if the OVF template doesn't exist, or doesn't match the layout.
if [ "${OUTPUT_FMT}" == "vmdk" ] ; then
  if [ ! -s "${OVF_TEMPLATE}" ] ; then
    echo "required OVF template not found: ${OVF_TEMPLATE}" >&2
    exit 1
  fi

  if [ "${PARTITION_PLAN}" == "split" ] ; then
    if ! grep -Fq '{{DATA_DISK}}' "${OVF_TEMPLATE}" ; then
      echo "Missing data disk in OVF template, which is required for 'split' layout." >&2
      exit 1
    fi
  fi

  if [ "${PARTITION_PLAN}" == "unified" ] ; then
    if grep -Fq '{{DATA_DISK}}' "${OVF_TEMPLATE}" ; then
      echo "Incorrect data disk in OVF template, which is not supported for 'unified' layout." >&2
      exit 1
    fi
  fi

  if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
    if ! grep -Fq '{{DB_CERT_DER_HEX}}' "${OVF_TEMPLATE}" ; then
      echo "Missing CA certificate field in OVF template, which is required for Secure Boot support." >&2
      exit 1
    fi
  fi
fi

# Store output artifacts in a versioned directory.
OUTPUT_DIR="${OUTPUT_DIR}/${VERSION_ID}-${BUILD_ID}"
mkdir -p "${OUTPUT_DIR}"

FILENAME_PREFIX="${IMAGE_NAME}-${VARIANT}-${ARCH}-${VERSION_ID}-${BUILD_ID}"
SYMLINK_PREFIX="${IMAGE_NAME}-${VARIANT}-${ARCH}"
VERSIONED_SYMLINK_PREFIX="${IMAGE_NAME}-${VARIANT}-${ARCH}-${VERSION_ID}"
FRIENDLY_VERSIONED_SYMLINK_PREFIX="${IMAGE_NAME}-${VARIANT}-${ARCH}-v${VERSION_ID}"

OS_IMAGE_NAME="${FILENAME_PREFIX}"
OS_IMAGE_SYMLINK="${SYMLINK_PREFIX}"
OS_IMAGE_VERSIONED_SYMLINK="${VERSIONED_SYMLINK_PREFIX}"
OS_IMAGE_FRIENDLY_VERSIONED_SYMLINK="${FRIENDLY_VERSIONED_SYMLINK_PREFIX}"

DATA_IMAGE_NAME="${FILENAME_PREFIX}-data"
DATA_IMAGE_SYMLINK="${SYMLINK_PREFIX}-data"
DATA_IMAGE_VERSIONED_SYMLINK="${VERSIONED_SYMLINK_PREFIX}-data"
DATA_IMAGE_FRIENDLY_VERSIONED_SYMLINK="${FRIENDLY_VERSIONED_SYMLINK_PREFIX}-data"

BOOT_IMAGE_NAME="${FILENAME_PREFIX}-boot.ext4.lz4"
BOOT_IMAGE_SYMLINK="${SYMLINK_PREFIX}-boot.ext4.lz4"
BOOT_IMAGE_VERSIONED_SYMLINK="${VERSIONED_SYMLINK_PREFIX}-boot.ext4.lz4"
BOOT_IMAGE_FRIENDLY_VERSIONED_SYMLINK="${FRIENDLY_VERSIONED_SYMLINK_PREFIX}-boot.ext4.lz4"

VERITY_IMAGE_NAME="${FILENAME_PREFIX}-root.verity.lz4"
VERITY_IMAGE_SYMLINK="${SYMLINK_PREFIX}-root.verity.lz4"
VERITY_IMAGE_VERSIONED_SYMLINK="${VERSIONED_SYMLINK_PREFIX}-root.verity.lz4"
VERITY_IMAGE_FRIENDLY_VERSIONED_SYMLINK="${FRIENDLY_VERSIONED_SYMLINK_PREFIX}-root.verity.lz4"

ROOT_IMAGE_NAME="${FILENAME_PREFIX}-root.ext4.lz4"
ROOT_IMAGE_SYMLINK="${SYMLINK_PREFIX}-root.ext4.lz4"
ROOT_IMAGE_VERSIONED_SYMLINK="${VERSIONED_SYMLINK_PREFIX}-root.ext4.lz4"
ROOT_IMAGE_FRIENDLY_VERSIONED_SYMLINK="${FRIENDLY_VERSIONED_SYMLINK_PREFIX}-root.ext4.lz4"

OS_IMAGE="$(mktemp)"
BOOT_IMAGE="$(mktemp)"
VERITY_IMAGE="$(mktemp)"
ROOT_IMAGE="$(mktemp)"
DATA_IMAGE="$(mktemp)"
EFI_IMAGE="$(mktemp)"
PRIVATE_IMAGE="$(mktemp)"
BOTTLEROCKET_DATA="$(mktemp)"

ROOT_MOUNT="$(mktemp -d)"
BOOT_MOUNT="$(mktemp -d)"
DATA_MOUNT="$(mktemp -d)"
EFI_MOUNT="$(mktemp -d)"
PRIVATE_MOUNT="$(mktemp -d)"

SBKEYS="${HOME}/sbkeys"

SELINUX_ROOT="/etc/selinux"
SELINUX_POLICY="fortified"
SELINUX_FILE_CONTEXTS="${ROOT_MOUNT}/${SELINUX_ROOT}/${SELINUX_POLICY}/contexts/files/file_contexts"

VERITY_VERSION=1
VERITY_HASH_ALGORITHM=sha256
VERITY_DATA_BLOCK_SIZE=4096
VERITY_HASH_BLOCK_SIZE=4096

# Bottlerocket has been experimentally shown to boot faster on EBS volumes when striping the root filesystem into 4MiB stripes.
# We use 4kb ext4 blocks. The stride and stripe should both be $STRIPE_SIZE / $EXT4_BLOCK_SIZE
ROOT_STRIDE=1024
ROOT_STRIPE_WIDTH=1024

case "${PARTITION_PLAN}" in
  split)
    truncate -s "${OS_IMAGE_SIZE_GIB}G" "${OS_IMAGE}"
    truncate -s "${DATA_IMAGE_SIZE_GIB}G" "${DATA_IMAGE}"
    ;;
  unified)
    truncate -s "$((OS_IMAGE_SIZE_GIB + DATA_IMAGE_SIZE_GIB))G" "${OS_IMAGE}"
    ;;
esac

declare -A partlabel parttype partguid partsize partoff
set_partition_sizes \
  "${OS_IMAGE_SIZE_GIB}" "${DATA_IMAGE_SIZE_GIB}" "${PARTITION_PLAN}" \
  partsize partoff
set_partition_labels partlabel
set_partition_types parttype
set_partition_uuids partguid "${PARTITION_PLAN}"

declare -a partargs
for part in \
  BIOS \
  EFI-A BOOT-A ROOT-A HASH-A RESERVED-A \
  EFI-B BOOT-B ROOT-B HASH-B RESERVED-B \
  PRIVATE DATA-A DATA-B ;
do
  # We create the DATA-B partition separately if we're using the split layout
  if [ "${part}" == "DATA-B" ] ; then
    continue
  fi

  # Each partition is aligned to a 1 MiB boundary, and extends to the sector
  # before the next partition starts. Specify the end point in sectors so we
  # can subtract a sector to fix the off-by-one error that comes from adding
  # start and size together. (1 MiB contains 2048 512-byte sectors.)
  part_start="${partoff[${part}]}"
  part_end="$((part_start + partsize[${part}]))"
  part_end="$((part_end * 2048 - 1))"

  partargs+=(-n "0:${part_start}M:${part_end}")
  partargs+=(-c "0:${partlabel[${part}]}")
  partargs+=(-t "0:${parttype[${part}]}")
  partargs+=(-u "0:${partguid[${part}]:-R}")

  # Boot partition attributes:
  #  48 = gptprio priority bit
  #  56 = gptprio successful bit
  case "${part}" in
    BOOT-A) partargs+=(-A 0:"set":48 -A 0:"set":56) ;;
    BOOT-B) partargs+=(-A 0:"clear":48 -A 0:"clear":56) ;;
  esac
done

sgdisk --clear "${partargs[@]}" --sort --print "${OS_IMAGE}"

# Partition the separate data disk, if we're using the split layout.
if [ "${PARTITION_PLAN}" == "split" ] ; then
  data_start="${partoff[DATA-B]}"
  data_end=$((data_start + partsize[DATA-B]))
  data_end=$((data_end * 2048 - 1))
  sgdisk --clear \
    -n "0:${data_start}M:${data_end}" \
    -c "0:${partlabel[DATA-B]}" \
    -t "0:${parttype[DATA-B]}" \
    -u "0:${partguid[DATA-B]}" \
    --sort --print "${DATA_IMAGE}"
fi

INSTALL_TIME="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
rpm -iv --ignorearch --root "${ROOT_MOUNT}" "${PACKAGE_DIR}"/*.rpm

# inventory installed packages
INVENTORY_QUERY="\{\"Name\":\"%{NAME}\"\
,\"Publisher\":\"Bottlerocket\"\
,\"Version\":\"${VERSION_ID}\"\
,\"Release\":\"${BUILD_ID}\"\
,\"InstalledTime\":\"${INSTALL_TIME}\"\
,\"ApplicationType\":\"%{GROUP}\"\
,\"Architecture\":\"%{ARCH}\"\
,\"Url\":\"%{URL}\"\
,\"Summary\":\"%{Summary}\"\}\n"

mapfile -t installed_rpms <<< "$(rpm -qa --root "${ROOT_MOUNT}" \
  --queryformat "${INVENTORY_QUERY}")"

# wrap installed_rpms mapfile into json
INVENTORY_DATA="$(jq --raw-output . <<<  "${installed_rpms[@]}")"
# remove the 'bottlerocket-' prefix from package names
INVENTORY_DATA="$(jq --arg PKG_PREFIX "bottlerocket-" \
                '(.Name) |= sub($PKG_PREFIX; "")' <<< "${INVENTORY_DATA}")"
# sort by package name and add 'Content' as top-level
INVENTORY_DATA="$(jq --slurp 'sort_by(.Name)' <<< "${INVENTORY_DATA}" | jq '{"Content": .}')"
printf "%s\n" "${INVENTORY_DATA}" > "${ROOT_MOUNT}/usr/share/bottlerocket/application-inventory.json"

# install licenses
install -p -m 0644 /host/{COPYRIGHT,LICENSE-APACHE,LICENSE-MIT} "${ROOT_MOUNT}"/usr/share/licenses/
mksquashfs \
  "${ROOT_MOUNT}"/usr/share/licenses \
  "${ROOT_MOUNT}"/usr/share/bottlerocket/licenses.squashfs \
  -no-exports -all-root -comp zstd
rm -rf "${ROOT_MOUNT}"/var/lib "${ROOT_MOUNT}"/usr/share/licenses/*

if [[ "${ARCH}" == "x86_64" ]]; then
  # MBR and BIOS-BOOT
  echo "(hd0) ${OS_IMAGE}" > "${ROOT_MOUNT}/boot/grub/device.map"
  "${ROOT_MOUNT}/sbin/grub-bios-setup" \
     --directory="${ROOT_MOUNT}/boot/grub" \
     --device-map="${ROOT_MOUNT}/boot/grub/device.map" \
     --root="hd0" \
     --skip-fs-probe \
     "${OS_IMAGE}"

  rm -vf "${ROOT_MOUNT}"/boot/grub/* "${ROOT_MOUNT}"/sbin/grub*
fi

# We also need an EFI partition, formatted FAT32 with the
# EFI binary at the correct path, e.g. /efi/boot. The grub
# package has placed the image in /boot/efi/EFI/BOOT.
mv "${ROOT_MOUNT}/boot/efi"/* "${EFI_MOUNT}"

# Do the setup required for `pesign` and `gpg` signing and
# verification to "just work" later on, regardless of which
# type of signing profile we have.
if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
  declare -a SHIM_SIGN_KEY
  declare -a CODE_SIGN_KEY

  # For an AWS profile, we expect a config file for the PKCS11
  # helper. Otherwise, there should be a local key and cert.
  if [ -s "${HOME}/.config/aws-kms-pkcs11/config.json" ] ; then
    # Set AWS environment variables from build secrets, if present.
    for var in AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN ; do
      val="${var,,}"
      val="${HOME}/.aws/${val//_/-}.env"
      [ -s "${val}" ] || continue
      declare -x "${var}=$(cat "${val}")"
    done
    # Verify that AWS credentials are functional.
    aws sts get-caller-identity
    # Log all PKCS11 helper activity, to simplify debugging.
    export AWS_KMS_PKCS11_DEBUG=1
    SB_KEY_SOURCE="aws"
    SHIM_SIGN_KEY=(-c shim-sign-key -t shim-sign-key)
    CODE_SIGN_KEY=(-c code-sign-key -t code-sign-key)
  else
    # Disable the PKCS11 helper.
    rm /etc/pkcs11/modules/aws-kms-pkcs11.module

    # Generate the PKCS12 archives for import.
    openssl pkcs12 \
      -export \
      -passout pass: \
      -inkey "${SBKEYS}/shim-sign.key" \
      -in "${SBKEYS}/shim-sign.crt" \
      -certfile "${SBKEYS}/db.crt" \
      -out "${SBKEYS}/shim-sign.p12"

    openssl pkcs12 \
      -export \
      -passout pass: \
      -inkey "${SBKEYS}/code-sign.key" \
      -in "${SBKEYS}/code-sign.crt" \
      -certfile "${SBKEYS}/vendor.crt" \
      -out "${SBKEYS}/code-sign.p12"

    # Import certificates and private key archive.
    PEDB="/etc/pki/pesign"

    certutil -d "${PEDB}" -A -n db -i "${SBKEYS}/db.crt" -t ",,C"
    certutil -d "${PEDB}" -A -n shim-sign-key -i "${SBKEYS}/shim-sign.crt" -t ",,P"
    pk12util -d "${PEDB}" -i "${SBKEYS}/shim-sign.p12" -W ""

    certutil -d "${PEDB}" -A -n vendor -i "${SBKEYS}/vendor.crt" -t ",,C"
    certutil -d "${PEDB}" -A -n code-sign-key -i "${SBKEYS}/code-sign.crt" -t ",,P"
    pk12util -d "${PEDB}" -i "${SBKEYS}/code-sign.p12" -W ""

    certutil -d "${PEDB}" -L
    SB_KEY_SOURCE="local"
    SHIM_SIGN_KEY=(-c shim-sign-key)
    CODE_SIGN_KEY=(-c code-sign-key)
  fi

  # Convert certificates from PEM format (ASCII) to DER (binary). This could be
  # done when the certificates are created, but the resulting binary files are
  # not as nice to store in source control.
  for cert in PK KEK db vendor ; do
    openssl x509 \
      -inform PEM -in "${SBKEYS}/${cert}.crt" \
      -outform DER -out "${SBKEYS}/${cert}.cer"
  done

  # For signing the grub config, we need to embed the GPG public key in binary
  # form, which is similarly awkward to store in source control.
  gpg --import "${SBKEYS}/config-sign.key"
  if [ "${SB_KEY_SOURCE}" == "aws" ] ; then
    gpg --card-status
  fi
  gpg --export > "${SBKEYS}/config-sign.pubkey"
  gpg --list-keys
fi

# shim expects the following data structure in `.vendor_cert`:
#
# struct {
#   uint32_t vendor_authorized_size;
#   uint32_t vendor_deauthorized_size;
#   uint32_t vendor_authorized_offset;
#   uint32_t vendor_deauthorized_offset;
# } cert_table;
#
cert_table() {
  local input output size offset uint32_t
  input="${1:?}"
  output="${2:?}"
  size="$(stat -c %s "${input}")"
  rm -f "${output}"
  # The cert payload is offset by four 4-byte uint32_t values in the header.
  offset="$((4 * 4))"
  for n in "${size}" 0 "${offset}" "$(( size + offset ))" ; do
    printf -v uint32_t '\\x%02x\\x%02x\\x%02x\\x%02x' \
      $((n & 255)) $((n >> 8 & 255)) $((n >> 16 & 255)) $((n >> 24 & 255))
    printf "${uint32_t}" >> "${output}"
  done
  cat "${input}" >> "${output}"
  # Zero-pad the output to the expected section size. Otherwise a subsequent
  # `objcopy` operation on the same section might fail to replace it, if the
  # new vendor certificate is larger than this one.
  truncate -s 4096 "${output}"
}

# Helper function to log the object layout before and after changes.
objdumpcopy() {
  local obj objdump objcopy
  obj="${1:?}"
  shift
  objdump="${ARCH}-bottlerocket-linux-gnu-objdump"
  objcopy="${ARCH}-bottlerocket-linux-gnu-objcopy"
  "${objdump}" -h "${obj}"
  "${objcopy}" "${@}" "${obj}"
  "${objdump}" -h "${obj}"
}

pushd "${EFI_MOUNT}/EFI/BOOT" >/dev/null
shims=(boot*.efi)
shim="${shims[0]}"
grubs=(grub*.efi)
grub="${grubs[0]}"
mokms=(mm*.efi)
mokm="${mokms[0]}"
if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
  # Convert the vendor certificate to the expected format.
  cert_table "${SBKEYS}/vendor.cer" "${SBKEYS}/vendor.obj"

  # Replace the embedded vendor certificate, then sign shim with the db key.
  objdumpcopy "${shim}" \
    --update-section ".vendor_cert=${SBKEYS}/vendor.obj"
  pesign -i "${shim}" -o "${shim}.signed" -s "${SHIM_SIGN_KEY[@]}"
  mv "${shim}.signed" "${shim}"
  pesigcheck -i "${shim}" -n 0 -c "${SBKEYS}/db.cer"

  # Sign the MOK manager as well.
  pesign -i "${mokm}" -o "${mokm}.signed" -s "${CODE_SIGN_KEY[@]}"
  mv "${mokm}.signed" "${mokm}"
  pesigcheck -i "${mokm}" -n 0 -c "${SBKEYS}/vendor.cer"

  # Replace the embedded gpg public key, then sign grub with the vendor key.
  objdumpcopy "${grub}" \
    --file-alignment 4096 \
    --update-section ".pubkey=${SBKEYS}/config-sign.pubkey"
  pesign -i "${grub}" -o "${grub}.signed" -s "${CODE_SIGN_KEY[@]}"
  mv "${grub}.signed" "${grub}"
  pesigcheck -i "${grub}" -n 0 -c "${SBKEYS}/vendor.cer"
else
  # Generate a zero-sized certificate in the expected format.
  cert_table /dev/null "${SBKEYS}/vendor.obj"

  # Replace the embedded vendor certificate with the zero-sized one, which shim
  # will ignore when Secure Boot is disabled.
  objdumpcopy "${shim}" \
    --update-section ".vendor_cert=${SBKEYS}/vendor.obj"

   # Remove the embedded gpg public key to disable GRUB's signature checks.
   objdumpcopy "${grub}" \
     --file-alignment 4096 \
     --remove-section ".pubkey"
fi
popd >/dev/null

dd if=/dev/zero of="${EFI_IMAGE}" bs=1M count="${partsize[EFI-A]}"
mkfs.vfat -I -S 512 "${EFI_IMAGE}" $((partsize[EFI-A] * 1024))
mmd -i "${EFI_IMAGE}" ::/EFI
mmd -i "${EFI_IMAGE}" ::/EFI/BOOT
mcopy -i "${EFI_IMAGE}" "${EFI_MOUNT}/EFI/BOOT"/*.efi ::/EFI/BOOT
if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
  # Make the signing certificate available on the EFI system partition so it
  # can be imported through the firmware setup UI on bare metal systems.
  mcopy -i "${EFI_IMAGE}" "${SBKEYS}"/db.{crt,cer} ::/EFI/BOOT
fi
dd if="${EFI_IMAGE}" of="${OS_IMAGE}" conv=notrunc bs=1M seek="${partoff[EFI-A]}"

# Ensure that the grub directory exists.
mkdir -p "${ROOT_MOUNT}/boot/grub"

# Now that we're done messing with /, move /boot out of it
mv "${ROOT_MOUNT}/boot"/* "${BOOT_MOUNT}"

if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
  pushd "${BOOT_MOUNT}" >/dev/null
  vmlinuz="vmlinuz"
  pesign -i "${vmlinuz}" -o "${vmlinuz}.signed" -s "${CODE_SIGN_KEY[@]}"
  mv "${vmlinuz}.signed" "${vmlinuz}"
  pesigcheck -i "${vmlinuz}" -n 0 -c "${SBKEYS}/vendor.cer"
  popd >/dev/null
fi

# Set the Bottlerocket variant, version, and build-id
SYS_ROOT="${ARCH}-bottlerocket-linux-gnu/sys-root"
VERSION="${VERSION_ID} (${VARIANT})"
cat <<EOF >> "${ROOT_MOUNT}/${SYS_ROOT}/usr/lib/os-release"
VERSION="${VERSION}"
PRETTY_NAME="${PRETTY_NAME} ${VERSION}"
VARIANT_ID=${VARIANT}
VERSION_ID=${VERSION_ID}
BUILD_ID=${BUILD_ID}
HOME_URL="https://github.com/bottlerocket-os/bottlerocket"
SUPPORT_URL="https://github.com/bottlerocket-os/bottlerocket/discussions"
BUG_REPORT_URL="https://github.com/bottlerocket-os/bottlerocket/issues"
EOF

# Set the BOTTLEROCKET-DATA Filesystem for creating/mounting
if [ "${XFS_DATA_PARTITION}" == "yes" ] ; then
  printf "%s\n" "DATA_PARTITION_FILESYSTEM=xfs" >> "${ROOT_MOUNT}/${SYS_ROOT}/usr/share/bottlerocket/image-features.env"
else
  printf "%s\n" "DATA_PARTITION_FILESYSTEM=ext4" >> "${ROOT_MOUNT}/${SYS_ROOT}/usr/share/bottlerocket/image-features.env"
fi

# BOTTLEROCKET-ROOT-A
mkdir -p "${ROOT_MOUNT}/lost+found"
ROOT_LABELS=$(setfiles -n -d -F -m -r "${ROOT_MOUNT}" \
    "${SELINUX_FILE_CONTEXTS}" "${ROOT_MOUNT}" \
    | awk -v root="${ROOT_MOUNT}" '{gsub(root"/","/"); gsub(root,"/"); print "ea_set", $1, "security.selinux", $4}')
mkfs.ext4 -E "lazy_itable_init=0,stride=${ROOT_STRIDE},stripe_width=${ROOT_STRIPE_WIDTH}" \
  -O ^has_journal -b "${VERITY_DATA_BLOCK_SIZE}" -d "${ROOT_MOUNT}" "${ROOT_IMAGE}" "${partsize[ROOT-A]}M"
echo "${ROOT_LABELS}" | debugfs -w -f - "${ROOT_IMAGE}"
resize2fs -M "${ROOT_IMAGE}"
dd if="${ROOT_IMAGE}" of="${OS_IMAGE}" conv=notrunc bs=1M seek="${partoff[ROOT-A]}"

# BOTTLEROCKET-VERITY-A
veritypart_mib="${partsize[HASH-A]}"
truncate -s "${veritypart_mib}M" "${VERITY_IMAGE}"
veritysetup_output="$(veritysetup format \
    --format "$VERITY_VERSION" \
    --hash "$VERITY_HASH_ALGORITHM" \
    --data-block-size "$VERITY_DATA_BLOCK_SIZE" \
    --hash-block-size "$VERITY_HASH_BLOCK_SIZE" \
    "${ROOT_IMAGE}" "${VERITY_IMAGE}" \
    | tee /dev/stderr)"
verityimage_size="$(stat -c %s "${VERITY_IMAGE}")"
veritypart_bytes="$((veritypart_mib * 1024 * 1024))"
if [ "${verityimage_size}" -gt "${veritypart_bytes}" ] ; then
    echo "verity content is larger than partition (${veritypart_mib}M)"
    exit 1
fi
VERITY_DATA_4K_BLOCKS="$(grep '^Data blocks:' <<<"${veritysetup_output}" | awk '{ print $NF }')"
VERITY_DATA_512B_BLOCKS="$((VERITY_DATA_4K_BLOCKS * 8))"
VERITY_ROOT_HASH="$(grep '^Root hash:' <<<"${veritysetup_output}" | awk '{ print $NF }')"
VERITY_SALT="$(grep '^Salt:' <<<"${veritysetup_output}" | awk '{ print $NF }')"
veritysetup verify "${ROOT_IMAGE}" "${VERITY_IMAGE}" "${VERITY_ROOT_HASH}"
dd if="${VERITY_IMAGE}" of="${OS_IMAGE}" conv=notrunc bs=1M seek="${partoff[HASH-A]}"

declare -a DM_VERITY_ROOT
DM_VERITY_ROOT=(
  "root,,,ro,0"
  "${VERITY_DATA_512B_BLOCKS}"
  "verity"
  "${VERITY_VERSION}"
  "PARTUUID=\$boot_uuid/PARTNROFF=1"
  "PARTUUID=\$boot_uuid/PARTNROFF=2"
  "${VERITY_DATA_BLOCK_SIZE}"
  "${VERITY_HASH_BLOCK_SIZE}"
  "${VERITY_DATA_4K_BLOCKS}"
  "1"
  "${VERITY_HASH_ALGORITHM}"
  "${VERITY_ROOT_HASH}"
  "${VERITY_SALT}"
  "2"
  "restart_on_corruption"
  "ignore_zero_blocks"
)

# write GRUB config
# If GRUB_SET_PRIVATE_VAR is set, include the parameters that support Boot Config
if [ "${GRUB_SET_PRIVATE_VAR}" == "yes" ] ; then
   BOOTCONFIG='bootconfig'
   INITRD="initrd (\$private)/bootconfig.data"
else
   BOOTCONFIG=""
   INITRD=""
fi

# If UEFI_SECURE_BOOT is set, disable interactive edits. Otherwise the intended
# kernel command line parameters could be changed if the boot fails. Disable
# signature checking as well, since grub.cfg will have already been verified
# before we reach this point. bootconfig.data is generated at runtime and can't
# be signed with a trusted key, so continuing to check signatures would prevent
# it from being read. If boot fails, trigger an automatic reboot, since nothing
# can be changed for troubleshooting purposes.
if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
   echo 'set superusers=""' > "${BOOT_MOUNT}/grub/grub.cfg"
   echo 'set check_signatures="no"' >> "${BOOT_MOUNT}/grub/grub.cfg"
   FALLBACK=$'   echo "rebooting in 30 seconds..."\n'
   FALLBACK+=$'   sleep 30\n'
   FALLBACK+=$'   reboot\n'
else
   FALLBACK=""
fi

cat <<EOF >> "${BOOT_MOUNT}/grub/grub.cfg"
set default="0"
set timeout="0"
set dm_verity_root="${DM_VERITY_ROOT[@]}"

menuentry "${PRETTY_NAME} ${VERSION_ID}" --unrestricted {
   linux (\$root)/vmlinuz \\
       ${KERNEL_PARAMETERS} \\
       ${BOOTCONFIG} \\
       root=/dev/dm-0 rootwait ro \\
       raid=noautodetect \\
       random.trust_cpu=on \\
       selinux=1 enforcing=1 \\
       dm-mod.create="\$dm_verity_root" \\
       -- \\
       systemd.log_target=journal-or-kmsg \\
       systemd.log_color=0 \\
       systemd.show_status=true
   ${INITRD}
   boot
   ${FALLBACK}
}
EOF

if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
  gpg --detach-sign "${BOOT_MOUNT}/grub/grub.cfg"
  gpg --verify "${BOOT_MOUNT}/grub/grub.cfg.sig"
fi

# BOTTLEROCKET-BOOT-A
mkdir -p "${BOOT_MOUNT}/lost+found"
chmod -R go-rwx "${BOOT_MOUNT}"
BOOT_LABELS=$(setfiles -n -d -F -m -r "${BOOT_MOUNT}" \
    "${SELINUX_FILE_CONTEXTS}" "${BOOT_MOUNT}" \
  | awk -v root="${BOOT_MOUNT}" '{gsub(root"/","/"); gsub(root,"/"); print "ea_set", $1, "security.selinux", $4}')
mkfs.ext4 -O ^has_journal -d "${BOOT_MOUNT}" "${BOOT_IMAGE}" "${partsize[BOOT-A]}M"
echo "${BOOT_LABELS}" | debugfs -w -f - "${BOOT_IMAGE}"
resize2fs -M "${BOOT_IMAGE}"
dd if="${BOOT_IMAGE}" of="${OS_IMAGE}" conv=notrunc bs=1M seek="${partoff[BOOT-A]}"

# BOTTLEROCKET-PRIVATE

# Generate an empty bootconfig file for the image, so grub doesn't pause and
# print an error that the file doesn't exist.
cat <<EOF > "${PRIVATE_MOUNT}/bootconfig.in"
kernel {}
init {}
EOF
touch "${PRIVATE_MOUNT}/bootconfig.data"
bootconfig -a "${PRIVATE_MOUNT}/bootconfig.in" "${PRIVATE_MOUNT}/bootconfig.data"
rm "${PRIVATE_MOUNT}/bootconfig.in"

# Targeted toward the current API server implementation.
# Relative to the ext4 defaults, we:
# - adjust the inode ratio since we expect lots of small files
# - retain the inode size to allow most settings to be stored inline
# - retain the block size to handle worse-case alignment for hardware
mkfs.ext4 -b 4096 -i 4096 -I 256 -d "${PRIVATE_MOUNT}" "${PRIVATE_IMAGE}" "${partsize[PRIVATE]}M"
dd if="${PRIVATE_IMAGE}" of="${OS_IMAGE}" conv=notrunc bs=1M seek="${partoff[PRIVATE]}"

# BOTTLEROCKET-DATA-A and BOTTLEROCKET-DATA-B

# If we build on a host with SELinux enabled, we could end up with labels that
# do not match our policy. Since we allow replacing the data volume at runtime,
# we can't count on these labels being correct in any case, and it's better to
# remove them all.
UNLABELED=$(find "${DATA_MOUNT}" \
    | awk -v root="${DATA_MOUNT}" '{gsub(root"/","/"); gsub(root,"/"); print "ea_rm", $1, "security.selinux"}')

mkfs_data() {
  local target size offset
  target="${1:?}"
  size="${2:?}"
  offset="${3:?}"
  # Create an XFS filesystem if requested
  if [ "${XFS_DATA_PARTITION}" == "yes" ] ; then
    echo "writing XFS filesystem for DATA"
    # Create a file to write the filesystem to first
    dd if=/dev/zero of="${BOTTLEROCKET_DATA}" bs=1M count=${size%?}
    # block size of 4096, directory block size of 16384
    # enable inotbtcount, bigtime, and reflink
    # use an internal log with starting size of 64m
    # use the minimal 2 Allocation groups, this still overprovisions when expanded
    # set strip units of 512k and sectsize to make EBS volumes align
    mkfs.xfs \
      -b size=4096 -n size=16384 \
      -m inobtcount=1,bigtime=1,reflink=1 \
      -l internal,size=64m \
      -d agcount=2,su=512k,sw=1,sectsize=4096 \
      -f "${BOTTLEROCKET_DATA}"
  else
    # default to ext4
    echo "writing ext4 filesystem for DATA"
    mkfs.ext4 -m 0 -d "${DATA_MOUNT}" "${BOTTLEROCKET_DATA}" "${size}"
    echo "${UNLABELED}" | debugfs -w -f - "${BOTTLEROCKET_DATA}"
  fi
  dd if="${BOTTLEROCKET_DATA}" of="${target}" conv=notrunc bs=1M seek="${offset}"
}

# Decide which data filesystem to create at build time based on layout.
#
# The DATA-A partition will always exist, but for the "split" layout, it will be
# too small to provide the desired filesystem parameters (inode count, etc) when
# it is grown later on. Hence this filesystem is only created for "unified".
#
# The DATA-B partition does not exist on the "unified" layout, which anticipates
# a single storage device. Hence this filesystem is only created for "split".
#
# If the other partition is available at runtime, the filesystem will be created
# during first boot instead, providing flexibility at the cost of a minor delay.
case "${PARTITION_PLAN}" in
  unified)
    mkfs_data "${OS_IMAGE}" "${partsize["DATA-A"]}M" "${partoff["DATA-A"]}"
    ;;
  split)
    mkfs_data "${DATA_IMAGE}" "${partsize["DATA-B"]}M" "${partoff["DATA-B"]}"
    ;;
esac

sgdisk -v "${OS_IMAGE}"
[ -s "${DATA_IMAGE}" ] && sgdisk -v "${DATA_IMAGE}"

symlink_image() {
  local ext what
  ext="${1}"
  what="${2}"
  ext="${ext:+.$ext}"
  target="${what^^}_NAME"
  for link in symlink versioned_symlink friendly_versioned_symlink ; do
    link="${what^^}_${link^^}"
    ln -s "${!target}${ext}" "${OUTPUT_DIR}/${!link}${ext}"
  done
}

if [[ ${OUTPUT_FMT} == "raw" ]]; then
  lz4 -vc "${OS_IMAGE}" >"${OUTPUT_DIR}/${OS_IMAGE_NAME}.img.lz4"
  symlink_image "img.lz4" "os_image"
  if [ -s "${DATA_IMAGE}" ] ; then
    lz4 -vc "${DATA_IMAGE}" >"${OUTPUT_DIR}/${DATA_IMAGE_NAME}.img.lz4"
    symlink_image "img.lz4" "data_image"
  fi
elif [[ ${OUTPUT_FMT} == "qcow2" ]]; then
  qemu-img convert -f raw -O qcow2 "${OS_IMAGE}" "${OUTPUT_DIR}/${OS_IMAGE_NAME}.qcow2"
  symlink_image "qcow2" "os_image"
  if [ -s "${DATA_IMAGE}" ] ; then
    qemu-img convert -f raw -O qcow2 "${DATA_IMAGE}" "${OUTPUT_DIR}/${DATA_IMAGE_NAME}.qcow2"
    symlink_image "qcow2" "data_image"
  fi
elif [[ ${OUTPUT_FMT} == "vmdk" ]]; then
  # Stream optimization is required for creating an Open Virtual Appliance (OVA)
  qemu-img convert -f raw -O vmdk -o subformat=streamOptimized "${OS_IMAGE}" "${OUTPUT_DIR}/${OS_IMAGE_NAME}.vmdk"
  symlink_image "vmdk" "os_image"
  if [ -s "${DATA_IMAGE}" ] ; then
    qemu-img convert -f raw -O vmdk -o subformat=streamOptimized "${DATA_IMAGE}" "${OUTPUT_DIR}/${DATA_IMAGE_NAME}.vmdk"
    symlink_image "vmdk" "data_image"
  fi
fi

# Now create the OVA if needed.
if [ "${OUTPUT_FMT}" == "vmdk" ] ; then
  os_vmdk="${OS_IMAGE_NAME}.vmdk"
  data_vmdk="${DATA_IMAGE_NAME}.vmdk"
  ovf="${OS_IMAGE_NAME}.ovf"
  ova_dir="$(mktemp -d)"

  # The manifest expects disk sizes in bytes.
  bytes_in_gib="$((1024 * 1024 * 1024))"
  os_disk_bytes="$((OS_IMAGE_PUBLISH_SIZE_GIB * bytes_in_gib))"
  data_disk_bytes="$((DATA_IMAGE_PUBLISH_SIZE_GIB * bytes_in_gib))"
  sed "${OVF_TEMPLATE}" \
     -e "s/{{OS_DISK}}/${os_vmdk}/g" \
     -e "s/{{DATA_DISK}}/${data_vmdk}/g" \
     -e "s/{{OS_DISK_BYTES}}/${os_disk_bytes}/g" \
     -e "s/{{DATA_DISK_BYTES}}/${data_disk_bytes}/g" \
     > "${ova_dir}/${ovf}"

  # The manifest templates for Secure Boot expect the cert data for
  # PK, KEK, db, and dbx.
  if [ "${UEFI_SECURE_BOOT}" == "yes" ] ; then
    pk_cert_der_hex="$(hexdump -ve '1/1 "%02x"' "${SBKEYS}/PK.cer")"
    kek_cert_der_hex="$(hexdump -ve '1/1 "%02x"' "${SBKEYS}/KEK.cer")"
    db_cert_der_hex="$(hexdump -ve '1/1 "%02x"' "${SBKEYS}/db.cer")"
    dbx_empty_hash_hex="$(sha256sum /dev/null | awk '{ print $1 }')"
    sed -i \
      -e "s/{{PK_CERT_DER_HEX}}/${pk_cert_der_hex}/g" \
      -e "s/{{KEK_CERT_DER_HEX}}/${kek_cert_der_hex}/g" \
      -e "s/{{DB_CERT_DER_HEX}}/${db_cert_der_hex}/g" \
      -e "s/{{DBX_EMPTY_HASH_HEX}}/${dbx_empty_hash_hex}/g" \
      "${ova_dir}/${ovf}"
  fi

  # Make sure we replaced all the '{{...}}' fields with real values.
  if grep -F -e '{{' -e '}}' "${ova_dir}/${ovf}" ; then
    echo "Failed to fully render the OVF template" >&2
    exit 1
  fi

  # Create the manifest file with the hashes of the VMDKs and the OVF.
  manifest="${OS_IMAGE_NAME}.mf"
  pushd "${OUTPUT_DIR}" >/dev/null
  os_sha256="$(sha256sum ${os_vmdk} | awk '{print $1}')"
  echo "SHA256(${os_vmdk})= ${os_sha256}" > "${ova_dir}/${manifest}"
  if [ -s "${DATA_IMAGE}" ] ; then
    data_sha256="$(sha256sum ${data_vmdk} | awk '{print $1}')"
    echo "SHA256(${data_vmdk})= ${data_sha256}" >> "${ova_dir}/${manifest}"
  fi
  popd >/dev/null
  pushd "${ova_dir}" >/dev/null
  ovf_sha256="$(sha256sum ${ovf} | awk '{print $1}')"
  echo "SHA256(${ovf})= ${ovf_sha256}" >> "${manifest}"
  popd >/dev/null

  # According to the OVF spec:
  # https://www.dmtf.org/sites/default/files/standards/documents/DSP0243_2.1.1.pdf,
  # the OVF must be first in the tar bundle.  Manifest is next, and then the
  # files must fall in the same order as listed in the References section of the
  # OVF file
  ova="${OS_IMAGE_NAME}.ova"
  tar -cf "${OUTPUT_DIR}/${ova}" -C "${ova_dir}" "${ovf}" "${manifest}"
  tar -rf "${OUTPUT_DIR}/${ova}" -C "${OUTPUT_DIR}" "${os_vmdk}"
  if [ -s "${DATA_IMAGE}" ] ; then
     tar -rf "${OUTPUT_DIR}/${ova}" -C "${OUTPUT_DIR}" "${data_vmdk}"
  fi

  symlink_image "ova" "os_image"
fi

lz4 -9vc "${BOOT_IMAGE}" >"${OUTPUT_DIR}/${BOOT_IMAGE_NAME}"
lz4 -9vc "${VERITY_IMAGE}" >"${OUTPUT_DIR}/${VERITY_IMAGE_NAME}"
lz4 -9vc "${ROOT_IMAGE}" >"${OUTPUT_DIR}/${ROOT_IMAGE_NAME}"

symlink_image "" "boot_image"
symlink_image "" "verity_image"
symlink_image "" "root_image"

find "${OUTPUT_DIR}" -type f -print -exec chown 1000:1000 {} \;

# Clean up temporary files to reduce size of layer.
rm -f "${PACKAGE_DIR}"/*.rpm
rm -rf /tmp/*