#!/bin/bash
# PXE Clonezilla: back up or restore a PC by its machine number.
#
# The iPXE menu (GetPxeScript.aspx, "Clonezilla: back up / restore this
# PC") boots Clonezilla live with:
#   ocs_prerun="wget ... -O /tmp/pxe-cz.sh http://<server>/clonezilla-scripts/pxe-cz.sh"
#   ocs_repository="smb3://clonezilla:<password>@<server>/clonezilla/"
#   ocs_live_run="bash /tmp/pxe-cz.sh"
#   pxe_cz_mode=save|restore pxe_machine=<typed text> pxe_end=1
# Clonezilla runs the prerun (this download), mounts the repository on
# /home/partimag, then runs this script as root on the console.
#
# Images are named <machine>-<YYYY-MM-DD-HHMM> (server time, site zone),
# with pxe-meta.txt (machine, serial, model, disk, date, Clonezilla
# version, 32/64-bit, check result) next to Clonezilla's own files. The
# PXE webapp (Tools > Clonezilla backups) reads both.
#
# Runs under Clonezilla live's bash (3.3.3 amd64 and 3.2.0-5 i686, both
# Debian with coreutils, util-linux, dmidecode, wget). Plain ASCII only.

set -u

# PXE_CZ_* overrides exist only for testing this script off Clonezilla.
OCSROOT="${PXE_CZ_OCSROOT:-/home/partimag}"
CMDLINE_FILE="${PXE_CZ_CMDLINE:-/proc/cmdline}"
SYS_BLOCK="${PXE_CZ_SYS_BLOCK:-/sys/block}"
SITE_TZ="America/New_York"
SERVER="$(sed -n 's/.*ocs_prerun="[^"]*http:\/\/\([^/":]*\).*/\1/p' "$CMDLINE_FILE" 2>/dev/null)"
[ -n "$SERVER" ] || SERVER=172.16.9.1
CMDLINE=" $(cat "$CMDLINE_FILE" 2>/dev/null) "

# --- output helpers -------------------------------------------------------

line() { echo "======================================================================"; }

banner() {
  echo
  line
  for l in "$@"; do echo "  $l"; done
  line
  echo
}

# Ask what to do at the end. Success: power off after 60 s unless a key is
# pressed. Failure: wait, the operator has to read the message.
finish() {
  local ok="$1" ans
  sync
  while true; do
    echo
    echo "  [Enter] power off     r = reboot     s = shell (type exit to come back)"
    if [ "$ok" = "yes" ]; then
      echo -n "  Powering off in 60 seconds... "
      if ! read -r -t 60 ans; then ans=""; echo; fi
    else
      echo -n "  Choice: "
      read -r ans || ans=""
    fi
    case "$ans" in
      r|R) umount "$OCSROOT" 2>/dev/null; sync; reboot; sleep 30 ;;
      s|S) bash -i ;;
      *)   umount "$OCSROOT" 2>/dev/null; sync; poweroff; sleep 30 ;;
    esac
  done
}

fail() {
  banner "FAILED: $1" "${@:2}"
  finish no
  exit 1
}

# --- machine number -------------------------------------------------------

# Letters, digits and dash, starting with a letter or digit, up to 32.
valid_machine() {
  [[ "$1" =~ ^[A-Za-z0-9][A-Za-z0-9-]{0,31}$ ]]
}

# The text typed at the iPXE prompt sits between pxe_machine= and the
# pxe_end=1 sentinel, so a typed space shows up here and is rejected
# instead of being split into another kernel parameter.
machine_from_cmdline() {
  local rest
  case "$CMDLINE" in
    *" pxe_machine="*" pxe_end=1 "*) ;;
    *) echo ""; return ;;
  esac
  rest="${CMDLINE#* pxe_machine=}"
  echo "${rest%% pxe_end=1 *}"
}

# Sets MACHINE (upper case). Not run in $(...): finish() reads the
# keyboard and must write to the console.
get_machine() {
  local m
  m="$(machine_from_cmdline)"
  if [ -z "$m" ] || ! valid_machine "$m"; then
    if [ -n "$m" ]; then
      banner "The machine number typed at the boot menu is not valid:" \
             "  [$m]" "Use letters, digits and dash only (for example 6501 or WJ-6501)."
    else
      banner "No machine number was given at the boot menu."
    fi
    while true; do
      echo -n "  Machine number (letters, digits, dash; empty = stop): "
      read -r m || m=""
      [ -n "$m" ] || finish no
      valid_machine "$m" && break
      echo "  [$m] is not a valid machine number: use letters, digits and dash only."
    done
  fi
  MACHINE="$(echo "$m" | tr '[:lower:]' '[:upper:]')"
}

# --- time -----------------------------------------------------------------

# The PC's clock is often wrong (dead CMOS battery on old PCs, local time
# on Windows PCs), so take the time from the PXE server's HTTP Date header
# and set the live system clock from it. Falls back to the PC clock.
TIME_SOURCE="pc-clock"
sync_clock() {
  local hdr epoch
  hdr="$(wget -q -S --spider -T 10 -t 2 "http://$SERVER/clonezilla-scripts/pxe-cz.sh" 2>&1 \
        | sed -n 's/^ *Date: *//p' | head -1)"
  [ -n "$hdr" ] || return 0
  epoch="$(date -u -d "$hdr" +%s 2>/dev/null)" || return 0
  [ -n "$epoch" ] || return 0
  date -u -s "@$epoch" >/dev/null 2>&1 && TIME_SOURCE="pxe-server"
}

site_date() {
  if [ -e "/usr/share/zoneinfo/$SITE_TZ" ]; then
    TZ="$SITE_TZ" date "$@"
  else
    date "$@"
  fi
}

# --- hardware -------------------------------------------------------------

dmi() {
  local v
  v="$(dmidecode -s "$1" 2>/dev/null | grep -v '^#' | head -1 | tr -d '\r' | sed 's/^ *//; s/ *$//')"
  echo "$v"
}

# Whole disks a backup can come from or a restore can go to: no CD, loop,
# RAM, zram, device-mapper, md, eMMC boot areas, or empty card readers.
list_disks() {
  local d n size
  for d in "$SYS_BLOCK"/*; do
    n="${d##*/}"
    case "$n" in
      loop*|sr*|ram*|zram*|dm-*|md*|fd*|nbd*|*boot[0-9]|*rpmb) continue ;;
    esac
    size="$(cat "$d/size" 2>/dev/null)"
    [ -n "$size" ] && [ "$size" -gt 0 ] 2>/dev/null || continue
    echo "$n"
  done
}

disk_bytes() { echo $(( $(cat "$SYS_BLOCK/$1/size" 2>/dev/null || echo 0) * 512 )); }

disk_model() {
  local m
  m="$(cat "$SYS_BLOCK/$1/device/model" 2>/dev/null | sed 's/^ *//; s/ *$//')"
  [ -n "$m" ] || m="$(lsblk -dn -o MODEL "/dev/$1" 2>/dev/null | sed 's/^ *//; s/ *$//')"
  [ -n "$m" ] || m="unknown model"
  echo "$m"
}

disk_tran() { lsblk -dn -o TRAN "/dev/$1" 2>/dev/null | sed 's/ //g'; }

human() {
  local b="$1"
  awk -v b="$b" 'BEGIN { if (b >= 1e12) printf "%.1f TB", b/1e12; else printf "%.0f GB", b/1e9 }'
}

describe_disk() {
  local n="$1" t
  t="$(disk_tran "$n")"
  echo "/dev/$n  $(human "$(disk_bytes "$n")")  $(disk_model "$n")${t:+  ($t)}"
}

# pick_disk "prompt" -> disk name on stdout. One disk: that one.
pick_disk() {
  local prompt="$1" disks n i ans
  mapfile -t disks < <(list_disks)
  if [ "${#disks[@]}" -eq 0 ]; then
    echo ""
    return
  fi
  if [ "${#disks[@]}" -eq 1 ]; then
    echo "${disks[0]}"
    return
  fi
  while true; do
    {
      echo
      echo "  $prompt"
      i=1
      for n in "${disks[@]}"; do
        echo "    $i) $(describe_disk "$n")"
        i=$((i + 1))
      done
      echo -n "  Number (1-${#disks[@]}): "
    } >&2
    read -r ans || ans=""
    if [[ "$ans" =~ ^[0-9]+$ ]] && [ "$ans" -ge 1 ] && [ "$ans" -le "${#disks[@]}" ]; then
      echo "${disks[$((ans - 1))]}"
      return
    fi
    echo "  Type a number from the list." >&2
  done
}

live_version() {
  local v=""
  if [ -e /etc/ocs/ocs-live.conf ]; then
    # shellcheck disable=SC1091
    v="$(. /etc/ocs/ocs-live.conf >/dev/null 2>&1; echo "${drbl_ocs_live_ver:-}")"
  fi
  [ -n "$v" ] || v="clonezilla $(dpkg-query -W -f='${Version}' clonezilla 2>/dev/null)"
  echo "$v"
}

bits() {
  case "$(uname -m)" in
    x86_64|amd64) echo 64 ;;
    *) echo 32 ;;
  esac
}

# set_meta FILE KEY VALUE: replace or add one key=value line.
set_meta() {
  local f="$1" k="$2" v="$3" tmp
  v="$(echo "$v" | tr -d '\r\n')"
  tmp="$(mktemp /tmp/pxe-meta.XXXXXX)"
  [ -e "$f" ] && grep -v "^$k=" "$f" > "$tmp"
  echo "$k=$v" >> "$tmp"
  cat "$tmp" > "$f"
  rm -f "$tmp"
}

get_meta() {
  sed -n "s/^$2=//p" "$1" 2>/dev/null | head -1
}

repo_ready() {
  mountpoint -q "$OCSROOT" 2>/dev/null || grep -qs " $OCSROOT cifs " /proc/mounts
}

# --- backup ---------------------------------------------------------------

do_save() {
  local m="$1" disk stamp img n meta rc chk
  disk="$(pick_disk "This PC has more than one disk. Which one do you want to back up?")"
  [ -n "$disk" ] || fail "No disk found to back up." "Check the disk is connected and shows in the BIOS."
  stamp="$(site_date +%Y-%m-%d-%H%M)"
  img="$m-$stamp"
  n=2
  while [ -e "$OCSROOT/$img" ]; do
    img="$m-$stamp-$n"
    n=$((n + 1))
  done

  banner "BACKUP of machine $m" \
         "Disk:   $(describe_disk "$disk")" \
         "Image:  $img" \
         "To:     \\\\$SERVER\\clonezilla\\$img"

  ocs-sr -q2 -batch -j2 -z9p -i 0 -sfsck -scs -senc -p command savedisk "$img" "$disk"
  rc=$?

  meta="$OCSROOT/$img/pxe-meta.txt"
  if [ ! -d "$OCSROOT/$img" ]; then
    fail "Clonezilla did not create the image $img (exit code $rc)." \
         "See the messages above. Nothing was saved."
  fi
  : > "$meta"
  set_meta "$meta" machine "$m"
  set_meta "$meta" date "$(site_date '+%Y-%m-%d %H:%M %Z')"
  set_meta "$meta" date_utc "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  set_meta "$meta" time_source "$TIME_SOURCE"
  set_meta "$meta" serial "$(dmi system-serial-number)"
  set_meta "$meta" manufacturer "$(dmi system-manufacturer)"
  set_meta "$meta" model "$(dmi system-product-name)"
  set_meta "$meta" disk "$disk"
  set_meta "$meta" disk_model "$(disk_model "$disk")"
  set_meta "$meta" disk_size_bytes "$(disk_bytes "$disk")"
  set_meta "$meta" clonezilla "$(live_version)"
  set_meta "$meta" bits "$(bits)"
  if [ -d /sys/firmware/efi ]; then set_meta "$meta" firmware uefi; else set_meta "$meta" firmware bios; fi
  if [ "$rc" -ne 0 ]; then
    set_meta "$meta" save failed
    set_meta "$meta" check not-run
    fail "Clonezilla could not save machine $m (exit code $rc)." \
         "The partial image $img is marked as failed on the share;" \
         "delete it in the webapp (Tools > Clonezilla backups) and try again."
  fi
  set_meta "$meta" save ok

  banner "Checking the saved image $img ..."
  ocs-chkimg -b "$img"
  chk=$?
  set_meta "$meta" check_date "$(site_date '+%Y-%m-%d %H:%M %Z')"
  if [ "$chk" -ne 0 ]; then
    set_meta "$meta" check failed
    fail "The image $img was saved but its CHECK FAILED (exit code $chk)." \
         "Do not rely on it. Back this PC up again."
  fi
  set_meta "$meta" check ok
  sync
  banner "SUCCESS: machine $m backed up and checked." \
         "Image:  $img" \
         "Serial: $(get_meta "$meta" serial)   Model: $(get_meta "$meta" model)"
  finish yes
}

# --- restore --------------------------------------------------------------

is_cz_image() {
  local d="$1"
  [ -d "$d" ] || return 1
  [ -f "$d/disk" ] || [ -f "$d/parts" ] || return 1
  local f
  for f in "$d"/*; do
    case "${f##*/}" in
      *-ptcl-img*|*.dd-img*|*-img.*|clonezilla-img) return 0 ;;
    esac
  done
  return 1
}

# Images of one machine, newest first: <machine>-YYYY-MM-DD-HHMM[-n] by
# name, plus any image whose pxe-meta.txt says machine=<machine>.
machine_images() {
  local m="$1" d n mm
  for d in "$OCSROOT"/*/; do
    d="${d%/}"
    n="${d##*/}"
    is_cz_image "$d" || continue
    mm="$(get_meta "$d/pxe-meta.txt" machine | tr '[:lower:]' '[:upper:]')"
    if [ "$mm" = "$m" ]; then
      echo "$n"
    elif [ -z "$mm" ] && echo "$n" | tr '[:lower:]' '[:upper:]' | grep -qE "^$m-[0-9]{4}-[0-9]{2}-[0-9]{2}-[0-9]{4}(-[0-9]+)?$"; then
      echo "$n"
    fi
  done | sort_newest
}

# Sort names by their date stamp (the last YYYY-MM-DD-HHMM[-n]), newest first.
sort_newest() {
  local n key
  while read -r n; do
    key="$(echo "$n" | grep -oE '[0-9]{4}-[0-9]{2}-[0-9]{2}-[0-9]{4}(-[0-9]+)?$')"
    echo "${key:-0000}|$n"
  done | sort -r | cut -d'|' -f2-
}

image_disks() { tr -s ' \t' '\n' < "$OCSROOT/$1/disk" 2>/dev/null | grep -v '^$'; }

image_disk_bytes() {
  local img="$1" src="$2" b sec
  b="$(get_meta "$OCSROOT/$img/pxe-meta.txt" disk_size_bytes)"
  if [[ "$b" =~ ^[0-9]+$ ]]; then echo "$b"; return; fi
  # Clonezilla's sfdisk dump: "last-lba: N" (GPT) or the partitions' end.
  sec="$(sed -n 's/^last-lba: *//p' "$OCSROOT/$img/$src-pt.sf" 2>/dev/null | head -1)"
  if [[ "$sec" =~ ^[0-9]+$ ]]; then echo $(( (sec + 34) * 512 )); return; fi
  echo 0
}

do_restore() {
  local m="$1" imgs i n ans img src srcb tgt tgtb meta chkopt rc mo se ck
  mapfile -t imgs < <(machine_images "$m")
  if [ "${#imgs[@]}" -eq 0 ]; then
    fail "No images for machine $m on \\\\$SERVER\\clonezilla." \
         "Check the machine number, or look in the webapp (Tools > Clonezilla backups)."
  fi
  while true; do
    echo
    echo "  Images of machine $m (newest first):"
    i=1
    for n in "${imgs[@]}"; do
      meta="$OCSROOT/$n/pxe-meta.txt"
      mo="$(get_meta "$meta" model)"; se="$(get_meta "$meta" serial)"; ck="$(get_meta "$meta" check)"
      echo "    $i) $n   ${mo:-model ?}  serial ${se:-?}  check: ${ck:-?}"
      i=$((i + 1))
    done
    echo -n "  Image to restore (1-${#imgs[@]}, q = quit): "
    read -r ans || ans=q
    case "$ans" in q|Q) finish no ;; esac
    if [[ "$ans" =~ ^[0-9]+$ ]] && [ "$ans" -ge 1 ] && [ "$ans" -le "${#imgs[@]}" ]; then
      img="${imgs[$((ans - 1))]}"
      break
    fi
    echo "  Type a number from the list."
  done
  meta="$OCSROOT/$img/pxe-meta.txt"
  if [ "$(get_meta "$meta" save)" = "failed" ]; then
    fail "The image $img is marked as a FAILED backup and cannot be restored."
  fi
  mapfile -t src < <(image_disks "$img")
  if [ "${#src[@]}" -ne 1 ]; then
    fail "The image $img holds ${#src[@]} disks (${src[*]:-none})." \
         "Restore it with Clonezilla: manual (boot menu)."
  fi
  srcb="$(image_disk_bytes "$img" "${src[0]}")"

  tgt="$(pick_disk "Which disk do you want to restore $img to? (ALL DATA ON IT WILL BE ERASED)")"
  [ -n "$tgt" ] || fail "No disk found to restore to."
  tgtb="$(disk_bytes "$tgt")"
  if [ "$srcb" -gt 0 ] && [ "$tgtb" -lt "$srcb" ]; then
    fail "The target disk is SMALLER than the disk the image came from." \
         "Image disk:  $(human "$srcb")" "Target disk: $(describe_disk "$tgt")" \
         "Use a disk at least as large, or Clonezilla: manual (expert mode, -icds)."
  fi

  banner "RESTORE machine $m" \
         "Image:   $img  (saved $(get_meta "$meta" date))" \
         "From:    $(get_meta "$meta" model)  serial $(get_meta "$meta" serial)  disk $(human "$srcb")" \
         "Target:  $(describe_disk "$tgt")" \
         "" \
         "EVERYTHING on /dev/$tgt will be ERASED and replaced by this image." \
         "This PC: $(dmi system-product-name)  serial $(dmi system-serial-number)"
  echo -n "  Type YES (capitals) to erase /dev/$tgt and restore: "
  read -r ans || ans=""
  if [ "$ans" != "YES" ]; then
    banner "Cancelled. Nothing was written."
    finish no
  fi

  # A larger target: -k1 recreates the partition table proportionally
  # (every partition scaled by target/source size) and -r grows each
  # filesystem to its new partition, so the extra space is used. The
  # default (-k0) copies the table as-is and leaves the rest unallocated.
  # The image was checked after it was saved; check again only if that
  # check did not pass.
  chkopt=""
  [ "$(get_meta "$meta" check)" = "ok" ] && chkopt="-scr"
  ocs-sr -g auto -e1 auto -e2 -r -j2 -k1 $chkopt -batch -p command restoredisk "$img" "$tgt"
  rc=$?
  if [ "$rc" -ne 0 ]; then
    fail "Clonezilla could not restore $img to /dev/$tgt (exit code $rc)." \
         "The disk may be partly written: do not boot from it until a restore succeeds."
  fi
  sync
  banner "SUCCESS: $img restored to /dev/$tgt." "At the next start, boot this PC from its disk (not the network)."
  finish yes
}

# --- main -----------------------------------------------------------------

MODE="$(echo "$CMDLINE" | sed -n 's/.* pxe_cz_mode=\([a-z]*\) .*/\1/p')"
case "$MODE" in
  save|restore) ;;
  *) fail "Unknown mode [$MODE] on the kernel command line (pxe_cz_mode=save or restore)." ;;
esac

if ! repo_ready; then
  fail "The backup share \\\\$SERVER\\clonezilla is not mounted on $OCSROOT." \
       "Check the PXE server is up and the clonezilla Samba account works" \
       "(playbook var clonezilla_smb_password), then boot this option again."
fi

sync_clock
MACHINE=""
get_machine
case "$MODE" in
  save)    do_save "$MACHINE" ;;
  restore) do_restore "$MACHINE" ;;
esac