mkvm/main.go

package main

import (
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/http"
	"strings"
	"sync"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"libvirt.org/go/libvirtxml"

	"mkvm/config"
	"mkvm/libvirtx"
)

var (
	wg        sync.WaitGroup
	nicSource = libvirtxml.DomainInterfaceSource{}
	rootCmd   = cobra.Command{
		Use:   "mkvm name [name [name]]",
		Short: "create virtual machine(s) via libvirt",
		Args:  cobra.MinimumNArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			logrus.SetLevel(logrus.DebugLevel)

			if err := config.Load(); err != nil {
				logrus.WithError(err).Fatal("error loading config")
			}

			for _, u := range argSSHKeyURLs {
				if err := downloadSSHKeys(u); err != nil {
					logrus.WithError(err).WithField("url", u).Fatal("error downloading SSH keys")
				}
			}

			conn, err := libvirtx.New()
			if err != nil {
				logrus.WithError(err).Fatal("error connecting to libvirt")
				return
			}
			defer conn.Close()

			serverInterfaceName := ""
			if config.C.Network != "" {
				nicSource.Network = &libvirtxml.DomainInterfaceSourceNetwork{Network: config.C.Network}
				libvirtnet, err := conn.LookupNetworkByName(config.C.Network)
				if err != nil {
					logrus.WithError(err).WithField("network", config.C.Network).Fatal("error finding libvirt network")
				}

				xmlstr, err := libvirtnet.GetXMLDesc(0)
				if err != nil {
					logrus.WithError(err).WithField("network", config.C.Network).Fatal("error getting network xml description")
				}

				var net libvirtxml.Network
				if err := net.Unmarshal(xmlstr); err != nil {
					logrus.WithError(err).WithField("network", config.C.Network).Fatal("error parsing network xml description")
				}

				serverInterfaceName = net.Bridge.Name
			} else if config.C.Bridge != "" {
				nicSource.Bridge = &libvirtxml.DomainInterfaceSourceBridge{Bridge: config.C.Bridge}
				serverInterfaceName = config.C.Bridge
			} else {
				logrus.Fatal("no network or bridge configured")
			}

			serverInterface, err := net.InterfaceByName(serverInterfaceName)
			if err != nil {
				logrus.WithError(err).Fatal("error finding local network interface to run server on")
			}

			serverInterfaceAddrs, err := serverInterface.Addrs()
			if err != nil {
				logrus.WithError(err).Fatal("error finding local network interface's IP")
			}

			if len(serverInterfaceAddrs) == 0 {
				logrus.WithField("interface", serverInterfaceName).Fatal("bridge interface does not have an IP on this machine")
			}

			serverBindIP, _, err := net.ParseCIDR(serverInterfaceAddrs[0].String())
			if err != nil {
				logrus.WithField("interface", serverInterfaceName).WithField("address", serverInterfaceAddrs[0].String()).WithError(err).Fatal("error parsing local address")
			}

			port := rand.Intn(65535-1025) + 1025
			serverBind := fmt.Sprintf("%s:%d", serverBindIP, port)

			go runHTTPServer(serverBind)

			for i, name := range args {
				metadataURL := fmt.Sprintf("http://%s/%d", serverBind, i)
				if err := buildCloudConfig(i, name, metadataURL); err != nil {
					logrus.WithError(err).WithField("vm", name).Error("error building cloudconfig for vm")
					continue
				}

				wg.Add(1)
				if err := mkvm(conn, metadataURL, name); err != nil {
					logrus.WithError(err).Error("unexpected error building VM")
					wg.Done()
				}
			}

			logrus.Info("waiting for VM(s) to finish provisioning")
			wg.Wait()
		},
	}

	argMemoryMB   int
	argCPUs       int
	argImage      string
	argDiskSizeGB int

	// cloudinit args
	argSSHKeys    []string
	argSSHKeyURLs []string
	argPackages   []string
)

func downloadSSHKeys(url string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		logrus.WithFields(logrus.Fields{
			"url":    url,
			"status": resp.Status,
			"body":   string(body),
		}).Error("non-200 response from SSH key URL")
		return fmt.Errorf("non-200 response from SSH key URL: %s", resp.Status)
	}

	count := 0
	for _, key := range strings.Split(string(body), "\n") {
		key = strings.TrimSpace(key)
		if key == "" {
			continue
		}
		argSSHKeys = append(argSSHKeys, key)
		count++
	}

	logrus.WithField("url", url).WithField("keys", count).Debug("downloaded SSH authorized keys")

	return nil
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		panic(err)
	}
}

func init() {
	rootCmd.Flags().IntVarP(&argMemoryMB, "memory", "m", 1024, "amount of memory (in MB) to assign to the VM")
	rootCmd.Flags().IntVarP(&argCPUs, "cpu", "c", 2, "the number of vCPU cores to assign to the VM")
	rootCmd.Flags().StringVarP(&argImage, "image", "", "https://cloud.debian.org/images/cloud/bookworm/latest/debian-12-generic-amd64.qcow2", "URL of the image to download")
	rootCmd.Flags().IntVarP(&argDiskSizeGB, "disk", "d", 25, "disk size (in GB)")

	rootCmd.Flags().StringArrayVar(&argSSHKeys, "ssh-keys", nil, "SSH key(s) authorzed to access the VM")
	rootCmd.Flags().StringArrayVarP(&argSSHKeyURLs, "ssh-key-urls", "s", nil, "URL(s) to SSH key(s) authorzed to access the VM. Expected in authorized_keys format.")
	rootCmd.Flags().StringArrayVarP(&argPackages, "packages", "p", nil, "packages to install on the VM")
}