From 174d69c3eae39e1cae7be11ead2518b7e1704701 Mon Sep 17 00:00:00 2001
From: Itxaka <itxaka.garcia@spectrocloud.com>
Date: Tue, 9 Jan 2024 16:42:49 +0100
Subject: [PATCH] Merge uki upgrade into the normal upgrade command (#197)

---
 internal/agent/upgrade.go | 13 ++++++--
 main.go                   | 64 ---------------------------------------
 2 files changed, 10 insertions(+), 67 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 274caf43..8fc138df 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -13,6 +13,7 @@ import (
 	config "github.com/kairos-io/kairos-agent/v2/pkg/config"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
 	"github.com/kairos-io/kairos-agent/v2/pkg/uki"
+	internalutils "github.com/kairos-io/kairos-agent/v2/pkg/utils"
 	events "github.com/kairos-io/kairos-sdk/bus"
 	"github.com/kairos-io/kairos-sdk/collector"
 	"github.com/kairos-io/kairos-sdk/utils"
@@ -61,6 +62,14 @@ func Upgrade(
 	source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) error {
 	bus.Manager.Initialize()
 
+	if internalutils.UkiBootMode() == internalutils.UkiRemovableMedia {
+		return upgradeUki(source, dirs, strictValidations)
+	} else {
+		return upgrade(source, force, strictValidations, dirs, preReleases, upgradeRecovery)
+	}
+}
+
+func upgrade(source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) error {
 	upgradeSpec, c, err := generateUpgradeSpec(source, force, strictValidations, dirs, preReleases, upgradeRecovery)
 	if err != nil {
 		return err
@@ -187,9 +196,7 @@ func getReleasesFromProvider(includePrereleases bool) ([]string, error) {
 	return result, nil
 }
 
-func UkiUpgrade(source string, dirs []string, strictValidations bool) error {
-	bus.Manager.Initialize()
-
+func upgradeUki(source string, dirs []string, strictValidations bool) error {
 	cliConf, err := generateUpgradeConfForCLIArgs(source, false)
 	if err != nil {
 		return err
diff --git a/main.go b/main.go
index 5a6b38c7..85f016d7 100644
--- a/main.go
+++ b/main.go
@@ -20,7 +20,6 @@ import (
 	"github.com/kairos-io/kairos-agent/v2/internal/webui"
 	agentConfig "github.com/kairos-io/kairos-agent/v2/pkg/config"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
-	"github.com/kairos-io/kairos-agent/v2/pkg/uki"
 	"github.com/kairos-io/kairos-sdk/bundles"
 	"github.com/kairos-io/kairos-sdk/collector"
 	"github.com/kairos-io/kairos-sdk/machine"
@@ -702,69 +701,6 @@ The validate command expects a configuration file as its only argument. Local fi
 			return nil
 		},
 	},
-	{
-		Name:        "uki",
-		Usage:       "UKI subcommands",
-		Description: "UKI subcommands",
-		// we could set the flag --source at this level so we could have the flag for all subcommands but that translates into an ugly command
-		// in which you need to put the source flag before the subcommand, which is a mess. Just bad UX.
-		// command level: kairos-agent uki --source oci:whatever install
-		// subcommand level: kairos-agent uki install --source oci:whatever
-		Subcommands: []*cli.Command{
-			{
-				Name: "upgrade",
-				Flags: []cli.Flag{
-					&cli.StringFlag{
-						Name:  "source",
-						Usage: "Source for upgrade. Composed of `type:address`. Accepts `file:`,`dir:` or `oci:` for the type of source.\nFor example `file:/var/share/myimage.tar`, `dir:/tmp/extracted` or `oci:repo/image:tag`",
-						Action: func(c *cli.Context, s string) error {
-							return validateSource(s)
-						},
-					},
-				},
-				Before: func(c *cli.Context) error {
-					if c.String("source") == "" {
-						return fmt.Errorf("On uki, source is required")
-					}
-					return nil
-				},
-				Action: func(c *cli.Context) error {
-					source := c.String("source")
-					return agent.UkiUpgrade(source, configScanDir, c.Bool("strict-validation"))
-				},
-			},
-			{
-				Name: "reset",
-				Flags: []cli.Flag{
-					&cli.StringFlag{
-						Name:  "source",
-						Usage: "Source for upgrade. Composed of `type:address`. Accepts `file:`,`dir:` or `oci:` for the type of source.\nFor example `file:/var/share/myimage.tar`, `dir:/tmp/extracted` or `oci:repo/image:tag`",
-						Action: func(c *cli.Context, s string) error {
-							return validateSource(s)
-						},
-					},
-				},
-				Before: func(c *cli.Context) error {
-					return fmt.Errorf("not implemented")
-				},
-				Action: func(c *cli.Context) error {
-					config, err := agentConfig.Scan(collector.Directories(configScanDir...), collector.NoLogs, collector.StrictValidation(c.Bool("strict-validation")))
-					if err != nil {
-						return err
-					}
-
-					// Load the spec from the config
-					resetSpec, err := agentConfig.ReadUkiResetSpecFromConfig(config)
-					if err != nil {
-						return err
-					}
-
-					resetAction := uki.NewResetAction(config, resetSpec)
-					return resetAction.Run()
-				},
-			},
-		},
-	},
 	{
 		Name:        "versioneer",
 		Usage:       "versioneer subcommands",