#!/bin/bash
set -euo pipefail

opts=$(getopt --name "$(basename "${0}")" --options 'hV:P:DFA' \
       --longoptions 'help,to-version:,to-payload:,force-dev-key,force-flatcar-key,disable-afterwards' -- "${@}")
eval set -- "${opts}"

USER_PAYLOAD=
PAYLOAD=
VERSION=
FORCE_DEV_KEY=
FORCE_FLATCAR_KEY=
DISABLE_AFTERWARDS=

while true; do
  case "$1" in
  -h|--help)
    echo "Usage: $(basename "${0}") --to-version VERSION [--to-payload FILENAME] [--force-dev-key|--force-flatcar-key|--disable-afterwards]"
    echo "  Updates Flatcar Container Linux through a temporary local update service on localhost."
    echo "  The update-engine service will be unmasked (to disable updates again use -A)."
    echo "  The reboot should be done after applying the update, either manually or through your reboot manager (check locksmithd/FLUO)."
    echo "  An error will be reported if a previously applied update wasn't booted into yet (you may discard it with 'update_engine_client -reset_status')."
    echo "  Warning: If you jump between channels, delete any GROUP configured in /etc/flatcar/update.conf for the new defaults to apply."
    echo "Options:"
    echo "  -V, --to-version <VERSION>	Updates to the version, by default using the matching release from update.release.flatcar-linux.net"
    echo "  -P, --to-payload <FILENAME>	Updates to the given update payload file instead of downloading it"
    echo "  -D, --force-dev-key		Bind-mounts the dev key over /usr/share/update_engine/update-payload-key.pub.pem"
    echo "  -F, --force-flatcar-key	Bind-mounts the Flatcar release key over /usr/share/update_engine/update-payload-key.pub.pem"
    echo "  -A, --disable-afterwards	Writes SERVER=disabled to /etc/flatcar/update.conf when done (this overwrites any custom SERVER)"
    echo
    echo "Example for updating to the latest Stable release and disabling automatic updates afterwards:"
    echo '  VER=$(curl -fsSL https://stable.release.flatcar-linux.net/amd64-usr/current/version.txt | grep FLATCAR_VERSION= | cut -d = -f 2)'
    echo "  $(basename "${0}") -V \$VER -A"
    exit 1
    ;;
  -V|--to-version)
    shift
    VERSION="$1"
    ;;
  -P|--to-payload)
    shift
    PAYLOAD="$1"
    USER_PAYLOAD=1
    if [ "$PAYLOAD" = "" ]; then
      echo "Error: --to-payload must not have an empty value" > /dev/stderr ; exit 1
    fi
    ;;
  -D|--force-dev-key)
    FORCE_DEV_KEY=1
    KEY="https://raw.githubusercontent.com/flatcar-linux/coreos-overlay/main/coreos-base/coreos-au-key/files/developer-v1.pub.pem"
    ;;
  -F|--force-flatcar-key)
    FORCE_FLATCAR_KEY=1
    KEY="https://raw.githubusercontent.com/flatcar-linux/coreos-overlay/flatcar-master/coreos-base/coreos-au-key/files/official-v2.pub.pem"
    ;;
  -A|--disable-afterwards)
    DISABLE_AFTERWARDS=1
    ;;
  --)
    shift
    break;;
  esac
  shift
done

if [ "${VERSION}" = "" ]; then
  echo "Error: must specify --to-version" > /dev/stderr ; exit 1
fi

if [ "${FORCE_DEV_KEY}" = "1" ] && [ "${FORCE_FLATCAR_KEY}" = "1" ]; then
  echo "Error: must only specify one of --force-dev-key or --force-flatcar-key" > /dev/stderr ; exit 1
fi

[ "$EUID" = "0" ] || { echo "Need to be root: sudo $0 $opts" > /dev/stderr ; exit 1 ; }

if mount | grep -q /usr/share/update_engine/update-payload-key.pub.pem; then
  echo "Warning: found a bind mount on /usr/share/update_engine/update-payload-key.pub.pem (will only unmount it if --force-dev-key|--force-flatcar-key is set)"
fi

if ss -tan | grep -q "LISTEN.*:9090"; then
  echo "Error: some process is using port 9090" > /dev/stderr ; exit 1
fi
if ss -tan | grep -q "LISTEN.*:9091"; then
  echo "Error: some process is using port 9091" > /dev/stderr ; exit 1
fi

# Migrate CoreOS machines to Flatcar
if [ -d "/etc/coreos" ]; then
  mv /etc/coreos /etc/flatcar
  ln -s flatcar /etc/coreos
fi

HARDCODED_GROUP=$(grep -m 1 -o '^GROUP=.*' /etc/flatcar/update.conf 2> /dev/null || true)
if [ "${HARDCODED_GROUP}" != "" ]; then
  echo "Warning: found hardcoded ${HARDCODED_GROUP} in /etc/flatcar/update.conf - make sure it fits the release channel you want to follow" > /dev/stderr
fi

systemctl unmask update-engine
systemctl start update-engine

STATUS=$(update_engine_client -status 2>/dev/null | { grep '^CURRENT_OP=UPDATE_STATUS_UPDATED_NEED_REBOOT$' || true ; })
if [ "$STATUS" != "" ]; then
  echo "Error: a previously downloaded update wasn't applied yet, you can discard it with 'update_engine_client -reset_status'" > /dev/stderr; exit 1
fi

touch /etc/flatcar/update.conf
PREV_SERVER=$(grep '^SERVER=' /etc/flatcar/update.conf || true)
sed -i "/SERVER=.*/d" /etc/flatcar/update.conf

echo "SERVER=http://localhost:9090/update" >> /etc/flatcar/update.conf
BOARD=$({ grep -m 1 BOARD= /usr/share/coreos/release || true ; } | cut -d = -f 2-)
if [ "$BOARD" = "" ]; then
  echo "Error: could not find board from /usr/share/coreos/release" > /dev/stderr ; exit 1
fi

SHA256_TO_CHECK=
if [ "$PAYLOAD" = "" ]; then
  PAYLOAD="/var/tmp/update_payload"
  rm -f "$PAYLOAD"
  echo "Downloading update payload..."
  curl -fsSL -o "$PAYLOAD" --retry-delay 1 --retry 60 --retry-connrefused --retry-max-time 60 --connect-timeout 20 "https://update.release.flatcar-linux.net/${BOARD}/${VERSION}/flatcar_production_update.gz"
  SHA256_TO_CHECK=$(curl -fsSL --retry-delay 1 --retry 60 --retry-connrefused --retry-max-time 60 --connect-timeout 20 "https://update.release.flatcar-linux.net/${BOARD}/${VERSION}/flatcar_production_update.gz.sha256" | cut -d " " -f 1)
  if [ "${SHA256_TO_CHECK}" = "" ]; then
    echo "Error: could not download sha256 checksum file" > /dev/stderr ; exit 1
  fi
  SHA256_HEX=$(sha256sum -b "$PAYLOAD" | cut -d " " -f 1)
  if [ "${SHA256_TO_CHECK}" != "${SHA256_HEX}" ]; then
    echo "Error: mismatch with downloaded SHA256 checksum (${SHA256_TO_CHECK})" > /dev/stderr ; exit 1
  fi
  echo "When restarting after an error you may reuse it with '--to-payload $PAYLOAD'"
fi

BASE="http://localhost:9091/"
HASH=$(openssl dgst -binary -sha1 < "$PAYLOAD" | base64)
SHA256=$(openssl dgst -binary -sha256 < "$PAYLOAD" | base64)
SIZE=$(stat --printf='%s\n' "$PAYLOAD")

rm -f /tmp/response
tee /tmp/response > /dev/null <<-EOF
	<response protocol="3.0" server="flatcar-update"><daystart elapsed_seconds="0"></daystart>
	<app appid="{e96281a6-d1af-4bde-9a0a-97b76e56dc57}" status="ok"><ping status="ok"></ping>
	<updatecheck status="ok"><urls><url codebase="${BASE}"></url></urls>
	<manifest version="${VERSION}"><packages><package name="flatcar_production_update.gz" hash="${HASH}" size="${SIZE}" required="true"></package></packages>
	<actions><action event="postinstall" sha256="${SHA256}" DisablePayloadBackoff="true"></action></actions></manifest>
	</updatecheck><event status="ok"></event></app></response>
EOF

ncat --keep-open -c "echo -en 'HTTP/1.1 200 OK\ncontent-type: application/gzip\ncontent-length: $SIZE\n\n'; cat \"$PAYLOAD\"" -l 9091 &
ncat --keep-open -c "echo -en 'HTTP/1.1 200 OK\ncontent-type: text/xml\ncontent-length: $(stat --printf='%s\n' /tmp/response)\n\n'; cat /tmp/response" -l 9090 &
trap "umount /usr/share/update_engine/update-payload-key.pub.pem 2> /dev/null || true; rm -f /tmp/response ; kill 0" EXIT INT

if [ "${FORCE_DEV_KEY}" = "1" ] || [ "${FORCE_FLATCAR_KEY}" = "1" ]; then
  rm -f /tmp/key
  curl -fsSL -o /tmp/key --retry-delay 1 --retry 60 --retry-connrefused --retry-max-time 60 --connect-timeout 20 "$KEY"
  umount /usr/share/update_engine/update-payload-key.pub.pem 2> /dev/null || true
  echo "Bind-mounting /usr/share/update_engine/update-payload-key.pub.pem"
  mount --bind /tmp/key /usr/share/update_engine/update-payload-key.pub.pem
fi

echo "Forcing update..."

# Force an update
if update_engine_client -update 2> /dev/null > /dev/null; then
  STATUS=$(update_engine_client -status 2>/dev/null | { grep '^CURRENT_OP=UPDATE_STATUS_UPDATED_NEED_REBOOT$' || true ; })
else
  STATUS=
fi

# Set previous or wanted SERVER setting
sed -i "/SERVER=.*/d" /etc/flatcar/update.conf
if [ "${DISABLE_AFTERWARDS}" = "1" ]; then
  echo "Setting SERVER=disabled in /etc/flatcar/update.conf"
  echo "SERVER=disabled" >> /etc/flatcar/update.conf
elif [ "${PREV_SERVER}" != "" ]; then
  echo "${PREV_SERVER}" >> /etc/flatcar/update.conf
fi

if [ "$STATUS" = "" ]; then
  echo "Error: update failed" > /dev/stderr; exit 1
fi

if [ "${USER_PAYLOAD}" = "" ]; then
  echo "Removing payload $PAYLOAD"
  rm -f "$PAYLOAD"
fi

echo "Done, please make sure to reboot either manually or through your reboot manager (check locksmithd/FLUO)"
