Provisioning/internal/image/isomod.go

package image

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/kdomanski/iso9660"
)

func ModifyISOForAutoInstall(isoPath, answerURL string) error {
	origConfig, err := readGrubCfgFromISO(isoPath)
	if err != nil {
		return err
	}
	if origConfig == nil {
		return fmt.Errorf("no grub.cfg found in ISO filesystem")
	}

	log.Printf("image: found grub.cfg (%d bytes), first 200 chars: %s",
		len(origConfig), truncate(string(origConfig), 200))

	newConfig := rewriteGrubConfig(string(origConfig), answerURL)

	if len(newConfig) > len(origConfig) {
		return fmt.Errorf("modified GRUB config (%d bytes) exceeds original (%d bytes)", len(newConfig), len(origConfig))
	}

	padded := make([]byte, len(origConfig))
	copy(padded, []byte(newConfig))
	for i := len(newConfig); i < len(padded); i++ {
		padded[i] = '\n'
	}

	offsets, err := findAllOccurrences(isoPath, origConfig)
	if err != nil {
		return fmt.Errorf("locate grub.cfg in ISO: %w", err)
	}

	log.Printf("image: found grub.cfg at %d location(s) in ISO: %v", len(offsets), offsets)

	f, err := os.OpenFile(isoPath, os.O_WRONLY, 0)
	if err != nil {
		return fmt.Errorf("open ISO for writing: %w", err)
	}
	defer f.Close()

	for _, offset := range offsets {
		if _, err := f.WriteAt(padded, offset); err != nil {
			return fmt.Errorf("write at offset %d: %w", offset, err)
		}
	}

	log.Printf("image: modified grub.cfg at %d location(s) for auto-install", len(offsets))
	log.Printf("image: new grub.cfg:\n%s", newConfig)
	return nil
}

func readGrubCfgFromISO(isoPath string) ([]byte, error) {
	f, err := os.Open(isoPath)
	if err != nil {
		return nil, fmt.Errorf("open ISO: %w", err)
	}
	defer f.Close()

	img, err := iso9660.OpenImage(f)
	if err != nil {
		return nil, fmt.Errorf("parse ISO: %w", err)
	}

	root, err := img.RootDir()
	if err != nil {
		return nil, fmt.Errorf("read ISO root: %w", err)
	}

	grubFile := findFileByName(root, "grub.cfg")
	if grubFile == nil {
		return nil, nil
	}

	reader := grubFile.Reader()
	data, err := readAll(reader)
	if err != nil {
		return nil, fmt.Errorf("read grub.cfg: %w", err)
	}
	return data, nil
}

func findFileByName(dir *iso9660.File, target string) *iso9660.File {
	children, err := dir.GetChildren()
	if err != nil {
		return nil
	}
	for _, child := range children {
		if child.IsDir() {
			if result := findFileByName(child, target); result != nil {
				return result
			}
		} else if strings.EqualFold(child.Name(), target) {
			return child
		}
	}
	return nil
}

func rewriteGrubConfig(original, answerURL string) string {
	lines := strings.Split(original, "\n")
	var result []string
	depth := 0
	firstMenuModified := false

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)

		// Strip comments to reclaim space for added kernel parameters
		if strings.HasPrefix(trimmed, "#") {
			continue
		}

		if strings.HasPrefix(trimmed, "set timeout=") {
			result = append(result, "set timeout=0")
			continue
		}
		if strings.HasPrefix(trimmed, "set default=") {
			result = append(result, "set default=0")
			continue
		}

		if strings.HasPrefix(trimmed, "menuentry") {
			depth++
		}
		if trimmed == "}" && depth > 0 {
			if depth == 1 {
				firstMenuModified = true
			}
			depth--
		}

		if depth > 0 && !firstMenuModified &&
			(strings.HasPrefix(trimmed, "linux ") || strings.HasPrefix(trimmed, "linux\t")) {
			if !strings.Contains(line, "proxmox-start-auto-installer") {
				line = strings.TrimRight(line, " \t") + " proxmox-start-auto-installer"
			}
			if answerURL != "" && !strings.Contains(line, "proxmox-auto-installer-answer-url") {
				line += " proxmox-auto-installer-answer-url=" + answerURL
			}
		}

		line = strings.TrimRight(line, " \t")
		result = append(result, line)
	}

	// Collapse consecutive blank lines
	var collapsed []string
	prevBlank := false
	for _, line := range result {
		if line == "" {
			if !prevBlank {
				collapsed = append(collapsed, line)
			}
			prevBlank = true
		} else {
			collapsed = append(collapsed, line)
			prevBlank = false
		}
	}

	return strings.Join(collapsed, "\n")
}

func findAllOccurrences(isoPath string, content []byte) ([]int64, error) {
	f, err := os.Open(isoPath)
	if err != nil {
		return nil, fmt.Errorf("open ISO: %w", err)
	}
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return nil, err
	}

	needle := content
	if len(needle) > 256 {
		needle = needle[:256]
	}

	const chunkSize = 8 * 1024 * 1024
	overlap := len(needle) - 1
	buf := make([]byte, chunkSize+overlap)

	var offsets []int64
	filePos := int64(0)

	for filePos < stat.Size() {
		readStart := filePos
		if filePos > 0 {
			readStart -= int64(overlap)
		}

		n, err := f.ReadAt(buf, readStart)
		if n == 0 {
			break
		}

		searchFrom := 0
		if filePos > 0 {
			searchFrom = overlap
		}

		chunk := buf[:n]
		pos := searchFrom
		for {
			idx := bytes.Index(chunk[pos:], needle)
			if idx == -1 {
				break
			}
			offsets = append(offsets, readStart+int64(pos)+int64(idx))
			pos += idx + 1
		}

		filePos += chunkSize
		if err != nil {
			break
		}
	}

	if len(offsets) == 0 {
		return nil, fmt.Errorf("grub.cfg content not found in ISO raw data (searched %d bytes with %d-byte needle)",
			stat.Size(), len(needle))
	}

	return offsets, nil
}

func readAll(r interface{ Read([]byte) (int, error) }) ([]byte, error) {
	var buf bytes.Buffer
	_, err := buf.ReadFrom(r)
	return buf.Bytes(), err
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}