#!/bin/bash
#
# prepare-boot-tools.sh — Download/extract boot files for PXE boot tools
#
# Downloads Clonezilla Live for PXE booting (with Debian's Microsoft-signed
# shim, so it boots with Secure Boot on), and extracts Blancco Drive Eraser
# from its ISO. No memory tester: Dell PCs have one built in (F12 at
# power-on, Diagnostics).
#
# Usage:
#   ./prepare-boot-tools.sh [/path/to/blancco.iso]
#
# Output directories:
#   boot-tools/clonezilla/   — vmlinuz, initrd.img, filesystem.squashfs
#   boot-tools/blancco/      — extracted boot files or ISO for memdisk

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "$0")"/.. && pwd)"
OUT_DIR="${OUT_DIR:-$REPO_ROOT/boot-tools}"
BLANCCO_ISO="${1:-}"

# Auto-detect Blancco ISO in project directory
if [ -z "$BLANCCO_ISO" ]; then
    BLANCCO_ISO=$(find "$REPO_ROOT" -maxdepth 1 -name '*DriveEraser*.iso' -o -name '*blancco*.iso' 2>/dev/null | head -1)
fi

mkdir -p "$OUT_DIR"/{clonezilla,clonezilla-i686,blancco}

echo "============================================"
echo "PXE Boot Tools Preparation"
echo "============================================"

# --- Clonezilla Live ---
echo ""
echo "[1/2] Clonezilla Live"

CLONEZILLA_VERSION="3.3.3-37"
CLONEZILLA_FILE="clonezilla-live-${CLONEZILLA_VERSION}-amd64.zip"
CLONEZILLA_URL="https://sourceforge.net/projects/clonezilla/files/clonezilla_live_stable/${CLONEZILLA_VERSION}/${CLONEZILLA_FILE}/download"

if [ -f "$OUT_DIR/clonezilla/vmlinuz" ] && [ -f "$OUT_DIR/clonezilla/filesystem.squashfs" ] \
   && [ -f "$OUT_DIR/clonezilla/shimx64.efi" ]; then
    echo "    Already prepared, skipping. Delete boot-tools/clonezilla/ to re-download."
else
    echo "    Downloading Clonezilla Live ${CLONEZILLA_VERSION}..."
    TMPDIR=$(mktemp -d)

    wget -q --show-progress -O "$TMPDIR/$CLONEZILLA_FILE" "$CLONEZILLA_URL" || {
        echo "    ERROR: Download failed. Trying alternative URL..."
        # Fallback: try OSDN mirror
        wget -q --show-progress -O "$TMPDIR/$CLONEZILLA_FILE" \
            "https://free.nchc.org.tw/clonezilla-live/stable/${CLONEZILLA_FILE}" || {
            echo "    ERROR: Could not download Clonezilla. Download manually and place in boot-tools/clonezilla/"
            echo "    Need: vmlinuz, initrd.img, filesystem.squashfs from the live ZIP"
        }
    }

    if [ -f "$TMPDIR/$CLONEZILLA_FILE" ]; then
        echo "    Extracting PXE boot files..."
        unzip -o -j "$TMPDIR/$CLONEZILLA_FILE" "live/vmlinuz" -d "$OUT_DIR/clonezilla/"
        unzip -o -j "$TMPDIR/$CLONEZILLA_FILE" "live/initrd.img" -d "$OUT_DIR/clonezilla/"
        unzip -o -j "$TMPDIR/$CLONEZILLA_FILE" "live/filesystem.squashfs" -d "$OUT_DIR/clonezilla/"
        # Debian's Microsoft-signed shim (the release's EFI/boot/bootx64.efi):
        # the iPXE menu starts the Debian-signed kernel through it, so
        # Clonezilla runs with Secure Boot on (like Blancco via Ubuntu's).
        unzip -o -p "$TMPDIR/$CLONEZILLA_FILE" "EFI/boot/bootx64.efi" > "$OUT_DIR/clonezilla/shimx64.efi"
        rm -rf "$TMPDIR"
        echo "    Done."
    fi
fi

ls -lh "$OUT_DIR/clonezilla/" 2>/dev/null | grep -E 'vmlinuz|initrd|squashfs|shimx64' | sed 's/^/    /'
# The shim must be Microsoft-signed (firmware trusts it) and the kernel
# Debian-signed (the shim trusts that); otherwise Secure Boot refuses them.
if command -v sbverify >/dev/null && [ -s "$OUT_DIR/clonezilla/shimx64.efi" ]; then
    if sbverify --list "$OUT_DIR/clonezilla/shimx64.efi" 2>/dev/null | grep -q "Microsoft"; then
        echo "    shimx64.efi: Microsoft-signed"
    else
        echo "    WARNING: clonezilla/shimx64.efi is not Microsoft-signed; Clonezilla will need Secure Boot off."
    fi
    if ! sbverify --list "$OUT_DIR/clonezilla/vmlinuz" 2>/dev/null | grep -q "Debian Secure Boot"; then
        echo "    WARNING: clonezilla/vmlinuz is not Debian-signed; Clonezilla will need Secure Boot off."
    fi
fi

# --- Clonezilla Live 32-bit (XP-era PCs) ---
# Clonezilla 3.2.0-5 is the last release with a 32-bit (i686) build; the
# menu picks it for CPUs without 64-bit support (cpuid --ext 29). Legacy
# BIOS only, so no shim. Plain i686 (not -pae) for the oldest CPUs.
echo ""
echo "[1b/2] Clonezilla Live 32-bit"
CZ32_VERSION="3.2.0-5"
CZ32_FILE="clonezilla-live-${CZ32_VERSION}-i686.zip"
if [ -f "$OUT_DIR/clonezilla-i686/vmlinuz" ] && [ -f "$OUT_DIR/clonezilla-i686/filesystem.squashfs" ]; then
    echo "    Already prepared, skipping. Delete boot-tools/clonezilla-i686/ to re-download."
else
    TMPDIR=$(mktemp -d)
    if wget -q --show-progress -O "$TMPDIR/$CZ32_FILE" \
         "https://sourceforge.net/projects/clonezilla/files/clonezilla_live_stable/${CZ32_VERSION}/${CZ32_FILE}/download" \
       || wget -q --show-progress -O "$TMPDIR/$CZ32_FILE" \
         "https://free.nchc.org.tw/clonezilla-live/oldstable/${CZ32_FILE}"; then
        unzip -o -j "$TMPDIR/$CZ32_FILE" "live/vmlinuz" "live/initrd.img" "live/filesystem.squashfs" \
            -d "$OUT_DIR/clonezilla-i686/"
        echo "    Done."
    else
        echo "    ERROR: could not download $CZ32_FILE; 32-bit PCs cannot run Clonezilla until it is staged."
    fi
    rm -rf "$TMPDIR"
fi
ls -lh "$OUT_DIR/clonezilla-i686/" 2>/dev/null | grep -E 'vmlinuz|initrd|squashfs' | sed 's/^/    /'

# --- Blancco Drive Eraser ---
echo ""
echo "[2/2] Blancco Drive Eraser"

if [ -n "$BLANCCO_ISO" ] && [ -f "$BLANCCO_ISO" ]; then
    echo "    Extracting from: $BLANCCO_ISO"
    echo "    Using 7z to extract (no root required)..."

    # Blancco is Arch Linux-based. We need:
    #   arch/boot/x86_64/vmlinuz-bde-linux
    #   arch/boot/x86_64/initramfs-bde-linux.img
    #   arch/boot/intel-ucode.img
    #   arch/boot/amd-ucode.img
    #   arch/boot/config.img
    #   arch/x86_64/airootfs.sfs
    TMPDIR=$(mktemp -d)
    7z x -o"$TMPDIR" "$BLANCCO_ISO" \
        "arch/boot/x86_64/vmlinuz-bde-linux" \
        "arch/boot/x86_64/initramfs-bde-linux.img" \
        "arch/boot/intel-ucode.img" \
        "arch/boot/amd-ucode.img" \
        "arch/boot/config.img" \
        "arch/x86_64/airootfs.sfs" \
        -r 2>/dev/null || {
        echo "    7z extraction failed. Install p7zip-full: apt install p7zip-full"
    }

    # Flatten into blancco/ directory for HTTP serving
    if [ -f "$TMPDIR/arch/boot/x86_64/vmlinuz-bde-linux" ]; then
        cp "$TMPDIR/arch/boot/x86_64/vmlinuz-bde-linux" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/x86_64/initramfs-bde-linux.img" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/intel-ucode.img" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/amd-ucode.img" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/config.img" "$OUT_DIR/blancco/"
        # airootfs.sfs needs to be in arch/x86_64/ path relative to HTTP root
        mkdir -p "$OUT_DIR/blancco/arch/x86_64"
        cp "$TMPDIR/arch/x86_64/airootfs.sfs" "$OUT_DIR/blancco/arch/x86_64/"
        echo "    Extracted Blancco boot files."

        # --- Patch config.img (size-preserving) ---
        # config.img is a CPIO archive. Its preferences.xml is what the native
        # kernel (BIOS) boot path reads; the UEFI path gets the same file over
        # HTTP. Both now carry playbook/blancco-preferences.xml, so the two
        # paths cannot drift. The vendor file is laid out as
        # <xml> [NUL padding] <!-- filler comment --> and its size is kept:
        # 7.20 pads to 32768 with ~15KB of NULs; 7.15 had none, so growth
        # there comes out of the filler comment.
        if [ -f "$OUT_DIR/blancco/config.img" ]; then
            echo "    Writing playbook preferences into config.img..."
            CFGTMP=$(mktemp -d)
            cd "$CFGTMP"
            cpio -id < "$OUT_DIR/blancco/config.img" 2>/dev/null

            if [ -f "$CFGTMP/preferences.xml" ]; then
                PREFS_SRC="$REPO_ROOT/playbook/blancco-preferences.xml" python3 << 'PYEOF'
import os, re, sys

with open("preferences.xml", "rb") as f:
    data = f.read()
with open(os.environ["PREFS_SRC"], "rb") as f:
    ours = f.read().strip() + b"\n"
orig_size = len(data)

def version(x):
    m = re.search(rb"<version>\s*<generation>(\d+)</generation>\s*"
                  rb"<major>(\d+)</major>", x)
    return m.groups() if m else None

# Our file must be a merge onto THIS release's preferences, or settings the
# new schema added or renamed are silently lost.
if version(ours) != version(data):
    sys.exit(f"    ERROR: playbook preferences are for {version(ours)}, ISO is "
             f"{version(data)}. Merge onto the new vendor preferences first.")
for need in (b"<hostname>172.16.9.1</hostname>", b"<path>blancco-reports</path>",
             b"<auto_backup>true</auto_backup>"):
    if need not in ours:
        sys.exit(f"    ERROR: playbook preferences missing {need.decode()}")

# Split off the trailing filler comment and any NUL padding before it.
end = data.find(b"</root>")
cstart = data.rfind(b"<!--")
if end < 0 or cstart < end:
    sys.exit("    ERROR: vendor preferences.xml layout not recognised")
gap, tail = data[end + 7:cstart], data[cstart:]
xml = ours

room = orig_size - len(xml) - len(tail)
if room >= 0:
    # Refill the gap with what it held: NULs if it was padding, else newlines.
    gap = (b"\x00" if b"\x00" in gap else b"\n") * room
else:
    # No room left: take it out of the filler comment body (keep "<!-- " + " -->").
    cut = -room
    if len(tail) - cut < len(b"<!--  -->\n"):
        sys.exit(f"    ERROR: preferences need {cut} more bytes than the "
                 "padding and filler comment can give")
    tail = tail[:5] + tail[5 + cut:]
    gap = b""
out = xml + gap + tail
if len(out) != orig_size:
    sys.exit(f"    ERROR: preferences.xml size changed {orig_size} -> {len(out)}")
with open("preferences.xml", "wb") as f:
    f.write(out)
print(f"    preferences.xml: {len(xml)} bytes of XML, {orig_size} total (preserved)")
PYEOF

                # Repack CPIO padded to a 512-byte boundary
                ls -1 "$CFGTMP" | (cd "$CFGTMP" && cpio -o -H newc 2>/dev/null) | \
                    dd bs=512 conv=sync 2>/dev/null > "$OUT_DIR/blancco/config.img"
                echo "    Reports: SMB blancco@172.16.9.1/blancco-reports (from playbook preferences)"
            fi
            cd "$REPO_ROOT"
            rm -rf "$CFGTMP"
        fi

        # --- Patch initramfs to keep network interfaces up after copytoram ---
        # Blancco uses copytoram=y which triggers archiso_pxe_common latehook to
        # flush all network interfaces. Binary-patch the check from "y" to "N" so
        # the condition never matches. IMPORTANT: full extract/repack BREAKS booting.
        echo "    Patching initramfs to preserve network after copytoram..."
        INITRAMFS="$OUT_DIR/blancco/initramfs-bde-linux.img" python3 << 'PYEOF'
import lzma, os, sys

initramfs = os.environ["INITRAMFS"]
with open(initramfs, "rb") as f:
    data = f.read()

# 7.20 prepends an uncompressed early cpio (NIC firmware) before the XZ
# archive; 7.15 was XZ from byte 0. Patch only the XZ part, keep the prefix.
# Walk the cpio headers to its trailer: the firmware files are themselves
# .xz, so searching for the XZ magic would land inside the prefix.
off = 0
while data.startswith(b"070701", off):
    namesize = int(data[off + 94:off + 102], 16)
    filesize = int(data[off + 54:off + 62], 16)
    name = data[off + 110:off + 110 + namesize - 1]
    off = (((off + 110 + namesize + 3) & ~3) + filesize + 3) & ~3
    if name == b"TRAILER!!!":
        while off < len(data) and data[off] == 0:
            off += 1
if not data.startswith(b"\xfd7zXZ\x00", off):
    sys.exit(f"    ERROR: initramfs layout not recognised (no XZ archive at {off})")
prefix, compressed = data[:off], data[off:]

raw = lzma.decompress(compressed)

# Binary patch: change the copytoram check from "y" to "N" in
# archiso_pxe_common's run_latehook. Exact match only: the looser
# '"${copytoram}" = "y"' appears in several hooks and patching the first
# one found is not safe.
old = b'"y" ]; then\n            for curif in /sys/class/net'
new = b'"N" ]; then\n            for curif in /sys/class/net'
count = raw.count(old)
if count != 1:
    sys.exit(f"    ERROR: copytoram network-flush pattern found {count} times, "
             "expected 1. PXE boot would lose network after copytoram.")
raw = raw.replace(old, new, 1)

# Recompress with same XZ settings as archiso
recompressed = lzma.compress(raw, format=lzma.FORMAT_XZ,
                              check=lzma.CHECK_CRC32, preset=6)
with open(initramfs, "wb") as f:
    f.write(prefix + recompressed)
print(f"    initramfs patched: copytoram network flush disabled "
      f"(early cpio {len(prefix)} bytes kept)")
PYEOF

        # --- Build GRUB EFI binary for Blancco chainload ---
        # Broadcom iPXE can't pass initrd to Linux kernels in UEFI mode.
        # Solution: iPXE chains to grubx64.efi, which loads kernel+initrd via TFTP.
        GRUB_CFG="$REPO_ROOT/boot-tools/blancco/grub-blancco.cfg"
        if [ -f "$GRUB_CFG" ]; then
            if command -v grub-mkstandalone &>/dev/null; then
                echo "    Building grubx64.efi (GRUB chainload for Blancco)..."
                grub-mkstandalone \
                    --format=x86_64-efi \
                    --output="$OUT_DIR/blancco/grubx64.efi" \
                    --modules="linux normal echo net efinet http tftp chain sleep all_video efi_gop" \
                    "boot/grub/grub.cfg=$GRUB_CFG" 2>/dev/null
                echo "    Built grubx64.efi ($(du -h "$OUT_DIR/blancco/grubx64.efi" | cut -f1))"
            else
                echo "    WARNING: grub-mkstandalone not found. Install grub-efi-amd64-bin:"
                echo "             sudo apt install grub-efi-amd64-bin grub-common"
                echo "             Then re-run this script to build grubx64.efi"
            fi
        else
            echo "    WARNING: grub-blancco.cfg not found at $GRUB_CFG"
        fi
    else
        echo "    Could not extract boot files from ISO."
    fi
    rm -rf "$TMPDIR"
else
    echo "    No Blancco ISO found. Provide path as argument or place in project directory."
    echo "    Usage: $0 /path/to/DriveEraser.iso"
fi

ls -lh "$OUT_DIR/blancco/" 2>/dev/null | grep -v '^total' | sed 's/^/    /'

# --- Summary ---
echo ""
echo "============================================"
echo "Boot tools prepared in: $OUT_DIR/"
echo "============================================"
echo ""

for tool in clonezilla clonezilla-i686 blancco; do
    COUNT=$(find "$OUT_DIR/$tool" -type f 2>/dev/null | wc -l)
    SIZE=$(du -sh "$OUT_DIR/$tool" 2>/dev/null | cut -f1)
    printf "  %-15s %s (%d files)\n" "$tool" "$SIZE" "$COUNT"
done

echo ""
echo "These files need to be copied to the PXE server's web root:"
echo "  /var/www/html/clonezilla/"
echo "  /var/www/html/clonezilla-i686/"
echo "  /var/www/html/blancco/"
echo ""
echo "The build-usb.sh script will include them automatically,"
echo "or copy them manually to the server."
echo ""