#!/bin/bash
#
# prepare-boot-tools.sh — Download/extract boot files for PXE boot tools
#
# Downloads Clonezilla Live and Memtest86+ for PXE booting,
# and extracts Blancco Drive Eraser from its ISO.
#
# 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
#   boot-tools/memtest/      — memtest.efi

set -euo pipefail

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

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

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

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

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

CLONEZILLA_VERSION="3.2.1-6"
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" ]; 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/"
        rm -rf "$TMPDIR"
        echo "    Done."
    fi
fi

ls -lh "$OUT_DIR/clonezilla/" 2>/dev/null | grep -E 'vmlinuz|initrd|squashfs' | sed 's/^/    /'

# --- Blancco Drive Eraser ---
echo ""
echo "[2/3] 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 containing preferences.xml (padded to 32768 bytes).
        # The CPIO itself must remain exactly 194560 bytes (380 x 512-byte blocks).
        # We use Python for byte-level replacement to preserve exact sizes.
        if [ -f "$OUT_DIR/blancco/config.img" ]; then
            echo "    Patching config.img for network report storage..."
            CFGTMP=$(mktemp -d)
            cd "$CFGTMP"
            cpio -id < "$OUT_DIR/blancco/config.img" 2>/dev/null

            if [ -f "$CFGTMP/preferences.xml" ]; then
                ORIG_SIZE=$(stat -c%s "$CFGTMP/preferences.xml")

                python3 << 'PYEOF'
import sys

with open("preferences.xml", "rb") as f:
    data = f.read()

orig_size = len(data)

# Set SMB share credentials and path
data = data.replace(
    b'<username encrypted="false"></username>',
    b'<username encrypted="false">blancco</username>'
)
data = data.replace(
    b'<password encrypted="false"></password>',
    b'<password encrypted="false">blancco</password>'
)
data = data.replace(
    b'<hostname></hostname>',
    b'<hostname>10.9.100.1</hostname>'
)
data = data.replace(
    b'<path></path>',
    b'<path>blancco-reports</path>'
)

# Enable auto-backup
data = data.replace(
    b'<auto_backup>false</auto_backup>',
    b'<auto_backup>true</auto_backup>'
)

# Enable bootable report
data = data.replace(
    b'<bootable_report>\n      <enabled>false</enabled>\n    </bootable_report>',
    b'<bootable_report>\n      <enabled>true</enabled>\n    </bootable_report>'
)

# Maintain exact file size by trimming trailing padding/whitespace
diff = len(data) - orig_size
if diff > 0:
    # The file has trailing whitespace/padding before the final XML closing tags
    # Trim from the padding area (spaces before closing comment or end of file)
    end_pos = data.rfind(b'<!-- ')
    if end_pos > 0:
        comment_end = data.find(b' -->', end_pos)
        if comment_end > 0:
            data = data[:comment_end - diff] + data[comment_end:]
    if len(data) > orig_size:
        # Fallback: trim trailing whitespace
        data = data.rstrip()
        data = data + b'\n' * (orig_size - len(data))
elif diff < 0:
    # Pad with spaces to maintain size
    data = data[:-1] + b' ' * (-diff) + data[-1:]

if len(data) != orig_size:
    print(f"    WARNING: Size mismatch ({len(data)} vs {orig_size}), padding to match")
    if len(data) > orig_size:
        data = data[:orig_size]
    else:
        data = data + b'\x00' * (orig_size - len(data))

with open("preferences.xml", "wb") as f:
    f.write(data)

print(f"    preferences.xml: {orig_size} bytes (preserved)")
PYEOF

                # Repack CPIO with exact 512-byte block alignment (194560 bytes)
                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@10.9.100.1/blancco-reports, bootable report enabled"
            fi
            cd "$SCRIPT_DIR"
            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..."
        python3 << PYEOF
import lzma, sys, os

initramfs = "$OUT_DIR/blancco/initramfs-bde-linux.img"

with open(initramfs, "rb") as f:
    compressed = f.read()

# Decompress XZ stream
try:
    raw = lzma.decompress(compressed)
except lzma.LZMAError:
    print("    WARNING: Could not decompress initramfs (not XZ?), skipping patch")
    sys.exit(0)

# Binary patch: change the copytoram check from "y" to "N"
old = b'"y" ]; then\n            for curif in /sys/class/net'
new = b'"N" ]; then\n            for curif in /sys/class/net'

if old not in raw:
    # Try alternate pattern with different whitespace
    old = b'"\${copytoram}" = "y"'
    new = b'"\${copytoram}" = "N"'

if old in raw:
    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(recompressed)
    print("    initramfs patched: copytoram network flush disabled")
else:
    print("    WARNING: copytoram pattern not found in initramfs, skipping patch")
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="$SCRIPT_DIR/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 tftp chain sleep" \
                    "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/^/    /'

# --- Memtest86+ ---
echo ""
echo "[3/3] Memtest86+"

MEMTEST_VERSION="7.20"
MEMTEST_URL="https://memtest.org/download/${MEMTEST_VERSION}/mt86plus_${MEMTEST_VERSION}.binaries.zip"

if [ -f "$OUT_DIR/memtest/memtest.efi" ]; then
    echo "    Already prepared, skipping."
else
    echo "    Downloading Memtest86+ v${MEMTEST_VERSION}..."
    TMPDIR=$(mktemp -d)

    wget -q --show-progress -O "$TMPDIR/memtest.zip" "$MEMTEST_URL" || {
        echo "    ERROR: Download failed. Download manually from https://memtest.org"
        TMPDIR=""
    }

    if [ -n "$TMPDIR" ] && [ -f "$TMPDIR/memtest.zip" ]; then
        echo "    Extracting EFI binary..."
        unzip -o -j "$TMPDIR/memtest.zip" "memtest64.efi" -d "$OUT_DIR/memtest/" 2>/dev/null || \
        unzip -o -j "$TMPDIR/memtest.zip" "mt86plus_${MEMTEST_VERSION}.x64.efi" -d "$OUT_DIR/memtest/" 2>/dev/null || \
        unzip -o "$TMPDIR/memtest.zip" -d "$TMPDIR/extract/"

        # Find the EFI file regardless of exact name
        EFI_FILE=$(find "$TMPDIR" "$OUT_DIR/memtest" -name '*.efi' -name '*64*' 2>/dev/null | head -1)
        if [ -n "$EFI_FILE" ] && [ ! -f "$OUT_DIR/memtest/memtest.efi" ]; then
            cp "$EFI_FILE" "$OUT_DIR/memtest/memtest.efi"
        fi
        rm -rf "$TMPDIR"
        echo "    Done."
    fi
fi

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

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

for tool in clonezilla blancco memtest; 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/blancco/"
echo "  /var/www/html/memtest/"
echo ""
echo "The build-usb.sh script will include them automatically,"
echo "or copy them manually to the server."
echo ""