Update main.go

Author: michal92299

source-code/main.go

@@ -1,19 +1,17 @@
 package main
 
 import (
-	"flag"
 	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
-	"strconv"
 	"strings"
 	"time"
 )
 
 const (
-	mountPoint    = "/mnt/hroot"
-	snapshotPrefix = "@pre-update-"
+	mountPoint     = "/mnt/hroot"
+	snapshotPrefix = "-pre-update-" // teraz używamy tej stałej
 	updateSnapshot = "@update"
 	btrfsDevice    = "/dev/sda1" // TODO: Detect or configure the Btrfs device
 	rootSubvolume  = "@"
@@ -34,11 +32,11 @@ func main() {
 	case "switch":
 		switchCmd()
 	case "rollback":
-		rollbackCmd(flag.Args()[1:])
+		rollbackCmd(os.Args[2:])
 	case "install":
-		installCmd(flag.Args()[1:])
+		installCmd(os.Args[2:])
 	case "remove":
-		removeCmd(flag.Args()[1:])
+		removeCmd(os.Args[2:])
 	case "clean":
 		cleanCmd()
 	case "status":
@@ -54,14 +52,14 @@ func usage() {
 Usage: hroot <command> [args]
 
 Commands:
-  snapshot          Create a read-only snapshot of the current root
-  update            Create and update a new snapshot offline
-  switch            Switch to the updated snapshot
-  rollback <name>   Rollback to a specific snapshot by name (e.g., @pre-update-20251130-2013)
-  install <pkg>     Install a package in the current root (non-atomic, for simplicity)
-  remove <pkg>      Remove a package from the current root (non-atomic)
-  clean             Clean up temporary files and unused snapshots
-  status            List available snapshots`)
+  snapshot Create a read-only snapshot of the current root
+  update   Create and update a new snapshot offline
+  switch   Switch to the updated snapshot (@update → default)
+  rollback <name> Rollback to a specific snapshot (e.g. @pre-update-20251130-2013)
+  install <pkg>... Install package(s) in the current root (non-atomic)
+  remove <pkg>...  Remove package(s) from the current root (non-atomic)
+  clean    Clean apt cache (snapshots must be deleted manually)
+  status   List available snapshots`)
 }
 
 func runCommand(name string, args ...string) error {
@@ -72,12 +70,32 @@ func runCommand(name string, args ...string) error {
 }
 
 func getSnapshotName() string {
-	return rootSubvolume + "pre-update-" + time.Now().Format("20060102-1504")
+	return rootSubvolume + snapshotPrefix + time.Now().Format("20060102-1504")
+}
+
+// Pobiera Subvolume ID dla podanej względnej ścieżki (np. @update, @pre-update-...)
+func getSubvolumeID(subvol string) (string, error) {
+	path := "/" + subvol // /@update, /@pre-update-20251130-2013 itd.
+	output, err := exec.Command("btrfs", "subvolume", "show", path).CombinedOutput()
+	if err != nil {
+		return "", fmt.Errorf("btrfs subvolume show %s failed: %v\n%s", path, err, output)
+	}
+
+	for _, line := range strings.Split(string(output), "\n") {
+		line = strings.TrimSpace(line)
+		if strings.HasPrefix(line, "Subvolume ID:") {
+			parts := strings.Fields(line)
+			if len(parts) >= 3 {
+				return parts[2], nil
+			}
+		}
+	}
+	return "", fmt.Errorf("Subvolume ID not found for %s", path)
 }
 
 func snapshotCmd() {
 	snapshotName := getSnapshotName()
-	fmt.Printf("Creating snapshot: %s\n", snapshotName)
+	fmt.Printf("Creating read-only snapshot: %s\n", snapshotName)
 	if err := runCommand("btrfs", "subvolume", "snapshot", "-r", "/", snapshotName); err != nil {
 		fmt.Fprintf(os.Stderr, "Error creating snapshot: %v\n", err)
 		os.Exit(1)
@@ -86,22 +104,22 @@ func snapshotCmd() {
 }
 
 func updateCmd() {
-	// Step 1: Create writable snapshot for update
+	// Krok 1: Tworzymy writable snapshot do aktualizacji
 	fmt.Println("Creating update snapshot:", updateSnapshot)
 	if err := runCommand("btrfs", "subvolume", "snapshot", "/", updateSnapshot); err != nil {
 		fmt.Fprintf(os.Stderr, "Error creating update snapshot: %v\n", err)
 		os.Exit(1)
 	}
 
-	// Step 2: Mount the snapshot
+	// Krok 2: Montujemy snapshot
 	os.MkdirAll(mountPoint, 0755)
 	if err := runCommand("mount", btrfsDevice, mountPoint, "-o", "subvol="+updateSnapshot); err != nil {
 		fmt.Fprintf(os.Stderr, "Error mounting update snapshot: %v\n", err)
 		os.Exit(1)
 	}
-	defer runCommand("umount", mountPoint) // Clean up on exit
+	defer runCommand("umount", mountPoint)
 
-	// Bind mount necessary filesystems
+	// Bind mount niezbędnych systemów plików
 	bindMounts := []string{"/proc", "/sys", "/dev", "/run"}
 	for _, m := range bindMounts {
 		target := filepath.Join(mountPoint, m[1:])
@@ -113,115 +131,113 @@ func updateCmd() {
 		defer runCommand("umount", target)
 	}
 
-	// Step 3: Chroot and perform update
+	// Krok 3: Chroot + aktualizacja
 	fmt.Println("Performing system update in chroot...")
-	chrootCmd := []string{"chroot", mountPoint, "apt", "update"}
-	if err := runCommand(chrootCmd[0], chrootCmd[1:]...); err != nil {
+	if err := runCommand("chroot", mountPoint, "apt", "update"); err != nil {
 		fmt.Fprintf(os.Stderr, "Error running apt update: %v\n", err)
 		os.Exit(1)
 	}
-	chrootCmd = []string{"chroot", mountPoint, "apt", "upgrade", "-y"}
-	if err := runCommand(chrootCmd[0], chrootCmd[1:]...); err != nil {
+	if err := runCommand("chroot", mountPoint, "apt", "upgrade", "-y"); err != nil {
 		fmt.Fprintf(os.Stderr, "Error running apt upgrade: %v\n", err)
 		os.Exit(1)
 	}
 
-	fmt.Println("Update completed in snapshot.")
+	fmt.Println("Update completed successfully in snapshot:", updateSnapshot)
+	fmt.Println("Run 'hroot switch' and reboot to apply.")
 }
 
 func switchCmd() {
-	// Get ID of update snapshot
-	out, err := exec.Command("btrfs", "subvolume", "find-new", "/", "9999999").CombinedOutput() // Hack to get current ID, but better to list
+	id, err := getSubvolumeID(updateSnapshot)
 	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error finding subvolume ID: %v\n", err)
+		fmt.Fprintf(os.Stderr, "Cannot get subvolume ID for %s: %v\n", updateSnapshot, err)
 		os.Exit(1)
 	}
-	// Parse ID from output - this is simplistic, assume we know path
-	fmt.Println("Switching default subvolume to", updateSnapshot)
-	if err := runCommand("btrfs", "subvolume", "set-default", updateSnapshot); err != nil {
+
+	fmt.Printf("Setting default subvolume to %s (ID: %s)\n", updateSnapshot, id)
+	if err := runCommand("btrfs", "subvolume", "set-default", id, "/"); err != nil {
 		fmt.Fprintf(os.Stderr, "Error setting default subvolume: %v\n", err)
 		os.Exit(1)
 	}
-	fmt.Println("Default subvolume switched. Reboot to apply.")
+	fmt.Println("Default subvolume changed. Reboot required.")
 }
 
 func rollbackCmd(args []string) {
 	if len(args) == 0 {
-		fmt.Fprintf(os.Stderr, "Rollback requires snapshot name\n")
+		fmt.Fprintf(os.Stderr, "Usage: hroot rollback <snapshot-name>\n")
 		os.Exit(1)
 	}
+
 	snapshotName := args[0]
-	fmt.Printf("Rolling back to %s\n", snapshotName)
-	if err := runCommand("btrfs", "subvolume", "set-default", snapshotName); err != nil {
+	id, err := getSubvolumeID(snapshotName)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Cannot get subvolume ID for %s: %v\n", snapshotName, err)
+		os.Exit(1)
+	}
+
+	fmt.Printf("Rolling back to %s (ID: %s)\n", snapshotName, id)
+	if err := runCommand("btrfs", "subvolume", "set-default", id, "/"); err != nil {
 		fmt.Fprintf(os.Stderr, "Error setting default subvolume: %v\n", err)
 		os.Exit(1)
 	}
-	fmt.Println("Rollback set. Reboot to apply.")
+	fmt.Println("Rollback successful. Reboot required.")
 }
 
-func installCmd(packages []string) {
-	if len(packages) == 0 {
-		fmt.Fprintf(os.Stderr, "Install requires package names\n")
+func installCmd(pkgs []string) {
+	if len(pkgs) == 0 {
+		fmt.Fprintf(os.Stderr, "Usage: hroot install <package>...\n")
 		os.Exit(1)
 	}
-	fmt.Printf("Installing packages: %v\n", packages)
-	args := append([]string{"install", "-y"}, packages...)
+	fmt.Printf("Installing packages (live system): %v\n", pkgs)
+	args := append([]string{"install", "-y"}, pkgs...)
 	if err := runCommand("apt", args...); err != nil {
 		fmt.Fprintf(os.Stderr, "Error installing packages: %v\n", err)
 		os.Exit(1)
 	}
 }
 
-func removeCmd(packages []string) {
-	if len(packages) == 0 {
-		fmt.Fprintf(os.Stderr, "Remove requires package names\n")
+func removeCmd(pkgs []string) {
+	if len(pkgs) == 0 {
+		fmt.Fprintf(os.Stderr, "Usage: hroot remove <package>...\n")
 		os.Exit(1)
 	}
-	fmt.Printf("Removing packages: %v\n", packages)
-	args := append([]string{"remove", "-y"}, packages...)
+	fmt.Printf("Removing packages (live system): %v\n", pkgs)
+	args := append([]string{"remove", "-y"}, pkgs...)
 	if err := runCommand("apt", args...); err != nil {
 		fmt.Fprintf(os.Stderr, "Error removing packages: %v\n", err)
 		os.Exit(1)
 	}
 }
 
 func cleanCmd() {
-	fmt.Println("Cleaning temporary files...")
+	fmt.Println("Cleaning apt cache...")
 	if err := runCommand("apt", "clean"); err != nil {
 		fmt.Fprintf(os.Stderr, "Error cleaning apt cache: %v\n", err)
 	}
-	// Optionally delete old snapshots - manual for now
-	fmt.Println("List snapshots with 'hroot status' and delete manually if needed.")
+	fmt.Println("Done. Delete old snapshots manually with 'btrfs subvolume delete /<name>'")
 }
 
 func statusCmd() {
-	out, err := exec.Command("btrfs", "subvolume", "list", "-p", "/").CombinedOutput()
+	output, err := exec.Command("btrfs", "subvolume", "list", "-p", "/").CombinedOutput()
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Error listing subvolumes: %v\n", err)
 		os.Exit(1)
 	}
-	lines := strings.Split(string(out), "\n")
+
 	fmt.Println("Available snapshots:")
+	lines := strings.Split(string(output), "\n")
 	for _, line := range lines {
-		if strings.Contains(line, rootSubvolume) {
+		if strings.Contains(line, rootSubvolume) || strings.Contains(line, updateSnapshot) {
 			fields := strings.Fields(line)
-			if len(fields) > 8 {
+			if len(fields) >= 9 {
 				id := fields[1]
-				path := fields[8]
-				if strings.HasPrefix(path, rootSubvolume) {
-					fmt.Printf("ID: %s, Path: %s\n", id, path)
-				}
+				path := fields[len(fields)-1]
+				fmt.Printf("  ID %-6s → %s\n", id, path)
 			}
 		}
 	}
 
-	// Get default
-	defOut, defErr := exec.Command("btrfs", "subvolume", "get-default", "/").CombinedOutput()
-	if defErr == nil {
-		defFields := strings.Fields(string(defOut))
-		if len(defFields) > 1 {
-			defID := defFields[1]
-			fmt.Printf("Current default ID: %s\n", defID)
-		}
+	defOutput, err := exec.Command("btrfs", "subvolume", "get-default", "/").CombinedOutput()
+	if err == nil {
+		fmt.Printf("\nCurrent default: %s", defOutput)
 	}
 }