#!/bin/bash

# Copyright (c) 2023 The Flatcar Maintainers.
# Use of this source code is governed by a BSD-style license that can
# be found in the LICENSE file.
#
# Script to generate sysext. See systemd-sysext(8). Prerequisite is
# that you've run build_packages and build_image.


SCRIPT_ROOT=$(dirname "$(readlink -f "$0")")
. "${SCRIPT_ROOT}/common.sh" || exit 1

# Script must run inside the chroot
assert_inside_chroot
assert_root_user

default_imagedir="$(readlink -f "${SCRIPT_ROOT}/../build/images")/<BOARD>/latest/"
default_install_root_basename='install-root'

# All these are used to set up the 'BUILD_DIR' variable
DEFINE_string board "${DEFAULT_BOARD}" \
  "The board to build a sysext for."
DEFINE_string metapkgs '' \
  "Comma-separated list of meta-packages to build from source and install into sysext image."
DEFINE_string squashfs_base '' \
  "The path to the squashfs base image. Defaults to the most current image built in '${default_imagedir}/${FLATCAR_PRODUCTION_IMAGE_SYSEXT_BASE}'."
DEFINE_string image_builddir '' \
  "Custom directory to build the sysext in. Defaults to a 'sysext' sub-directory of the directory the squashfs base image resides in; '${default_imagedir}/sysext' by default."
DEFINE_boolean strip_binaries "${FLAGS_FALSE}" \
  "After installation, scan sysext root for unstripped binaries and strip these. WARNING - this can subtly break some packages, e.g. Docker (see https://github.com/moby/moby/blob/master/project/PACKAGERS.md#stripping-binaries)."
DEFINE_string manglefs_script '' \
  "A path to executable that will customize the rootfs of the sysext image."
DEFINE_boolean generate_pkginfo "${FLAGS_FALSE}" \
  "Generate an additional squashfs '<sysext_name>_pkginfo.raw' with portage package meta-information (/var/db ...). Useful for creating sysext dependencies; see 'base_pkginfo' below."
DEFINE_string base_pkginfo "" \
  "Colon-separated list of pkginfo squashfs paths / files generated via 'generate_pkginfo' to base this sysext on. The corresponding base sysexts are expected to be merged with the sysext generated."
DEFINE_string compression "lz4hc" \
  "Compression to use for sysext EROFS image. Options: 'lz4', 'lz4hc', 'zstd', or 'none'. Default is 'lz4hc'."
DEFINE_string mkerofs_opts "" \
  "Additional mkfs.erofs options to pass via SYSTEMD_REPART_MKFS_OPTIONS_EROFS. If not specified, defaults are used based on compression type."
DEFINE_boolean ignore_version_mismatch "${FLAGS_FALSE}" \
  "Ignore version mismatch between SDK board packages and base squashfs. DANGEROUS."
DEFINE_string install_root_basename "${default_install_root_basename}" \
  "Name of a root directory where packages will be installed. ${default_install_root_basename@Q} by default."

FLAGS_HELP="USAGE: build_sysext [flags] <sysext_name> <binary_package> [<binary_package> ...]

This script is used to build a Flatcar sysext image.
The sysext will be based on an OS image build's sysext base squashfs, i.e. it is specific to
a Flatcar build or release.
The base squashfs can either come from a local build or downloaded from an official release.
By default, the sysext will be built in a 'sysext' sub-dir of the directory the squashfs base image
is in, but this can be changed with the --image_builddir option.

Examples:

Builds a sysext image named 'interpreters' with 'dev-lang/python' and 'dev-lang/perl' packages for the
most recent production image (default architecture, likely amd64) in the defaut build directory:

sudo build_sysext \\
     interpreters dev-lang/python dev-lang/perl


Builds a sysext image named 'oem-azure' in the 'oem-images' sub-directory with
metapackage 'coreos-base/oem-azure' for the arm64 squashfs base at
'build/artifacts/flatcar_production_image_sysext.squashfs':

sudo build_sysext \\
     --board=arm64-usr \\
     --metapkgs=coreos-base/oem-azure \\
     --mangle_fs=sdk_container/src/third_party/coreos-overlay/coreos-base/oem-azure/files/manglefs.sh \\
     --squashfs_base=build/artifacts/flatcar_production_image_sysext.squashfs \\
     --image_builddir=oem-images \\
     oem-azure


Mandatory command line parameters:
  <sysext_name>     - name of the sysext output file.
  <binary_package>  - List of existing binary packages to install. Can be omitted if --metapkgs was specified.
"

show_help_if_requested "$@"

# Parse command line
FLAGS "$@" || exit 1

eval set -- "${FLAGS_ARGV}"

# Only now can we die on error.  shflags functions leak non-zero error codes,
# so will die prematurely if 'switch_to_strict_mode' is specified before now.
switch_to_strict_mode -uo pipefail

# Validate command line parameters

SYSEXTNAME="${1:-}"
if [[ -z "${SYSEXTNAME}" ]]; then
  die "No sysext name provided."
fi
shift

if [[ -z "$FLAGS_squashfs_base" ]] ; then
  FLAGS_squashfs_base="$(readlink -f "${SCRIPT_ROOT}/../build/images/${FLAGS_board}/latest/${FLATCAR_PRODUCTION_IMAGE_SYSEXT_BASE}")"
fi
if [[ ! -f "${FLAGS_squashfs_base}" ]] ; then
  die "Squashfs base '${FLAGS_squashfs_base}' not found."
fi

if [[ -z "${FLAGS_image_builddir}" ]]; then
  FLAGS_image_builddir="$(dirname "${FLAGS_squashfs_base}")/sysext"
fi
BUILD_DIR=$(realpath "${FLAGS_image_builddir}")
mkdir -p "${BUILD_DIR}"

source "${BUILD_LIBRARY_DIR}/toolchain_util.sh" || exit 1
source "${BUILD_LIBRARY_DIR}/board_options.sh" || exit 1
source "${BUILD_LIBRARY_DIR}/reports_util.sh" || exit 1

# Architecture values are taken from systemd.unit(5).
declare -A SYSEXT_ARCHES
SYSEXT_ARCHES['amd64-usr']='x86-64'
SYSEXT_ARCHES['arm64-usr']='arm64'

declare -r SYSEXT_ARCHES

# Usage: _get_sysext_arch board [board...]
_get_sysext_arch() {
    local board=${1}
    if [[ ${#SYSEXT_ARCHES["${board}"]} -ne 0 ]]; then
        echo "${SYSEXT_ARCHES["${board}"]}"
    else
        die "Unknown board '${board}'"
    fi
}

cleanup() {
  local dirs=(
    "${BUILD_DIR}/fs-root"
    "${BUILD_DIR}/${FLAGS_install_root_basename}"
    "${BUILD_DIR}/workdir"
    "${BUILD_DIR}/img-rootfs"
  )
  umount "${dirs[@]}" 2>/dev/null || true
  rm -rf "${dirs[@]}" || true
  if [[ -d "${BUILD_DIR}/base-pkginfo" ]] ; then
    umount "${BUILD_DIR}/base-pkginfo"/* 2>/dev/null || true
    rm -rf "${BUILD_DIR}/base-pkginfo" || true
  fi
  rm -rf "${BUILD_DIR}/img-pkginfo"
}

# Set up trap to execute cleanup() on script exit
trap cleanup EXIT

ARCH=$(_get_sysext_arch "${FLAGS_board}")
cleanup

# If we need to handle pkginfo squashfs files, create mount points under
# ${BUILD_DIR}/base-pkginfo, mount the squashfs images, and add the mount paths to
# the list of lowerdirs.
pkginfo_lowerdirs=""
if [[ -n "${FLAGS_base_pkginfo}" ]] ; then
  for entry in $(echo ${FLAGS_base_pkginfo} | sed 's/:/ /g'); do
    ppath="$(readlink -f "${entry}")"
    if [[ ! -f "${ppath}" ]] ; then
      error "--base_pkginfo contains invalid entries."
      error "Pkginfo file '${ppath}' does not exist."
      die "Full --base_pkginfo: '${FLAGS_base_pkginfo}'"
    fi

    pfile="$(basename "${ppath}")"
    pmdir="${BUILD_DIR}/base-pkginfo/${pfile}"
    mkdir -p "${pmdir}"
    mount -rt squashfs -o loop,nodev "${ppath}" "${pmdir}"
    pkginfo_lowerdirs="${pkginfo_lowerdirs}:${pmdir}"
    info "Added packageinfo from '${ppath}' to base layers."
  done
fi

mkdir "${BUILD_DIR}/fs-root"
mount -rt squashfs -o loop,nodev "${FLAGS_squashfs_base}" "${BUILD_DIR}/fs-root"
mkdir "${BUILD_DIR}/${FLAGS_install_root_basename}"
mkdir "${BUILD_DIR}/workdir"
mount -t overlay overlay -o lowerdir="${BUILD_DIR}/fs-root${pkginfo_lowerdirs}",upperdir="${BUILD_DIR}/${FLAGS_install_root_basename}",workdir="${BUILD_DIR}/workdir" "${BUILD_DIR}/${FLAGS_install_root_basename}"

REPO_BUILD_ID=$(source "${REPO_MANIFESTS_DIR}/version.txt"; echo "$FLATCAR_BUILD_ID")
REPO_FLATCAR_VERSION=$(source "${REPO_MANIFESTS_DIR}/version.txt"; echo "$FLATCAR_VERSION")
VERSION_BOARD=$(source "${BUILD_DIR}/fs-root/usr/lib/os-release" && echo "$VERSION")

if [[ -z $REPO_BUILD_ID ]] && [[ ${COREOS_OFFICIAL:-0} -ne 1 ]]; then
    BASE_SQUASHFS_BUILD_ID=$(source "${BUILD_DIR}/fs-root/usr/lib/os-release" && echo -n "$BUILD_ID")
    info "This is a dev rebuild of an official release tag: No BUILD ID set in '${REPO_MANIFESTS_DIR}/version.txt'.  Will use base squashfs BUILD ID for version check."
    info "Repo root FLATCAR_VERSION is '$REPO_FLATCAR_VERSION', squashfs build ID is '$BASE_SQUASHFS_BUILD_ID'"
    FLATCAR_VERSION="${REPO_FLATCAR_VERSION}${BASE_SQUASHFS_BUILD_ID:++}${BASE_SQUASHFS_BUILD_ID}"
    info "Setting FLATCAR_VERSION to '$FLATCAR_VERSION'"
fi

if [ "$VERSION_BOARD" != "$FLATCAR_VERSION" ]; then
  warn "Base squashfs version: $VERSION_BOARD"
  warn "SDK board packages version: $FLATCAR_VERSION"
  if [[ "${FLAGS_ignore_version_mismatch}" = "${FLAGS_TRUE}" ]] ; then
    warn "Ignoring version mismatch as requested."
  else
    die "Version mismatch between board flatcar release and SDK container flatcar release."
  fi
fi

if [[ -n "${FLAGS_metapkgs}" ]]; then
  mapfile -t metapkgs < <(tr ',' '\n' <<<"${FLAGS_metapkgs}")
  "emerge-${FLAGS_board}" --nodeps --buildpkgonly --usepkg n --verbose "${metapkgs[@]}"
  set -- "${metapkgs[@]}" "${@}"
fi

if [[ ${#} -lt 1 ]]; then
   error 'No packages or meta packages to install.'
   show_help_if_requested -h
fi

info "Building '${SYSEXTNAME}' sysext with (meta-)packages '${@}' in '${BUILD_DIR}' using '${FLAGS_compression}' compression".

for package; do
  echo "Installing package into sysext image: $package"
  FEATURES="-ebuild-locks binpkg-multi-instance" emerge \
    --root="${BUILD_DIR}/${FLAGS_install_root_basename}" \
    --config-root="/build/${FLAGS_board}"  \
    --sysroot="/build/${FLAGS_board}"  \
    --usepkgonly \
    --binpkg-respect-use=y \
    --getbinpkg \
    --verbose \
    --jobs=${NUM_JOBS} \
    "${package}"
done

# Make squashfs generation more reproducible.
export SOURCE_DATE_EPOCH=$(stat -c '%Y' "${BUILD_DIR}/fs-root/usr/lib/os-release")

# Unmount in order to get rid of the overlay
umount "${BUILD_DIR}/${FLAGS_install_root_basename}"
umount "${BUILD_DIR}/fs-root"

if [[ "$FLAGS_generate_pkginfo" = "${FLAGS_TRUE}" ]] ; then
  info "  Creating pkginfo squashfs '${BUILD_DIR}/${SYSEXTNAME}_pkginfo.raw'"
  mkdir -p "${BUILD_DIR}/img-pkginfo/var/db"
  cp -R "${BUILD_DIR}/${FLAGS_install_root_basename}/var/db/pkg" "${BUILD_DIR}/img-pkginfo/var/db/"
  mksquashfs "${BUILD_DIR}/img-pkginfo" "${BUILD_DIR}/${SYSEXTNAME}_pkginfo.raw" \
              -noappend -xattrs-exclude '^btrfs.' -comp zstd -Xcompression-level 22 -b 512k
fi

info "Writing ${SYSEXTNAME}_packages.txt"
ROOT="${BUILD_DIR}/${FLAGS_install_root_basename}" PORTAGE_CONFIGROOT="${BUILD_DIR}/${FLAGS_install_root_basename}" \
      equery --no-color list --format '$cpv::$repo' '*' > "${BUILD_DIR}/${SYSEXTNAME}_packages.txt"


if [[ "${FLAGS_strip_binaries}" = "${FLAGS_TRUE}" ]]; then
    chost="$("portageq-${BOARD}" envvar CHOST)"
    strip="${chost}-strip"

    info "Stripping all non-stripped binaries in sysext using '${strip}'"

    # Find all non-stripped binaries, remove ':' from filepath, and strip 'em
    find "${BUILD_DIR}/${FLAGS_install_root_basename}" -exec file \{\} \; \
         | awk '/not stripped/ {print substr($1, 1, length($1)-1)}' \
         | while read bin; do
                info "     ${strip} ${bin}"
                "${strip}" "${bin}"
           done
fi

if [[ -n "${FLAGS_manglefs_script}" ]]; then
  if [[ ! -x "${FLAGS_manglefs_script}" ]]; then
    die "${FLAGS_manglefs_script} is not executable"
  fi
  "${FLAGS_manglefs_script}" "${BUILD_DIR}/${FLAGS_install_root_basename}"
fi

info "Removing non-/usr directories from sysext image"
for entry in "${BUILD_DIR}/${FLAGS_install_root_basename}"/*; do
  if [[ "${entry}" = */usr ]]; then
    continue
  fi
  info "  Removing ${entry##*/}"
  rm -rf "${entry}"
done
mkdir -p "${BUILD_DIR}/${FLAGS_install_root_basename}/usr/lib/extension-release.d"
version_field="${VERSION_FIELD_OVERRIDE:-VERSION_ID=${FLATCAR_VERSION_ID}}"
all_fields=(
  'ID=flatcar'
  "${version_field}"
  "ARCHITECTURE=${ARCH}"
)
printf '%s\n' "${all_fields[@]}" >"${BUILD_DIR}/${FLAGS_install_root_basename}/usr/lib/extension-release.d/extension-release.${SYSEXTNAME}"

info "Removing opaque directory markers to always merge all contents"
find "${BUILD_DIR}/${FLAGS_install_root_basename}" -xdev -type d -exec sh -c 'if [ "$(attr -R -q -g overlay.opaque {} 2>/dev/null)" = y ]; then attr -R -r overlay.opaque {}; fi' \;

info "Checking for invalid file ownership"
invalid_files=$(find "${BUILD_DIR}/${FLAGS_install_root_basename}" -user sdk -or -group sdk)
if [[ -n "${invalid_files}" ]]; then
  die "Invalid file ownership: ${invalid_files}"
fi

# Set up EROFS compression options based on compression type
if [[ "${FLAGS_compression}" != "none" ]]; then
  export SYSTEMD_REPART_MKFS_OPTIONS_EROFS="-z${FLAGS_compression}"

  if [[ -n "${FLAGS_mkerofs_opts}" ]]; then
    # User provided custom options
    export SYSTEMD_REPART_MKFS_OPTIONS_EROFS="${SYSTEMD_REPART_MKFS_OPTIONS_EROFS} ${FLAGS_mkerofs_opts}"
  elif [[ "${FLAGS_compression}" = "lz4hc" ]]; then
    # Default options for lz4hc
    export SYSTEMD_REPART_MKFS_OPTIONS_EROFS="${SYSTEMD_REPART_MKFS_OPTIONS_EROFS},12 -C65536 -Efragments,ztailpacking"
  elif [[ "${FLAGS_compression}" = "zstd" ]]; then
    # Default options for zstd
    export SYSTEMD_REPART_MKFS_OPTIONS_EROFS="${SYSTEMD_REPART_MKFS_OPTIONS_EROFS},level=22 -C524288 -Efragments,ztailpacking"
  fi
  info "Building sysext with ${FLAGS_compression} compression"
else
  info "Building sysext without compression (built-in sysexts)"
fi

systemd-repart \
  --private-key="${SYSEXT_SIGNING_KEY_DIR}/sysexts.key" \
  --certificate="${SYSEXT_SIGNING_KEY_DIR}/sysexts.crt" \
  --make-ddi=sysext \
  --copy-source="${BUILD_DIR}/${FLAGS_install_root_basename}" \
  "${BUILD_DIR}/${SYSEXTNAME}.raw"

rm -rf "${BUILD_DIR}"/{fs-root,"${FLAGS_install_root_basename}",workdir}

# Generate reports
mkdir "${BUILD_DIR}/img-rootfs"
systemd-dissect --read-only \
  --mount \
  --mkdir \
  --image-policy='root=encrypted+unprotected+absent:usr=encrypted+unprotected+absent' \
  "${BUILD_DIR}/${SYSEXTNAME}.raw" \
  "${BUILD_DIR}/img-rootfs"

write_contents "${BUILD_DIR}/img-rootfs" "${BUILD_DIR}/${SYSEXTNAME}_contents.txt"
write_contents_with_technical_details "${BUILD_DIR}/img-rootfs" "${BUILD_DIR}/${SYSEXTNAME}_contents_wtd.txt"
write_disk_space_usage_in_paths "${BUILD_DIR}/img-rootfs" "${BUILD_DIR}/${SYSEXTNAME}_disk_usage.txt"
systemd-dissect --umount --rmdir "${BUILD_DIR}/img-rootfs"