#!/usr/bin/env bash
set -euo pipefail

[[ $# -eq 4 ]] || {
  echo "usage: prepare-jd-native-probe-once.sh <world-root> <probe-image> <recovery-root> <expected-image-sha256>" >&2
  exit 64
}
[[ ${EUID} -eq 0 ]] || {
  echo "must run as root" >&2
  exit 77
}

world_root=$(readlink -f "$1")
candidate=$(readlink -f "$2")
recovery_root=$(readlink -m "$3")
expected_candidate_sha=$4
disk=/dev/vda
expected_bootloader_sha=32291877e37647ee5b778f968676cc69b59fb5b07604366a21cbfd8e4a4b87a8
expected_empty_38_sha=c597b1b463b7cec49548fac00d7588c9652bb67cc2b1b0c88676bf8f1558571e

[[ ${expected_candidate_sha} =~ ^[0-9a-f]{64}$ ]]
[[ ! -e ${recovery_root} ]]
grep -q '^node_id: JD-FD-PRIMARY$' "${world_root}/CURRENT.hldp"
/guanghu/bin/ghctl authorize "${world_root}" write_bootloader_and_system_partitions >/dev/null
[[ $(stat -c %s "${candidate}") -eq 14848 ]]
candidate_sha=$(sha256sum "${candidate}" | awk '{print $1}')
[[ ${candidate_sha} == "${expected_candidate_sha}" ]]

partition_dump=$(sfdisk -d "${disk}")
grep -q '^label: dos$' <<<"${partition_dump}"
grep -q '^label-id: 0xf4bb72a5$' <<<"${partition_dump}"
grep -Eq '^/dev/vda1 : start= *2048, size= *209713119, type=83$' <<<"${partition_dump}"

bootloader_sha=$(dd if="${disk}" bs=512 count=105 status=none | sha256sum | awk '{print $1}')
[[ ${bootloader_sha} == "${expected_bootloader_sha}" ]]
empty_sha=$(dd if="${disk}" bs=512 skip=105 count=38 status=none | sha256sum | awk '{print $1}')
[[ ${empty_sha} == "${expected_empty_38_sha}" ]]

mkdir -m 0700 "${recovery_root}"
printf '%s\n' "${partition_dump}" >"${recovery_root}/sfdisk-before.txt"
dd if="${disk}" of="${recovery_root}/first-2MiB.before.bin" bs=1M count=2 status=none
disk_sectors=$(blockdev --getsz "${disk}")
dd if="${disk}" of="${recovery_root}/last-2MiB.before.bin" \
  bs=512 skip=$((disk_sectors - 4096)) count=4096 status=none
cp /boot/grub/grub.cfg "${recovery_root}/grub.cfg.before"
cp /boot/grub/grubenv "${recovery_root}/grubenv.before"
cp /etc/default/grub "${recovery_root}/default-grub.before"
cp /etc/grub.d/40_custom "${recovery_root}/40_custom.before"
cp "${candidate}" "${recovery_root}/jd-native-probe-stage9.img"
cat /proc/sys/kernel/random/boot_id >"${recovery_root}/linux-boot-id.before"

install -m 0755 /dev/stdin /etc/grub.d/41_guanghu_jd_native_once <<'EOF'
#!/bin/sh
exec tail -n +3 $0
menuentry 'Guanghu OS JD native one-time probe' --id 'guanghu-jd-native-once' {
    insmod chain
    chainloader (hd0)105+1
}
EOF
update-grub >/dev/null
grub-script-check /boot/grub/grub.cfg
grep -q "menuentry 'Guanghu OS JD native one-time probe'.*'guanghu-jd-native-once'" \
  /boot/grub/grub.cfg

bootloader_sha_after_grub=$(dd if="${disk}" bs=512 count=105 status=none |
  sha256sum | awk '{print $1}')
[[ ${bootloader_sha_after_grub} == "${expected_bootloader_sha}" ]]
dd if="${candidate}" of="${disk}" bs=512 seek=105 count=29 \
  conv=notrunc,fsync status=none
readback_sha=$(dd if="${disk}" bs=512 skip=105 count=29 status=none |
  sha256sum | awk '{print $1}')
[[ ${readback_sha} == "${candidate_sha}" ]]
tail_sha=$(dd if="${disk}" bs=512 skip=134 count=9 status=none |
  sha256sum | awk '{print $1}')
expected_tail_sha=$(head -c $((9 * 512)) /dev/zero | sha256sum | awk '{print $1}')
[[ ${tail_sha} == "${expected_tail_sha}" ]]

observed_at=$(date --iso-8601=seconds)
cat >"${recovery_root}/PREPARE-RECEIPT.hldp" <<EOF
schema: guanghu.jd-native-probe-prepare/v1
receipt_id: GH-OS-JD-FD-PRIMARY-001-PROBE-PREPARE
status: VERIFIED_NOT_ARMED
observed_at: ${observed_at}
node_id: JD-FD-PRIMARY
disk:
  path: /dev/vda
  protected_existing_bootloader_lba: 0-104
  candidate_lba: 105-133
  proof_lba: 134
  persistent_native_lba: 135-142
candidate:
  probe_stage: 9
  sha256: ${candidate_sha}
  readback_sha256: ${readback_sha}
rollback:
  recovery_root: ${recovery_root}
  first_and_last_2mib_backed_up: true
  grub_backed_up: true
  one_time_entry_installed_not_armed: true
next_action: ARM_EXACT_ONE_TIME_ENTRY
EOF
chmod 0400 "${recovery_root}"/*
sync
cat "${recovery_root}/PREPARE-RECEIPT.hldp"