main.go

// Copyright 2025 Laurynas Četyrkinas
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"strings"

	"github.com/kdomanski/iso9660"
	"gopkg.in/yaml.v3"
)

// User data structures
type User struct {
	Name              string   `yaml:"name"`
	SSHAuthorizedKeys []string `yaml:"ssh_authorized_keys"`
}

type UserData struct {
	Users          []User   `yaml:"users"`
	Hostname       string   `yaml:"hostname"`
	PackageUpgrade *bool    `yaml:"package_upgrade,omitempty"`
	BootCmd        []string `yaml:"bootcmd,omitempty"`
	RunCmd         []string `yaml:"runcmd,omitempty"`
}

// Network config structures
type Subnet struct {
	Type           string   `yaml:"type"`
	Address        string   `yaml:"address,omitempty"`
	Gateway        string   `yaml:"gateway,omitempty"`
	DNSNameservers []string `yaml:"dns_nameservers,omitempty"`
}

type NetworkInterface struct {
	Type    string   `yaml:"type"`
	Name    string   `yaml:"name"`
	Subnets []Subnet `yaml:"subnets"`
}

type NetworkConfig struct {
	Version int                `yaml:"version"`
	Config  []NetworkInterface `yaml:"config"`
}

func main() {
	sshKey := flag.String("ssh-key", "", "SSH public key (required)")
	ipv4 := flag.String("ipv4", "", "Static IPv4 with CIDR")
	ipv4Gateway := flag.String("ipv4-gateway", "", "IPv4 gateway")
	ipv6 := flag.String("ipv6", "", "Static IPv6 with CIDR")
	ipv6Gateway := flag.String("ipv6-gateway", "fe80::1", "IPv6 gateway")
	dns := flag.String("dns", "95.85.95.85,2.56.220.2,2a03:90c0:999d::1,2a03:90c0:9992::1", "DNS servers (comma-separated)")
	iface := flag.String("interface", "eth0", "Network interface name")
	hostname := flag.String("hostname", "localhost", "Hostname")
	upgrade := flag.Bool("upgrade", true, "Upgrade packages")
	disableSerial := flag.Bool("disable-serial", false, "Disable serial console service")
	output := flag.String("output", "cloud-init.iso", "Output filename")

	flag.Parse()

	if *sshKey == "" {
		fmt.Fprintf(os.Stderr, "Error: -ssh-key is required\n")
		os.Exit(1)
	}

	if err := createISO(*sshKey, *ipv4, *ipv4Gateway, *ipv6, *ipv6Gateway, *dns, *iface, *hostname, *upgrade, *disableSerial, *output); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Created: %s\n", *output)
}

func createISO(sshKey, ipv4, ipv4Gateway, ipv6, ipv6Gateway, dns, iface, hostname string, upgrade, disableSerial bool, output string) error {
	userData, err := generateUserData(sshKey, hostname, upgrade, disableSerial)
	if err != nil {
		return fmt.Errorf("failed to generate user data: %w", err)
	}

	metaData := fmt.Sprintf("instance-id: %s\nlocal-hostname: %s\n", hostname, hostname)

	writer, err := iso9660.NewWriter()
	if err != nil {
		return err
	}
	defer writer.Cleanup()

	if err := writer.AddFile(strings.NewReader(userData), "user-data"); err != nil {
		return err
	}

	if err := writer.AddFile(strings.NewReader(metaData), "meta-data"); err != nil {
		return err
	}

	if ipv4 != "" || ipv6 != "" {
		networkConfig, err := generateNetworkConfig(ipv4, ipv4Gateway, ipv6, ipv6Gateway, dns, iface)
		if err != nil {
			return fmt.Errorf("failed to generate network config: %w", err)
		}
		if err := writer.AddFile(strings.NewReader(networkConfig), "network-config"); err != nil {
			return err
		}
	}

	file, err := os.Create(output)
	if err != nil {
		return err
	}
	defer file.Close()

	return writer.WriteTo(file, "CIDATA")
}

func generateUserData(sshKey, hostname string, upgrade, disableSerial bool) (string, error) {
	config := UserData{
		Users: []User{
			{
				Name:              "root",
				SSHAuthorizedKeys: []string{sshKey},
			},
		},
		Hostname: hostname,
	}

	if upgrade {
		config.PackageUpgrade = &upgrade
	}

	if disableSerial {
		config.BootCmd = []string{
			"systemctl disable --now serial-getty@ttyS0.service || true",
			"systemctl mask serial-getty@ttyS0.service || true",
		}
	}

	yamlData, err := yaml.Marshal(&config)
	if err != nil {
		return "", err
	}

	return "#cloud-config\n" + string(yamlData), nil
}

func generateNetworkConfig(ipv4, ipv4Gateway, ipv6, ipv6Gateway, dns, iface string) (string, error) {
	config := NetworkConfig{
		Version: 1,
		Config: []NetworkInterface{
			{
				Type:    "physical",
				Name:    iface,
				Subnets: []Subnet{},
			},
		},
	}

	ipv4DNSServers, ipv6DNSServers := separateDNSByIPVersion(dns)

	if ipv4 != "" {
		ipv4Subnet := Subnet{
			Type:    "static",
			Address: ipv4,
		}
		if ipv4Gateway != "" {
			ipv4Subnet.Gateway = ipv4Gateway
		}
		if len(ipv4DNSServers) > 0 {
			ipv4Subnet.DNSNameservers = ipv4DNSServers
		}
		config.Config[0].Subnets = append(config.Config[0].Subnets, ipv4Subnet)
	}

	if ipv6 != "" {
		ipv6Subnet := Subnet{
			Type:    "static6",
			Address: ipv6,
		}
		if ipv6Gateway != "" {
			ipv6Subnet.Gateway = ipv6Gateway
		}
		if len(ipv6DNSServers) > 0 {
			ipv6Subnet.DNSNameservers = ipv6DNSServers
		}
		config.Config[0].Subnets = append(config.Config[0].Subnets, ipv6Subnet)
	}

	yamlData, err := yaml.Marshal(&config)
	if err != nil {
		return "", err
	}

	return string(yamlData), nil
}

// separateDNSByIPVersion parses DNS servers and returns separate slices for IPv4 and IPv6
func separateDNSByIPVersion(dns string) (ipv4DNS, ipv6DNS []string) {
	if dns == "" {
		return nil, nil
	}

	for _, server := range strings.Split(dns, ",") {
		server = strings.TrimSpace(server)
		if ip := net.ParseIP(server); ip != nil {
			if ip.To4() != nil {
				ipv4DNS = append(ipv4DNS, server)
			} else {
				ipv6DNS = append(ipv6DNS, server)
			}
		}
	}
	return ipv4DNS, ipv6DNS
}