vmix.nix/lib/images/windows/helpers/customizeImage.nix

# Customize Windows images via offline registry merge and/or Audit Mode script execution.
#
# Templates can provide:
#   windowsRegistry — merged offline via virt-win-reg (fast, no boot needed)
#   auditScript     — injected into image and run in Audit Mode via QEMU boot
#   cdroms          — ISO files to attach as CDs when booting for auditScript
#
# Both can be combined: registry is applied first (offline), then the script runs (online).
{ pkgs, lib, ... }:
  originalImage: {
    name ? "",
    diskSize ? "",
    impure ? true,
    # Offline: merge .reg file into registry via virt-win-reg
    windowsRegistry ? "",
    # Online: boot into Audit Mode and run this script
    auditScript ? "",
    # CD-ROMs to attach when booting for auditScript (e.g. VirtIO ISO)
    cdroms ? [],
    # Files to upload into the image before running auditScript
    # List of { source = <drv or path>; dest = "/Windows/path"; }
    uploads ? [],
    # QEMU settings for auditScript boot
    vncDisplay ? null,
    smp ? 4,
    memSize ? 4096,
    nicModel ? null,
    # Flatten COW chain into a standalone qcow2 (removes backing file dependency)
    compact ? false,
    # QEMU timeout in seconds (default 30 min, increase for Windows Update)
    qemuTimeout ? 1800,
    # Blank disk attached for the Audit Mode boot, e.g. { size = "100G"; }.
    # Windows sees it as disk 1, which is what lets a template partition it and
    # relocate profiles onto it in that same boot instead of deferring OOBE to
    # real hardware. Emitted as the derivation's `data` output, so whatever the
    # template writes to it survives the build.
    extraDisk ? null,
  }:
  let
    originalImageName = lib.strings.removeSuffix "-vmix" (lib.strings.removeSuffix ".qcow2" originalImage.name);
    customImageName = (if name != "" then name else "custom") + "-${originalImageName}-vmix.qcow2";
    resultImg = "./disk.qcow2";

    isAHCI = originalImage.useAHCI or false;
    hasRegistry = windowsRegistry != "";
    hasAuditScript = auditScript != "";

    # Offline registry merge
    windowsRegFile = pkgs.writeText "${name}-registry.reg" windowsRegistry;
    virtWinRegMerge = lib.optionalString hasRegistry ''

      echo "=== vmix: merging registry entries (${name}) ==="
      virt-win-reg --merge ${resultImg} ${windowsRegFile}
    '';

    # Audit Mode script injection + QEMU boot
    auditScriptFile = pkgs.writeText "${name}-audit.cmd" auditScript;
    wrapperScript = pkgs.writeText "${name}-audit-wrapper.cmd" ''
      @echo off
      echo === vmix audit: ${name} ===
      call C:\vmix-audit-script.cmd
      echo === vmix audit: ${name} complete ===
      del /q C:\vmix-audit-script.cmd 2>nul
      shutdown /s /t 5 /c "vmix: ${name} complete" 2>nul
      del /q C:\vmix-audit-wrapper.cmd 2>nul
    '';
    runOnceReg = pkgs.writeText "${name}-runonce.reg" (lib.concatStringsSep "\n" [
      "Windows Registry Editor Version 5.00"
      ""
      ''[HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce]''
      ''"vmixAudit"="C:\\vmix-audit-wrapper.cmd"''
      ""
    ]);

    cdromArgs = lib.concatMapStringsSep " \\\n      " (cd: "-drive file=${cd},media=cdrom,readonly=on") cdroms;
    extraDiskImg = "./extra.qcow2";
    extraDiskArgs = lib.optionalString (extraDisk != null)
      (if isAHCI
       then "-drive file=${extraDiskImg},format=qcow2,if=none,id=disk1 -device ide-hd,drive=disk1"
       else "-drive file=${extraDiskImg},format=qcow2,if=virtio");

    displayArg = if vncDisplay != null then "-vnc ${vncDisplay}" else null;

    auditBootCommands = lib.optionalString hasAuditScript ''

      echo "=== vmix: injecting audit script (${name}) ==="
      virt-customize -a ${resultImg} \
        --upload ${auditScriptFile}:/vmix-audit-script.cmd \
        --upload ${wrapperScript}:/vmix-audit-wrapper.cmd \
        ${lib.concatMapStringsSep " \\\n        " (u: let dir = builtins.dirOf u.dest; in "${lib.optionalString (dir != "/") "--mkdir ${dir}"} --upload ${u.source}:${u.dest}") uploads}

      echo "=== vmix: adding RunOnce entry ==="
      virt-win-reg --merge ${resultImg} ${runOnceReg}

      # Boot into Audit Mode to run the script
      cp ${pkgs.OVMF.fd}/FV/OVMF_VARS.fd vars.fd
      chmod +w vars.fd

      VMIX_DISPLAY="-nographic"
      ${lib.optionalString (vncDisplay != null) ''VMIX_DISPLAY="-vnc ${vncDisplay}"''}
      ${lib.optionalString (vncDisplay == null) ''
        # find, not a glob: under nullglob a non-matching /tmp/.vmix-display-*
        # disappears entirely, leaving `ls -t` to list the build directory and
        # hand back nix's own env-vars dump. Exporting that as DISPLAY bloats
        # the environment until every exec dies with E2BIG.
        VMIX_DF=$(find /tmp -maxdepth 1 -type f -name '.vmix-display-*' -printf '%T@ %p\n' 2>/dev/null | sort -rn | head -1 | cut -d' ' -f2-)
        if [ -n "$VMIX_DF" ]; then
          export DISPLAY=$(sed -n '1p' "$VMIX_DF")
          export XAUTHORITY=$(sed -n '2p' "$VMIX_DF")
          export HOME=$(mktemp -d)
          export XDG_RUNTIME_DIR=$HOME
          export SDL_VIDEODRIVER=x11
          VMIX_DISPLAY="-display sdl"
        fi
      ''}

      echo "=== vmix: booting Audit Mode for ${name} ==="
      QEMU_ARGS="-accel kvm -m ${toString memSize} -smp ${toString smp} -cpu host -machine type=q35 \
        -drive if=pflash,format=raw,readonly=on,file=${pkgs.OVMF.fd}/FV/OVMF_CODE.fd \
        -drive if=pflash,format=raw,file=vars.fd \
        -rtc base=localtime,clock=host \
        -device qemu-xhci -device usb-tablet \
        ${if isAHCI
          then "-drive file=${resultImg},format=qcow2,if=none,id=disk0 -device ide-hd,drive=disk0"
          else "-drive file=${resultImg},format=qcow2,if=virtio"} \
        ${cdromArgs} \
        ${extraDiskArgs} \
        -nic user,model=${if nicModel != null then nicModel else if isAHCI then "e1000" else "virtio-net-pci"}"

      timeout ${toString qemuTimeout} qemu-system-x86_64 $VMIX_DISPLAY $QEMU_ARGS || \
        if [[ "$VMIX_DISPLAY" == "-display sdl" ]]; then
          echo "=== vmix: SDL failed, retrying headless ==="
          cp ${pkgs.OVMF.fd}/FV/OVMF_VARS.fd vars.fd && chmod +w vars.fd
          timeout ${toString qemuTimeout} qemu-system-x86_64 -nographic $QEMU_ARGS
        else
          exit 1
        fi

      echo "=== vmix: audit script ${name} complete ==="
    '';

    builderCommand = ''
      # create resulting image backed by original image
      qemu-img create -f qcow2 -b ${originalImage} -F qcow2 ${resultImg}
      [ -n "${diskSize}" ] && qemu-img resize ${resultImg} ${diskSize}
      ${lib.optionalString (extraDisk != null) ''
      echo "=== vmix: creating extra disk (${extraDisk.size}) ==="
      qemu-img create -f qcow2 ${extraDiskImg} ${extraDisk.size}
      ''}
      ${virtWinRegMerge}
      ${auditBootCommands}
      ${lib.optionalString compact ''
      echo "=== vmix: compacting image ==="
      qemu-img convert -O qcow2 ${resultImg} compact.qcow2
      mv compact.qcow2 ${resultImg}
      ''}
      mv ${resultImg} $out
      ${lib.optionalString (extraDisk != null) "mv ${extraDiskImg} $data"}
    '';
  builtImage = pkgs.runCommand customImageName ({
    nativeBuildInputs = with pkgs; [ qemu perl guestfs-tools ];
    requiredSystemFeatures = [ "kvm" ];
  } // lib.optionalAttrs impure { __noChroot = true; }
    # A second output rather than a directory, so ${image} keeps meaning the OS
    # qcow2 for every existing consumer and the fold can still back onto it.
    // lib.optionalAttrs (extraDisk != null) { outputs = [ "out" "data" ]; }) builderCommand;
  in
    builtImage // { _vmixOsType = "windows"; useAHCI = isAHCI; }