From 034fe688cd7fa9b097a1c0be393db3ee8d3424ba Mon Sep 17 00:00:00 2001
From: markmental <marky611@gmail.com>
Date: Mon, 15 Dec 2025 16:28:30 -0500
Subject: [PATCH] Adds cloud-init SSH key updating

---
 main.go | 88 ++++++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 78 insertions(+), 10 deletions(-)

diff --git a/main.go b/main.go
index 8d9dec8..486930e 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
         "fmt"
         "io"
         "net"
+        "time"
         "net/http"
         "os"
         "os/exec"
@@ -74,6 +75,7 @@ func 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 update-cloud <name> --pubkey-path <path>")
         fmt.Println("  mnmivm delete <name>")
         fmt.Println("  mnmivm list")
         os.Exit(1)
@@ -145,28 +147,87 @@ func readPublicKey(path string) string {
         return key
 }
 
+func updateCloudVM(name, pubKeyPath string) {
+        dir := vmDir(name)
+
+        if _, err := os.Stat(dir); err != nil {
+                panic("VM does not exist: " + name)
+        }
+
+        if vmRunning(name) {
+                panic("VM is currently running. Stop it before updating cloud-init.")
+        }
+
+        osName := readFileTrim(filepath.Join(dir, "os.name"))
+        if osName == "-" {
+                panic("cannot determine OS for VM")
+        }
+
+        // overwrite stored pubkey
+        keyData, err := os.ReadFile(pubKeyPath)
+        if err != nil {
+                panic(err)
+        }
+
+        if err := os.WriteFile(vmPubKeyPath(name), keyData, 0644); err != nil {
+                panic(err)
+        }
+
+        // regenerate seed ISO
+        createCloudInitSeed(name, osName, vmPubKeyPath(name))
+
+        fmt.Println("Cloud-init updated for VM:", name)
+        fmt.Println("Changes will apply on next boot.")
+}
+
 func createCloudInitSeed(name, osName, pubKeyPath string) {
         img := osImages[osName]
         sshKey := readPublicKey(pubKeyPath)
 
+        instanceID := fmt.Sprintf("%s-%d", name, time.Now().Unix())
+
         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)
+write_files:
+  - path: /home/%s/.ssh/authorized_keys
+    owner: %s:%s
+    permissions: '0600'
+    content: |
+      %s
 
-        tmpDir, _ := os.MkdirTemp("", "cloudinit")
+runcmd:
+  - chown -R %s:%s /home/%s/.ssh
+`, img.User, img.User, img.User, img.User, sshKey,
+   img.User, img.User, img.User)
+
+        metaData := fmt.Sprintf(
+                "instance-id: %s\nlocal-hostname: %s\n",
+                instanceID,
+                name,
+        )
+
+        tmpDir, err := os.MkdirTemp("", "cloudinit")
+        if err != nil {
+                panic(err)
+        }
         defer os.RemoveAll(tmpDir)
 
-        os.WriteFile(filepath.Join(tmpDir, "user-data"), []byte(userData), 0644)
-        os.WriteFile(filepath.Join(tmpDir, "meta-data"), []byte(metaData), 0644)
+        userDataPath := filepath.Join(tmpDir, "user-data")
+        metaDataPath := filepath.Join(tmpDir, "meta-data")
+
+        if err := os.WriteFile(userDataPath, []byte(userData), 0644); err != nil {
+                panic(err)
+        }
+        if err := os.WriteFile(metaDataPath, []byte(metaData), 0644); err != nil {
+                panic(err)
+        }
 
         cmd := exec.Command(
                 "genisoimage",
@@ -174,15 +235,16 @@ disable_root: true
                 "-volid", "cidata",
                 "-joliet",
                 "-rock",
-                filepath.Join(tmpDir, "user-data"),
-                filepath.Join(tmpDir, "meta-data"),
+                userDataPath,
+                metaDataPath,
         )
 
         if out, err := cmd.CombinedOutput(); err != nil {
-                panic(string(out))
+                panic("cloud-init ISO creation failed:\n" + string(out))
         }
 }
 
+
 func renderBoxTable(headers []string, rows [][]string) {
         colWidths := make([]int, len(headers))
 
@@ -436,6 +498,12 @@ func main() {
                 deleteVM(os.Args[2])
         case "list":
                 listVMs()
+	case "update-cloud":
+		if len(os.Args) < 3 {
+			usage()
+		}
+		updateCloudVM(os.Args[2], parseArg(os.Args, "--pubkey-path"))
+
         default:
                 usage()
         }