pxe-server/scripts/build-usb.sh

#!/bin/bash
#
# build-usb.sh — Build a bootable PXE-server installer USB
#
# Creates a two-partition USB:
#   Partition 1: Ubuntu Server 24.04 installer (ISO contents)
#   Partition 2: CIDATA volume (autoinstall config, .debs, playbook)
#
# The target machine boots from this USB, Ubuntu auto-installs with
# cloud-init (user-data/meta-data from CIDATA), installs offline .debs,
# and on first boot runs the Ansible playbook to configure PXE services.
#
# Usage:
#   sudo ./build-usb.sh /dev/sdX /path/to/ubuntu-24.04-live-server-amd64.iso
#
# WARNING: This will ERASE the target USB device.

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "$0")"/.. && pwd)"
AUTOINSTALL_DIR="$REPO_ROOT/autoinstall"
PLAYBOOK_DIR="$REPO_ROOT/playbook"
OFFLINE_PKG_DIR="$REPO_ROOT/offline-packages"

# --- Validate arguments ---
if [ $# -lt 2 ]; then
    echo "Usage: sudo $0 /dev/sdX /path/to/ubuntu-24.04.iso [/path/to/winpe-images]"
    echo ""
    echo "  The optional third argument is the path to WinPE deployment content"
    echo "  (e.g., the mounted Media Creator LITE USB). If provided, the images"
    echo "  will be bundled onto the CIDATA partition for automatic import."
    echo ""
    echo "Available removable devices:"
    lsblk -d -o NAME,SIZE,TRAN,RM | grep -E '^\S+\s+\S+\s+(usb)\s+1'
    exit 1
fi

USB_DEV="$1"
ISO_PATH="$2"
WINPE_SOURCE="${3:-}"

# Safety checks
if [ "$(id -u)" -ne 0 ]; then
    echo "ERROR: Must run as root (sudo)."
    exit 1
fi

if [ ! -b "$USB_DEV" ]; then
    echo "ERROR: $USB_DEV is not a block device."
    exit 1
fi

if [ ! -f "$ISO_PATH" ]; then
    echo "ERROR: ISO not found at $ISO_PATH"
    exit 1
fi

# Verify it's a removable device (safety against wiping system disks)
REMOVABLE=$(lsblk -nd -o RM "$USB_DEV" 2>/dev/null || echo "0")
if [ "$REMOVABLE" != "1" ]; then
    echo "WARNING: $USB_DEV does not appear to be a removable device."
    read -rp "Are you SURE you want to erase $USB_DEV? (type YES): " CONFIRM
    if [ "$CONFIRM" != "YES" ]; then
        echo "Aborted."
        exit 1
    fi
fi

# Verify required source files exist
if [ ! -f "$AUTOINSTALL_DIR/user-data" ]; then
    echo "ERROR: user-data not found at $AUTOINSTALL_DIR/user-data"
    exit 1
fi

if [ ! -f "$AUTOINSTALL_DIR/meta-data" ]; then
    echo "ERROR: meta-data not found at $AUTOINSTALL_DIR/meta-data"
    exit 1
fi

if [ ! -f "$PLAYBOOK_DIR/pxe_server_setup.yml" ]; then
    echo "ERROR: pxe_server_setup.yml not found at $PLAYBOOK_DIR/"
    exit 1
fi

echo "============================================"
echo "PXE Server USB Builder"
echo "============================================"
echo "USB Device : $USB_DEV"
echo "ISO        : $ISO_PATH"
echo "Source Dir : $REPO_ROOT"
echo ""
echo "This will ERASE all data on $USB_DEV."
read -rp "Continue? (y/N): " PROCEED
if [[ ! "$PROCEED" =~ ^[Yy]$ ]]; then
    echo "Aborted."
    exit 1
fi

# --- Unmount any existing partitions ---
echo ""
echo "[1/6] Unmounting existing partitions on $USB_DEV..."
for part in "${USB_DEV}"*; do
    umount "$part" 2>/dev/null || true
done

# --- Rebuild ISO with 'autoinstall' kernel parameter ---
echo "[2/6] Rebuilding ISO with autoinstall kernel parameter..."

ISO_WORK=$(mktemp -d)
7z -y x "$ISO_PATH" -o"$ISO_WORK/iso" >/dev/null 2>&1
mv "$ISO_WORK/iso/[BOOT]" "$ISO_WORK/BOOT"
chmod -R u+w "$ISO_WORK/iso"

# Patch grub.cfg: add 'autoinstall' to kernel cmdline, reduce timeout
sed -i 's|linux\t/casper/vmlinuz  ---|linux\t/casper/vmlinuz autoinstall ---|' "$ISO_WORK/iso/boot/grub/grub.cfg"
sed -i 's/^set timeout=30/set timeout=5/' "$ISO_WORK/iso/boot/grub/grub.cfg"

PATCHED_ISO="$ISO_WORK/patched.iso"
cd "$ISO_WORK/iso"
xorriso -as mkisofs -r \
    -V 'Ubuntu-Server 24.04.3 LTS amd64' \
    -o "$PATCHED_ISO" \
    --grub2-mbr ../BOOT/1-Boot-NoEmul.img \
    --protective-msdos-label \
    -partition_cyl_align off \
    -partition_offset 16 \
    --mbr-force-bootable \
    -append_partition 2 28732ac11ff8d211ba4b00a0c93ec93b ../BOOT/2-Boot-NoEmul.img \
    -appended_part_as_gpt \
    -iso_mbr_part_type a2a0d0ebe5b9334487c068b6b72699c7 \
    -c '/boot.catalog' \
    -b '/boot/grub/i386-pc/eltorito.img' \
    -no-emul-boot \
    -boot-load-size 4 \
    -boot-info-table \
    --grub2-boot-info \
    -eltorito-alt-boot \
    -e '--interval:appended_partition_2:::' \
    -no-emul-boot \
    -boot-load-size 10160 \
    . 2>/dev/null
cd "$REPO_ROOT"
echo "    ISO rebuilt with 'autoinstall' kernel param and 5s GRUB timeout"

echo "    Writing patched ISO to $USB_DEV..."
dd if="$PATCHED_ISO" of="$USB_DEV" bs=4M status=progress oflag=sync
sync
ISO_SIZE=$(stat -c%s "$PATCHED_ISO")
rm -rf "$ISO_WORK"

# --- Find the end of the ISO to create CIDATA partition ---
echo "[3/6] Creating CIDATA partition after ISO data..."

# Get ISO size in bytes and calculate the start sector for the new partition
SECTOR_SIZE=512
# Start the CIDATA partition 1MB after the ISO ends (alignment)
START_SECTOR=$(( (ISO_SIZE / SECTOR_SIZE) + 2048 ))

# Use sfdisk to append a new partition
echo "    ISO size: $((ISO_SIZE / 1024 / 1024)) MB"
echo "    CIDATA partition starts at sector $START_SECTOR"

# Add a new partition using sfdisk --append
echo "${START_SECTOR},+,L" | sfdisk --append "$USB_DEV" --no-reread 2>/dev/null || true
partprobe "$USB_DEV"
sleep 2

# Determine the new partition name (could be sdX3, sdX4, etc.)
CIDATA_PART=""
for part in "${USB_DEV}"[0-9]*; do
    # Find the partition that starts at or after our start sector
    PART_START=$(sfdisk -d "$USB_DEV" 2>/dev/null | grep "$part" | grep -o 'start=[[:space:]]*[0-9]*' | grep -o '[0-9]*')
    if [ -n "$PART_START" ] && [ "$PART_START" -ge "$START_SECTOR" ]; then
        CIDATA_PART="$part"
        break
    fi
done

# Fallback: use the last partition
if [ -z "$CIDATA_PART" ]; then
    CIDATA_PART=$(lsblk -ln -o NAME "$USB_DEV" | tail -1)
    CIDATA_PART="/dev/$CIDATA_PART"
fi

echo "    CIDATA partition: $CIDATA_PART"

# --- Format CIDATA partition ---
echo "[4/6] Formatting $CIDATA_PART as FAT32 (label: CIDATA)..."
mkfs.vfat -F 32 -n CIDATA "$CIDATA_PART"

# --- Mount and copy files ---
echo "[5/6] Copying autoinstall config, packages, and playbook to CIDATA..."
MOUNT_POINT=$(mktemp -d)
mount "$CIDATA_PART" "$MOUNT_POINT"

# Copy cloud-init files
cp "$AUTOINSTALL_DIR/user-data" "$MOUNT_POINT/"
cp "$AUTOINSTALL_DIR/meta-data" "$MOUNT_POINT/"

# Copy offline .deb packages into packages/ subdirectory
mkdir -p "$MOUNT_POINT/packages"
DEB_COUNT=0

if [ -d "$OFFLINE_PKG_DIR" ]; then
    for deb in "$OFFLINE_PKG_DIR"/*.deb; do
        if [ -f "$deb" ]; then
            cp "$deb" "$MOUNT_POINT/packages/"
            DEB_COUNT=$((DEB_COUNT + 1))
        fi
    done
fi
echo "    Copied $DEB_COUNT .deb packages to packages/"

# Copy playbook directory
cp -r "$PLAYBOOK_DIR" "$MOUNT_POINT/playbook"
echo "    Copied playbook/"

# Copy webapp
WEBAPP_DIR="$REPO_ROOT/webapp"
if [ -d "$WEBAPP_DIR" ]; then
    mkdir -p "$MOUNT_POINT/webapp"
    cp -r "$WEBAPP_DIR/app.py" "$WEBAPP_DIR/requirements.txt" "$MOUNT_POINT/webapp/"
    cp -r "$WEBAPP_DIR/templates" "$WEBAPP_DIR/static" "$MOUNT_POINT/webapp/"
    echo "    Copied webapp/"
fi

# Copy pip wheels for offline Flask install
PIP_WHEELS_DIR="$REPO_ROOT/pip-wheels"
if [ ! -d "$PIP_WHEELS_DIR" ]; then
    echo "    pip-wheels/ not found — downloading now..."
    mkdir -p "$PIP_WHEELS_DIR"
    if pip3 download -d "$PIP_WHEELS_DIR" flask lxml 2>/dev/null; then
        echo "    Downloaded pip wheels successfully."
    else
        echo "    WARNING: Failed to download pip wheels (no internet?)"
        echo "    The PXE server will need internet to install Flask later,"
        echo "    or manually copy wheels to pip-wheels/ and rebuild."
        rmdir "$PIP_WHEELS_DIR" 2>/dev/null || true
    fi
fi
if [ -d "$PIP_WHEELS_DIR" ]; then
    cp -r "$PIP_WHEELS_DIR" "$MOUNT_POINT/pip-wheels"
    WHEEL_COUNT=$(find "$PIP_WHEELS_DIR" -name '*.whl' | wc -l)
    echo "    Copied pip-wheels/ ($WHEEL_COUNT wheels)"
fi

# Copy WinPE boot files (wimboot, boot.wim, BCD, ipxe.efi, etc.)
BOOT_FILES_DIR="$REPO_ROOT/boot-files"
if [ -d "$BOOT_FILES_DIR" ]; then
    BOOT_FILE_COUNT=0
    for bf in "$BOOT_FILES_DIR"/*; do
        if [ -f "$bf" ]; then
            cp "$bf" "$MOUNT_POINT/"
            BOOT_FILE_COUNT=$((BOOT_FILE_COUNT + 1))
        fi
    done
    BOOT_FILES_SIZE=$(du -sh "$BOOT_FILES_DIR" | cut -f1)
    echo "    Copied $BOOT_FILE_COUNT boot files ($BOOT_FILES_SIZE) — wimboot, boot.wim, ipxe.efi, etc."
else
    echo "    WARNING: No boot-files/ found (copy WinPE boot files from Media Creator)"
fi

# Copy boot tools (Clonezilla, Blancco, Memtest) if prepared
BOOT_TOOLS_DIR="$REPO_ROOT/boot-tools"
if [ -d "$BOOT_TOOLS_DIR" ]; then
    cp -r "$BOOT_TOOLS_DIR" "$MOUNT_POINT/boot-tools"
    TOOLS_SIZE=$(du -sh "$MOUNT_POINT/boot-tools" | cut -f1)
    echo "    Copied boot-tools/ ($TOOLS_SIZE)"
else
    echo "    No boot-tools/ found (run prepare-boot-tools.sh first)"
fi

# Copy enrollment directory (PPKGs, run-enrollment.ps1) if present
# FAT32 has a 4GB max file size; files larger than that are split into chunks
# that the playbook reassembles with `cat`.
ENROLLMENT_DIR="$REPO_ROOT/enrollment"
FAT32_MAX=$((3500 * 1024 * 1024))   # 3500 MiB chunks, safely under 4GiB FAT32 limit
if [ -d "$ENROLLMENT_DIR" ]; then
    mkdir -p "$MOUNT_POINT/enrollment"
    SPLIT_COUNT=0
    for f in "$ENROLLMENT_DIR"/*; do
        [ -e "$f" ] || continue
        bn="$(basename "$f")"
        if [ -f "$f" ] && [ "$(stat -c%s "$f")" -gt "$FAT32_MAX" ]; then
            echo "    Splitting $bn (>$((FAT32_MAX / 1024 / 1024))M) into chunks..."
            split -b "$FAT32_MAX" -d -a 2 "$f" "$MOUNT_POINT/enrollment/${bn}.part."
            SPLIT_COUNT=$((SPLIT_COUNT + 1))
        else
            cp -r "$f" "$MOUNT_POINT/enrollment/"
        fi
    done
    PPKG_COUNT=$(find "$ENROLLMENT_DIR" -maxdepth 1 -name '*.ppkg' 2>/dev/null | wc -l)
    ENROLL_SIZE=$(du -sh "$MOUNT_POINT/enrollment" | cut -f1)
    echo "    Copied enrollment/ ($ENROLL_SIZE, $PPKG_COUNT PPKGs, $SPLIT_COUNT split)"
else
    echo "    No enrollment/ directory found (PPKGs can be uploaded via webapp later)"
fi

# Copy BIOS update binaries if staged
BIOS_DIR="$REPO_ROOT/bios-staging"
if [ -d "$BIOS_DIR" ] && [ "$(ls -A "$BIOS_DIR" 2>/dev/null)" ]; then
    echo "    Copying BIOS update binaries from bios-staging/..."
    mkdir -p "$MOUNT_POINT/bios"
    cp -r "$BIOS_DIR"/* "$MOUNT_POINT/bios/" 2>/dev/null || true
    BIOS_COUNT=$(find "$MOUNT_POINT/bios" -name '*.exe' 2>/dev/null | wc -l)
    BIOS_SIZE=$(du -sh "$MOUNT_POINT/bios" | cut -f1)
    echo "    Copied bios/ ($BIOS_SIZE, $BIOS_COUNT files)"
else
    echo "    No bios-staging/ found (BIOS updates can be pushed via download-drivers.py later)"
fi

# Copy Dell driver packs if staged
# Files larger than the FAT32 4GB limit are split into chunks; the playbook
# reassembles them on the server.
DRIVERS_DIR="$REPO_ROOT/drivers-staging"
if [ -d "$DRIVERS_DIR" ] && [ "$(ls -A "$DRIVERS_DIR" 2>/dev/null)" ]; then
    echo "    Copying Dell driver packs from drivers-staging/..."
    mkdir -p "$MOUNT_POINT/drivers"
    DRV_SPLIT=0
    # Mirror directory tree first (fast)
    (cd "$DRIVERS_DIR" && find . -type d -exec mkdir -p "$MOUNT_POINT/drivers/{}" \;)
    # Copy files <4GB directly, split files >=4GB into chunks
    while IFS= read -r f; do
        rel="${f#$DRIVERS_DIR/}"
        dest="$MOUNT_POINT/drivers/$rel"
        if [ "$(stat -c%s "$f")" -gt "$FAT32_MAX" ]; then
            echo "      Splitting $rel..."
            split -b "$FAT32_MAX" -d -a 2 "$f" "${dest}.part."
            DRV_SPLIT=$((DRV_SPLIT + 1))
        else
            cp "$f" "$dest"
        fi
    done < <(find "$DRIVERS_DIR" -type f)
    DRIVERS_SIZE=$(du -sh "$MOUNT_POINT/drivers" | cut -f1)
    echo "    Copied drivers/ ($DRIVERS_SIZE, $DRV_SPLIT split)"
else
    echo "    No drivers-staging/ found (drivers can be downloaded later)"
fi

# Optionally copy WinPE deployment images
if [ -n "$WINPE_SOURCE" ] && [ -d "$WINPE_SOURCE" ]; then
    echo "    Copying WinPE deployment content from $WINPE_SOURCE..."
    mkdir -p "$MOUNT_POINT/images"
    cp -r "$WINPE_SOURCE"/* "$MOUNT_POINT/images/" 2>/dev/null || true
    IMG_SIZE=$(du -sh "$MOUNT_POINT/images" | cut -f1)
    echo "    Copied WinPE images ($IMG_SIZE)"
elif [ -n "$WINPE_SOURCE" ]; then
    echo "    WARNING: WinPE source path not found: $WINPE_SOURCE (skipping)"
fi

# List what's on CIDATA
echo ""
echo "    CIDATA contents:"
ls -lh "$MOUNT_POINT/" | sed 's/^/      /'

# --- Cleanup ---
echo ""
echo "[6/6] Syncing and unmounting..."
sync
umount "$MOUNT_POINT"
rmdir "$MOUNT_POINT"

echo ""
echo "============================================"
echo "USB build complete!"
echo "============================================"
echo ""
echo "Next steps:"
echo "  1. Insert USB into target machine"
echo "  2. Boot from USB (F12 / boot menu)"
echo "  3. Ubuntu will auto-install and configure the PXE server"
echo "  4. After reboot, move the NIC to the isolated PXE network"
echo ""