#! /bin/ksh

# Copyright (c) 2021 Stefan Kreutz <mail@skreutz.com>
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

# shellcheck disable=SC2001

set -o errexit
set -o nounset
set -o pipefail

verbosity=0
tmp=
vnode=

function main {
  trap 'cleanup' EXIT ERR INT

  while getopts :v option; do
    case "$option" in
      v) verbosity=$((verbosity+1)) ;;
      :) err_exit "missing argument for option -$OPTARG" ;;
      ?) err_exit "illegal option: -$OPTARG" ;;
    esac
  done
  shift $((OPTIND-1))

  if [ $# -eq 0 ]; then
    err_exit "missing command"
  fi
  local _cmd="$1"
  shift 1

  tmp="$( readlink -f "$( mktemp -d -t )" )"
  debug "Created temporary directory $tmp"

  case "$_cmd" in
    fetch|patch) "$_cmd" "$@" ;;
    *) err_exit "undefined command: $_cmd" ;;
  esac
}

function fetch {
  local _force=0
  local _mirror _arch _release _pubkey _output _url _ftp_options _rel _iso

  while getopts :fm:r:p:o: option; do
    case "$option" in
      f) _force=1 ;;
      m) _mirror="$OPTARG" ;;
      r) _release="$OPTARG" ;;
      p) _pubkey="$OPTARG" ;;
      o) _output="$OPTARG" ;;
      :) err_exit "missing argument for option -$OPTARG" ;;
      ?) err_exit "illegal option: -$OPTARG" ;;
    esac
  done
  shift $((OPTIND-1))

  if [ $# -ne 0 ]; then
    err_exit "too many arguments: $*"
  fi

  _arch="$( uname -p )"
  if [ "$_arch" != amd64 ]; then
    err_exit "unsupported architecture: $_arch"
  fi

  _mirror="${_mirror:-"$( get_mirror )"}"
  case "$_mirror" in
    http://*|https://*) ;;
    *) err_exit "unsupported mirror: $_mirror" ;;
  esac

  _url="${_mirror%/}/${_release:-snapshots}/$_arch"
  info "Fetching from $_url/"

  case $verbosity in
    0) _ftp_options="-VM" ;;
    1) _ftp_options="-V" ;;
    *) _ftp_options="-v" ;;
  esac

  debug "Downloading SHA256.sig ..."
  ftp "$_ftp_options" -o "$tmp/SHA256.sig" "$_url/SHA256.sig"
  debug "Downloaded SHA256.sig"

  _rel="$( sed -n 's/^SHA256 (install\([0-9][0-9]\).iso) = .*/\1/p' "$tmp/SHA256.sig" | head -n 1 )"
  _iso="install${_rel}.iso"
  _output="${_output:-"$_iso"}"

  if [ -e "$_output" ] && [ $_force -eq 0 ]; then
    err_exit "file already exists: $_output"
  fi

  debug "Downloading $_iso ..."
  ftp "$_ftp_options" -o "$tmp/$_iso" "$_url/$_iso"
  debug "Downloaded $_iso"

  _pubkey="${_pubkey:-"/etc/signify/openbsd-${_rel}-base.pub"}"
  debug "Verifying $_iso using $_pubkey ..."
  ( cd "$tmp" && signify -C -p "$_pubkey" -x SHA256.sig -q -- "$_iso" )
  debug "Verified $_iso"

  mv "$tmp/$_iso" "$_output"
  if [ "$_iso" = "$_output" ]; then
    info "Fetched $_iso"
  else
    info "Fetched $_iso as $_output"
  fi
}

function patch {
  local _force=0
  local _compressed=0
  local _install_conf _upgrade_conf _site_dir _input_iso _output_iso
  local _release _rel _arch _mime _rd_path

  while getopts :fi:u:s: option; do
    case "$option" in
      f) _force=1 ;;
      i) _install_conf="$OPTARG" ;;
      u) _upgrade_conf="$OPTARG" ;;
      s) _site_dir="$OPTARG" ;;
      :) err_exit "missing argument for option -$OPTARG" ;;
      ?) err_exit "illegal option: -$OPTARG" ;;
    esac
  done
  shift $((OPTIND-1))

  if [ $# -lt 2 ]; then
    err_exit "missing arguments"
  fi
  _input_iso="$1"
  _output_iso="$2"
  shift 2

  if [ $# -ne 0 ]; then
    err_exit "too many arguments: $*"
  fi

  info "Unpacking ISO 9660 image ..."
  vnode="$( vnconfig "$_input_iso" )"
  debug "Configured $_input_iso as vnode disk $vnode"
  mkdir "$tmp/mnt"
  mount -t cd9660 "/dev/${vnode}c" "$tmp/mnt"
  debug "Mounted vnode disk $vnode"
  mkdir "$tmp/cd"
  tar -C "$tmp/mnt" -c -f - . | tar -C "$tmp/cd" -x -p -f -
  debug "Copied contents of $_input_iso to disk"
  umount "$tmp/mnt"
  debug "Unmounted vnode disk $vnode"
  vnconfig -u "$vnode"
  vnode=
  debug "Unconfigured vnode disk $vnode"
  info "Unpacked ISO 9660 image"

  debug "Inspecting directory layout ..."
  _rd_path="$( cd "$tmp/cd" && echo */*/bsd.rd )"
  _release="$( echo "$_rd_path" | sed 's,^\([0-9]\)\.\([0-9]\)/\([^/]*\)/bsd.rd,\1.\2,' )"
  _rel="$( echo "$_rd_path" | sed 's,^\([0-9]\)\.\([0-9]\)/\([^/]*\)/bsd.rd,\1\2,' )"
  debug "Inferred release $_release resp. $_rel"
  _arch="$( echo "$_rd_path" | sed 's,^\([0-9]\)\.\([0-9]\)/\([^/]*\)/bsd.rd,\3,' )"
  debug "Inferred arch $_arch"

  if [ "$_arch" != amd64 ]; then
    err_exit "unsupported architecture: $_arch"
  fi

  if [ -e "$_output_iso" ] && [ $_force -eq 0 ]; then
    err_exit "file already exists: $_output_iso"
  fi

  info "Patching RAMDISK kernel bsd.rd ..."
  _mime="$( file --brief --mime "$tmp/cd/$_release/$_arch/bsd.rd" )"
  if [ "$_mime" = "application/x-gzip" ]; then
    _compressed=1
    gzip -d -o "$tmp/bsd.rd" "$tmp/cd/$_release/$_arch/bsd.rd"
    debug "Uncompressed RAMDISK kernel bsd.rd"
  else
    cp "$tmp/cd/$_release/$_arch/bsd.rd" "$tmp/bsd.rd"
  fi
  rdsetroot -x "$tmp/bsd.rd" "$tmp/disk.fs"
  debug "Extracted disk image from RAMDISK kernel bsd.rd"
  vnode="$( vnconfig "$tmp/disk.fs" )"
  debug "Configured RAMDISK kernel's disk image as vnode disk $vnode"
  mount "/dev/${vnode}a" "$tmp/mnt"
  debug "Mounted vnode disk $vnode"
  if [ -n "${_install_conf:-}" ]; then
    install -o root -g wheel -m 0644 -C "$_install_conf" "$tmp/mnt/auto_install.conf"
    debug "Installed autoinstall(8) install response file $_install_conf"
  fi
  if [ -n "${_upgrade_conf:-}" ]; then
    install -o root -g wheel -m 0644 -C "$_upgrade_conf" "$tmp/mnt/auto_upgrade.conf"
    debug "Installed autoinstall(8) upgrade response file $_upgrade_conf"
  fi
  umount "$tmp/mnt"
  debug "Unmounted vnode disk $vnode"
  vnconfig -u "$vnode"
  vnode=
  debug "Unconfigured vnode disk $vnode"
  rdsetroot "$tmp/bsd.rd" "$tmp/disk.fs"
  debug "Inserted modified disk image into RAMDISK kernel bsd.rd"
  debug "Patched RAMDISK kernel"
  if [ $_compressed -eq 1 ]; then
    gzip -9fnq "$tmp/bsd.rd"
    mv "$tmp/bsd.rd.gz" "$tmp/bsd.rd"
    debug "Compressed patched RAMDISK kernel bsd.rd"
  fi

  info "Patching ISO 9660 image contents ..."
  install -o root -g 2000 -m 0755 -C "$tmp/bsd.rd" "$tmp/cd/$_release/$_arch/bsd.rd"
  debug "Installed modified RAMDISK kernel bsd.rd"
  if [ -n "${_site_dir:-}" ]; then
    if [ -e "$_site_dir/install.site" ] && [ ! -x "$_site_dir/install.site" ]; then
      warn "$_site_dir/install.site is not executable"
    fi
    ( cd "$_site_dir" && tar -c -z -f "$tmp/cd/$_release/$_arch/site${_rel}.tgz" . )
    ( cd "$tmp/cd/$_release/$_arch" && ls -l > index.txt )
    debug "Installed site-specific file set $_site_dir"
  fi
  debug "Patched ISO 9660 image contents"

  info "Creating bootable ISO 9660 image ..."
  # Source: https://cvsweb.openbsd.org/src/distrib/amd64/iso/Makefile?rev=1.47&content-type=text/x-cvsweb-markup
  if [ -e "${tmp}/cd/${_release}/${_arch}/eficdboot" ] ; then
    _efi=1
  fi
  if ! mkhybrid -a -R -T -L -l -d -D -N -o "${_output_iso}" \
    -A "Custom OpenBSD ${_release} ${_arch} Install CD" \
    -P "Copyright (c) $(date +%Y) Theo de Raadt, The OpenBSD project" \
    -p "Generated using installiso(8)" \
    -b "${_release}/${_arch}/cdbr" -c "${_release}/${_arch}/boot.catalog" \
    ${_efi:+-e "${_release}/${_arch}/eficdboot"} \
    "${tmp}/cd"
  then
    rm -f "$_output_iso"
    err_exit "failed to create bootable ISO 9660 image"
  fi
  info "Created bootable ISO 9660 image $_output_iso"
}

function cleanup {
  trap - EXIT ERR INT
  set +o errexit
  debug "Cleaning up ..."
  if [ -n "${vnode:-}" ]; then
    if mount | grep -qe "^/dev/$vnode" ; then
      umount "/dev/$vnode" && debug "Unmounted vnode disk $vnode"
    fi
    vnconfig -u "$vnode" && debug "Unconfigured vnode disk $vnode"
  fi
  if [ -n "${tmp:-}" ]; then
    rm -rf "$tmp" && debug "Removed temporary directory $tmp"
  fi
}

function get_mirror {
  grep -v -e '^[:space:]*$' -e '^#' /etc/installurl | head -n 1 ||
    echo "https://cdn.openbsd.org/pub/OpenBSD/"
}

function err_exit {
  print -u2 -- "$*"
  exit 1
}

function warn {
  print -u2 -- "$*"
}

function info {
  if [ $verbosity -lt 1 ]; then
    return
  fi
  print -- "$*"
}

function debug {
  if [ $verbosity -lt 2 ]; then
    return
  fi
  print -- "$*"
}

main "$@"