lab/bastion/tests/integration/helpers/pxe-vm.ts

// Create a blank UEFI VM for PXE boot testing.
// Unlike cloud image VMs, these have an empty disk and boot from network.
// Each VM gets a serial console on a TCP port for debugging without network/SSH.

import { execSync, spawnSync, type SpawnSyncReturns } from "node:child_process";
import { existsSync } from "node:fs";
import { join } from "node:path";
import { createConnection } from "node:net";
import { log } from "./libvirt.js";

const IMAGE_DIR = "/var/lib/libvirt/images";
const IS_ROOT = process.getuid?.() === 0;

function run(cmd: string, opts?: { timeout?: number }): string {
  const full = IS_ROOT ? cmd : `sudo ${cmd}`;
  return execSync(full, { encoding: "utf-8", stdio: "pipe", timeout: opts?.timeout ?? 60_000 });
}

function virsh(...args: string[]): SpawnSyncReturns<string> {
  const cmd = IS_ROOT ? "virsh" : "sudo";
  const finalArgs = IS_ROOT ? args : ["virsh", ...args];
  return spawnSync(cmd, finalArgs, { encoding: "utf-8", stdio: "pipe", timeout: 30_000 });
}

export interface PxeVmConfig {
  name: string;
  memory: number;       // MB
  vcpus: number;
  diskSize: number;     // GB
  network: string;      // libvirt network name
  arch?: "x86_64" | "aarch64";
}

/** Create a blank UEFI VM that PXE boots from the network. */
export function createPxeVm(config: PxeVmConfig): void {
  destroyPxeVm(config.name);

  const arch = config.arch ?? "x86_64";
  log(`Creating PXE VM: ${config.name} (${arch}, ${config.memory}MB RAM, ${config.vcpus} vCPU, ${config.diskSize}GB disk)`);

  // Create blank disk
  const diskPath = join(IMAGE_DIR, `${config.name}.qcow2`);
  run(`qemu-img create -f qcow2 "${diskPath}" ${config.diskSize}G`);

  // UEFI firmware paths (Fedora)
  if (arch === "aarch64") {
    const aavmf = "/usr/share/edk2/aarch64/QEMU_EFI.fd";
    if (!existsSync(aavmf)) {
      throw new Error(`AAVMF firmware not found at ${aavmf}. Install: sudo dnf install edk2-aarch64`);
    }
  } else {
    const ovmf = "/usr/share/edk2/ovmf/OVMF_CODE.fd";
    if (!existsSync(ovmf)) {
      throw new Error(`OVMF firmware not found at ${ovmf}. Install: sudo dnf install edk2-ovmf`);
    }
  }

  const virtInstallArgs = [
    "virt-install",
    `--name=${config.name}`,
    `--memory=${config.memory}`,
    `--vcpus=${config.vcpus}`,
    `--disk=path=${diskPath},format=qcow2,bus=virtio`,
    `--network=network=${config.network},model=virtio`,
    // UEFI firmware — required for PXE boot in modern mode
    `--boot=uefi,network`,
    // No OS to install — PXE provides everything
    "--os-variant=generic",
    "--noautoconsole",
    "--wait=0",
    // Graphics for debugging (VNC, connect with virt-viewer if needed)
    "--graphics=vnc,listen=127.0.0.1",
    // Serial console via TCP — allows exec without network/SSH
    // Connect: socat - TCP:127.0.0.1:4555
    "--serial=tcp,host=127.0.0.1:4555,mode=bind,protocol=telnet",
  ];

  if (arch === "aarch64") {
    virtInstallArgs.push("--arch=aarch64", "--machine=virt");
  }

  log(`Running: virt-install --name=${config.name} --boot=uefi,network ...`);
  run(virtInstallArgs.join(" "), { timeout: 30_000 });
  log(`PXE VM ${config.name} created (serial: telnet 127.0.0.1 4555)`);
}

/** Destroy a PXE VM and clean up its disk. */
export function destroyPxeVm(name: string): void {
  const result = virsh("dominfo", name);
  if (result.status !== 0) return;

  log(`Destroying PXE VM: ${name}`);
  virsh("destroy", name);
  virsh("undefine", name, "--remove-all-storage", "--nvram");
}

/** Get the MAC address of a VM's first NIC. */
export function getVmMac(name: string): string | null {
  const result = virsh("domiflist", name);
  if (result.status !== 0) return null;
  // Output format: Interface  Type  Source  Model  MAC
  const match = result.stdout.match(/([0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2})/i);
  return match ? match[1].toLowerCase() : null;
}

/** Reboot a VM (force off + start). */
export function rebootPxeVm(name: string): void {
  log(`Rebooting PXE VM: ${name}`);
  virsh("destroy", name);
  // Brief pause to let resources release
  spawnSync("sleep", ["2"]);
  virsh("start", name);
  log(`PXE VM ${name} restarted`);
}

/** Change VM boot order to disk first (skip PXE on next boot). */
export function setBootDisk(name: string): void {
  log(`Setting ${name} boot order to disk first`);
  virsh("destroy", name);
  spawnSync("sleep", ["2"]);
  // Get current XML, replace boot dev='network' with boot dev='hd'
  // This preserves UEFI loader/nvram settings (virt-xml --boot hd can break them)
  const dumpXml = virsh("dumpxml", name);
  if (dumpXml.status !== 0) throw new Error("Failed to dump VM XML");
  let xml = dumpXml.stdout;
  // Replace any <boot dev='...' /> entries with hd
  xml = xml.replace(/<boot dev='[^']*'\/>/g, "<boot dev='hd'/>");
  // If no boot dev entry, add one before </os>
  if (!xml.includes("<boot dev=")) {
    xml = xml.replace("</os>", "  <boot dev='hd'/>\n  </os>");
  }
  const xmlPath = `/tmp/${name}-bootfix.xml`;
  const { writeFileSync: writeFs, unlinkSync: unlinkFs } = require("node:fs") as typeof import("node:fs");
  writeFs(xmlPath, xml);
  run(`virsh define "${xmlPath}"`);
  try { unlinkFs(xmlPath); } catch { /* ignore */ }
  virsh("start", name);
  log(`${name} restarted with disk boot (UEFI preserved)`);
}

export interface IsoVmConfig {
  name: string;
  memory: number;       // MB
  vcpus: number;
  diskSize: number;     // GB
  network: string;      // libvirt network name
  isoPath: string;      // path to boot ISO
  arch?: "x86_64" | "aarch64";
}

/** Create a UEFI VM that boots from a CD-ROM ISO (not PXE). */
export function createIsoVm(config: IsoVmConfig): void {
  destroyPxeVm(config.name);

  const arch = config.arch ?? "x86_64";
  log(`Creating ISO boot VM: ${config.name} (${arch}, ${config.memory}MB RAM, ${config.vcpus} vCPU, ${config.diskSize}GB disk)`);

  // Create blank disk
  const diskPath = join(IMAGE_DIR, `${config.name}.qcow2`);
  run(`qemu-img create -f qcow2 "${diskPath}" ${config.diskSize}G`);

  const virtInstallArgs = [
    "virt-install",
    `--name=${config.name}`,
    `--memory=${config.memory}`,
    `--vcpus=${config.vcpus}`,
    `--disk=path=${diskPath},format=qcow2,bus=virtio`,
    // Boot ISO as CD-ROM
    `--disk=path=${config.isoPath},device=cdrom,readonly=on`,
    `--network=network=${config.network},model=virtio`,
    // UEFI firmware, boot from cdrom (not network)
    "--boot=uefi,cdrom",
    "--os-variant=generic",
    "--noautoconsole",
    "--wait=0",
    "--graphics=vnc,listen=127.0.0.1",
    // Serial console via TCP (port 4556 to avoid conflict with PXE VM)
    "--serial=tcp,host=127.0.0.1:4556,mode=bind,protocol=telnet",
  ];

  if (arch === "aarch64") {
    virtInstallArgs.push("--arch=aarch64", "--machine=virt");
  }

  log(`Running: virt-install --name=${config.name} --boot=uefi,cdrom ...`);
  run(virtInstallArgs.join(" "), { timeout: 60_000 });
  log(`ISO boot VM ${config.name} created (serial: telnet 127.0.0.1 4556)`);
}

/**
 * Execute a command on a VM via its serial console (telnet).
 * Works even when the VM has no network/SSH.
 * Returns the output after the command's echo.
 */
export async function serialExec(
  port: number,
  command: string,
  timeoutMs = 10_000,
): Promise<string> {
  return new Promise((resolve, reject) => {
    const timer = setTimeout(() => {
      sock.destroy();
      reject(new Error(`Serial exec timeout after ${timeoutMs}ms`));
    }, timeoutMs);

    const sock = createConnection({ host: "127.0.0.1", port });
    let buffer = "";
    let sentCommand = false;
    // Random marker to delimit command output
    const marker = `__SERIAL_END_${Date.now()}__`;

    sock.on("connect", () => {
      // Wait for login prompt or shell prompt, then send command
      setTimeout(() => {
        // Send a newline first to get a prompt
        sock.write("\r\n");
      }, 500);
    });

    sock.on("data", (data: Buffer) => {
      buffer += data.toString();

      if (!sentCommand && (buffer.includes("login:") || buffer.includes("# ") || buffer.includes("$ "))) {
        if (buffer.includes("login:")) {
          // Auto-login as root
          sock.write("root\r\n");
          sentCommand = false; // wait for shell prompt after login
          buffer = "";
          return;
        }
        // At shell prompt — send command with marker
        sentCommand = true;
        buffer = "";
        sock.write(`${command}; echo "${marker}"\r\n`);
      }

      if (sentCommand && buffer.includes(marker)) {
        clearTimeout(timer);
        // Extract output between command echo and marker
        const markerIdx = buffer.indexOf(marker);
        const output = buffer.substring(0, markerIdx).trim();
        // Remove the command echo (first line)
        const lines = output.split("\n");
        const result = lines.slice(1).join("\n").trim();
        sock.destroy();
        resolve(result);
      }
    });

    sock.on("error", (err) => {
      clearTimeout(timer);
      reject(new Error(`Serial connection failed: ${err.message}`));
    });
  });
}