mnmivm/main.go

package main

import (
        "fmt"
        "io"
        "net"
        "net/http"
        "os"
        "os/exec"
        "path/filepath"
        "strings"
)

const (
        rootDir   = "/var/lib/microvm"
        imageDir  = rootDir + "/images"
        vmDirBase = rootDir + "/vms"

        memMB = "1024"
        cpus  = "1"
)

type OSImage struct {
        Name     string
        URL      string
        Filename string
        User     string
}

var osImages = map[string]OSImage{
        "ubuntu": {
                Name:     "ubuntu",
                URL:      "https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img",
                Filename: "ubuntu-24.04.qcow2",
                User:     "ubuntu",
        },
        "debian": {
                Name:     "debian",
                URL:      "https://cloud.debian.org/images/cloud/trixie/20251117-2299/debian-13-generic-amd64-20251117-2299.qcow2",
                Filename: "debian-13.qcow2",
                User:     "debian",
        },
        "fedora": {
                Name:     "fedora",
                URL:      "https://download.fedoraproject.org/pub/fedora/linux/releases/43/Cloud/x86_64/images/Fedora-Cloud-Base-Generic-43-1.6.x86_64.qcow2",
                Filename: "fedora-43.qcow2",
                User:     "fedora",
        },
        "alpine": {
                Name:     "alpine",
                URL:      "https://dl-cdn.alpinelinux.org/alpine/v3.22/releases/cloud/generic_alpine-3.22.2-x86_64-bios-cloudinit-r0.qcow2",
                Filename: "alpine-3.22.qcow2",
                User:     "alpine",
        },
}

func ensureDirs() {
        for _, d := range []string{imageDir, vmDirBase} {
                if err := os.MkdirAll(d, 0755); err != nil {
                        panic(err)
                }
        }
}

func usage() {
        fmt.Println("usage:")
        fmt.Println("  mnmivm create <name> --os <ubuntu|debian|fedora|alpine> --pubkey-path <path>")
        fmt.Println("  mnmivm start <name>")
        fmt.Println("  mnmivm stop <name>")
        fmt.Println("  mnmivm delete <name>")
        fmt.Println("  mnmivm list")
        os.Exit(1)
}

func vmDir(name string) string        { return filepath.Join(vmDirBase, name) }
func pidFile(name string) string      { return filepath.Join(vmDir(name), "vm.pid") }
func diskPath(name string) string     { return filepath.Join(vmDir(name), "disk.qcow2") }
func seedISOPath(name string) string  { return filepath.Join(vmDir(name), "seed.iso") }
func vmPubKeyPath(name string) string { return filepath.Join(vmDir(name), "pubkey.pub") }

func imagePath(osName string) string {
        return filepath.Join(imageDir, osImages[osName].Filename)
}

func parseArg(args []string, key string) string {
        for i := 0; i < len(args)-1; i++ {
                if args[i] == key {
                        return args[i+1]
                }
        }
        panic("missing required argument: " + key)
}

func findFreePort() int {
        l, err := net.Listen("tcp", "127.0.0.1:0")
        if err != nil {
                panic(err)
        }
        defer l.Close()
        return l.Addr().(*net.TCPAddr).Port
}

func ensureBaseImage(osName string) {
        img := osImages[osName]
        path := imagePath(osName)

        if _, err := os.Stat(path); err == nil {
                return
        }

        fmt.Println("Downloading", img.Name, "cloud image...")
        resp, err := http.Get(img.URL)
        if err != nil {
                panic(err)
        }
        defer resp.Body.Close()

        out, err := os.Create(path)
        if err != nil {
                panic(err)
        }
        defer out.Close()

        if _, err := io.Copy(out, resp.Body); err != nil {
                panic(err)
        }
}

func readPublicKey(path string) string {
        data, err := os.ReadFile(path)
        if err != nil {
                panic(err)
        }
        key := strings.TrimSpace(string(data))
        if !strings.HasPrefix(key, "ssh-") {
                panic("invalid SSH public key")
        }
        return key
}

func createCloudInitSeed(name, osName, pubKeyPath string) {
        img := osImages[osName]
        sshKey := readPublicKey(pubKeyPath)

        userData := fmt.Sprintf(`#cloud-config
users:
  - name: %s
    sudo: ALL=(ALL) NOPASSWD:ALL
    shell: /bin/sh
    ssh_authorized_keys:
      - %s
ssh_pwauth: false
disable_root: true
`, img.User, sshKey)

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

        tmpDir, _ := os.MkdirTemp("", "cloudinit")
        defer os.RemoveAll(tmpDir)

        os.WriteFile(filepath.Join(tmpDir, "user-data"), []byte(userData), 0644)
        os.WriteFile(filepath.Join(tmpDir, "meta-data"), []byte(metaData), 0644)

        cmd := exec.Command(
                "genisoimage",
                "-output", seedISOPath(name),
                "-volid", "cidata",
                "-joliet",
                "-rock",
                filepath.Join(tmpDir, "user-data"),
                filepath.Join(tmpDir, "meta-data"),
        )

        if out, err := cmd.CombinedOutput(); err != nil {
                panic(string(out))
        }
}

func createVM(name, osName, pubKeyPath string) {
        if _, ok := osImages[osName]; !ok {
                panic("unsupported OS: " + osName)
        }

        ensureBaseImage(osName)
        os.MkdirAll(vmDir(name), 0755)

        keyData, _ := os.ReadFile(pubKeyPath)
        os.WriteFile(vmPubKeyPath(name), keyData, 0644)
        os.WriteFile(filepath.Join(vmDir(name), "os.name"), []byte(osName), 0644)

        cmd := exec.Command(
                "qemu-img", "create",
                "-f", "qcow2",
                "-F", "qcow2",
                "-b", imagePath(osName),
                diskPath(name),
        )

        if out, err := cmd.CombinedOutput(); err != nil {
                panic(string(out))
        }

        createCloudInitSeed(name, osName, vmPubKeyPath(name))
        fmt.Println("VM created:", name, "OS:", osName)
}

func startVM(name string) {
        if data, err := os.ReadFile(pidFile(name)); err == nil {
                if _, err := os.Stat("/proc/" + strings.TrimSpace(string(data))); err == nil {
                        panic("VM already running")
                }
                os.Remove(pidFile(name))
        }

        sshPort := findFreePort()
        vncPort := findFreePort()
        vncDisplay := vncPort - 5900
        if vncDisplay < 0 {
                panic("invalid VNC port")
        }

        cmd := exec.Command(
                "qemu-system-x86_64",
                "-enable-kvm",
                "-machine", "q35",
                "-cpu", "host",
                "-m", memMB,
                "-smp", cpus,
                "-vga", "virtio",
                "-display", fmt.Sprintf("vnc=0.0.0.0:%d", vncDisplay),
                "-drive", fmt.Sprintf("file=%s,if=virtio,format=qcow2", diskPath(name)),
                "-drive", fmt.Sprintf("file=%s,if=virtio,format=raw", seedISOPath(name)),
                "-netdev", fmt.Sprintf("user,id=net0,hostfwd=tcp::%d-:22", sshPort),
                "-device", "virtio-net-pci,netdev=net0",
                "-pidfile", pidFile(name),
                "-daemonize",
        )

        if out, err := cmd.CombinedOutput(); err != nil {
                panic(string(out))
        }

        os.WriteFile(filepath.Join(vmDir(name), "ssh.port"), []byte(fmt.Sprint(sshPort)), 0644)
        os.WriteFile(filepath.Join(vmDir(name), "vnc.port"), []byte(fmt.Sprint(vncPort)), 0644)

        fmt.Println("VM started:", name)
        fmt.Printf("SSH: ssh %s@localhost -p %d\n", osImages[readFileTrim(filepath.Join(vmDir(name), "os.name"))].User, sshPort)
        fmt.Printf("VNC: vnc://0.0.0.0:%d\n", vncPort)
}

func stopVM(name string) {
        data, _ := os.ReadFile(pidFile(name))
        exec.Command("kill", strings.TrimSpace(string(data))).Run()
        os.Remove(pidFile(name))
        fmt.Println("VM stopped:", name)
}

func deleteVM(name string) {
        dir := vmDir(name)
        if _, err := os.Stat(dir); err != nil {
                panic("VM does not exist")
        }

        if data, err := os.ReadFile(pidFile(name)); err == nil {
                exec.Command("kill", strings.TrimSpace(string(data))).Run()
                os.Remove(pidFile(name))
        }

        fmt.Printf("\nWARNING: Permanently delete VM \"%s\"?\nType YES to confirm: ", name)
        var confirm string
        fmt.Scanln(&confirm)
        if confirm != "YES" {
                fmt.Println("Aborted.")
                return
        }

        os.RemoveAll(dir)
        fmt.Println("VM deleted:", name)
}

func readFileTrim(path string) string {
        data, err := os.ReadFile(path)
        if err != nil {
                return "-"
        }
        return strings.TrimSpace(string(data))
}

func vmRunning(name string) bool {
        data, err := os.ReadFile(pidFile(name))
        if err != nil {
                return false
        }
        _, err = os.Stat("/proc/" + strings.TrimSpace(string(data)))
        return err == nil
}

func listVMs() {
        entries, _ := os.ReadDir(vmDirBase)

        fmt.Printf("%-12s %-8s %-8s %-16s %-16s %s\n",
                "NAME", "STATE", "OS", "SSH", "VNC", "PUBKEY")
        fmt.Println(strings.Repeat("-", 90))

        for _, e := range entries {
                name := e.Name()
                dir := vmDir(name)

                state := "stopped"
                if vmRunning(name) {
                        state = "running"
                }

                osName := readFileTrim(filepath.Join(dir, "os.name"))
                sshPort := readFileTrim(filepath.Join(dir, "ssh.port"))
                vncPort := readFileTrim(filepath.Join(dir, "vnc.port"))

                ssh := "-"
                if sshPort != "-" {
                        ssh = "localhost:" + sshPort
                }

                vnc := "-"
                if vncPort != "-" {
                        vnc = "0.0.0.0:" + vncPort
                }

                fmt.Printf("%-12s %-8s %-8s %-16s %-16s %s\n",
                        name, state, osName, ssh, vnc, vmPubKeyPath(name))
        }
}

func main() {
        if len(os.Args) < 2 {
                usage()
        }

        ensureDirs()

        switch os.Args[1] {
        case "create":
                createVM(os.Args[2], parseArg(os.Args, "--os"), parseArg(os.Args, "--pubkey-path"))
        case "start":
                startVM(os.Args[2])
        case "stop":
                stopVM(os.Args[2])
        case "delete":
                deleteVM(os.Args[2])
        case "list":
                listVMs()
        default:
                usage()
        }
}