From 7d18e6504a7e54d554bba6801d6a4860da67c884 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Tue, 26 Sep 2023 17:09:24 +0300
Subject: [PATCH 01/23] [WIP] Add `--source` to various commands

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/install.go             | 39 ++++++++++++----
 internal/agent/interactive_install.go |  7 +--
 main.go                               | 65 ++++++++++++++++++++++-----
 3 files changed, 87 insertions(+), 24 deletions(-)

diff --git a/internal/agent/install.go b/internal/agent/install.go
index 673c629..d2c99d0 100644
--- a/internal/agent/install.go
+++ b/internal/agent/install.go
@@ -53,16 +53,16 @@ func displayInfo(agentConfig *Config) {
 	}
 }
 
-func ManualInstall(c, device string, reboot, poweroff, strictValidations bool) error {
+func ManualInstall(c, sourceImg, device string, reboot, poweroff, strictValidations bool) error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	source, err := prepareConfiguration(ctx, c)
+	ConfigSource, err := prepareConfiguration(ctx, c)
 	if err != nil {
 		return err
 	}
 
-	cc, err := config.Scan(collector.Directories(source), collector.MergeBootLine, collector.StrictValidation(strictValidations), collector.NoLogs)
+	cc, err := config.Scan(collector.Directories(ConfigSource), collector.MergeBootLine, collector.StrictValidation(strictValidations), collector.NoLogs)
 	if err != nil {
 		return err
 	}
@@ -80,10 +80,10 @@ func ManualInstall(c, device string, reboot, poweroff, strictValidations bool) e
 		// Override from flags!
 		cc.Install.Device = device
 	}
-	return RunInstall(cc)
+	return RunInstall(cc, sourceImg)
 }
 
-func Install(dir ...string) error {
+func Install(sourceImg string, dir ...string) error {
 	var cc *config.Config
 	var err error
 
@@ -120,7 +120,7 @@ func Install(dir ...string) error {
 	// runs the installation
 	cc, err = config.Scan(collector.Directories(dir...), collector.MergeBootLine)
 	if err == nil && cc.Install != nil && cc.Install.Auto {
-		err = RunInstall(cc)
+		err = RunInstall(cc, sourceImg)
 		if err != nil {
 			return err
 		}
@@ -184,7 +184,7 @@ func Install(dir ...string) error {
 	pterm.Info.Println("Starting installation")
 
 	cc.Logger.Debugf("Runinstall with cc: %s\n", litter.Sdump(cc))
-	if err := RunInstall(cc); err != nil {
+	if err := RunInstall(cc, sourceImg); err != nil {
 		return err
 	}
 
@@ -213,7 +213,7 @@ func Install(dir ...string) error {
 	return nil
 }
 
-func RunInstall(c *config.Config) error {
+func RunInstall(c *config.Config, sourceImg string) error {
 	utils.SetEnv(c.Env)
 	utils.SetEnv(c.Install.Env)
 
@@ -250,7 +250,28 @@ func RunInstall(c *config.Config) error {
 	// Set our cloud-init to the file we just created
 	installSpec.CloudInit = append(installSpec.CloudInit, f.Name())
 	// Get the source of the installation if we are overriding it
-	if c.Install.Image != "" {
+	if sourceImg != "" {
+		imgSource, err := v1.NewSrcFromURI(sourceImg)
+		if err != nil {
+			return err
+		}
+		installSpec.Active.Source = imgSource
+
+		// TODO: Why only setting active source above? What about size?
+		// TODO: These 2 blocks are identical, DRY them.
+
+		// size, err := GetSourceSize(cfg, imgSource)
+		// if err != nil {
+		// 	c.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+		// }
+
+		// installSpec.Active.Source = imgSource
+		// installSpec.Passive.Source = imgSource
+		// installSpec.Recovery.Source = imgSource
+		// installSpec.Active.Size = uint(size)
+		// installSpec.Passive.Size = uint(size)
+		// installSpec.Recovery.Size = uint(size)
+	} else if c.Install.Image != "" {
 		imgSource, err := v1.NewSrcFromURI(c.Install.Image)
 		if err != nil {
 			return err
diff --git a/internal/agent/interactive_install.go b/internal/agent/interactive_install.go
index 7279300..8b65b94 100644
--- a/internal/agent/interactive_install.go
+++ b/internal/agent/interactive_install.go
@@ -129,7 +129,7 @@ func detectDevice() string {
 	return preferedDevice
 }
 
-func InteractiveInstall(debug, spawnShell bool) error {
+func InteractiveInstall(debug, spawnShell bool, sourceImg string) error {
 	var sshUsers []string
 	bus.Manager.Initialize()
 
@@ -229,7 +229,7 @@ func InteractiveInstall(debug, spawnShell bool) error {
 	}
 
 	if !isYes(allGood) {
-		return InteractiveInstall(debug, spawnShell)
+		return InteractiveInstall(debug, spawnShell, sourceImg)
 	}
 
 	usersToSet := map[string]schema.User{}
@@ -283,6 +283,7 @@ func InteractiveInstall(debug, spawnShell bool) error {
 			fmt.Printf("could not write event cloud init: %s\n", err.Error())
 		}
 		// override cc with our new config object from the scan, so it's updated for the RunInstall function
+		// TODO: Alternative solution: pass a reader here (the new feature) and add the image source
 		cc, _ = config.Scan(collector.Directories(tmpdir), collector.MergeBootLine, collector.NoLogs)
 	}
 
@@ -291,7 +292,7 @@ func InteractiveInstall(debug, spawnShell bool) error {
 	ccString, _ := cc.String()
 	pterm.Info.Println(ccString)
 
-	err = RunInstall(cc)
+	err = RunInstall(cc, sourceImg)
 	if err != nil {
 		pterm.Error.Println(err.Error())
 	}
diff --git a/main.go b/main.go
index 3d34a36..0649b75 100644
--- a/main.go
+++ b/main.go
@@ -110,15 +110,8 @@ See https://kairos.io/docs/upgrade/manual/ for documentation.
 			},
 		},
 		Before: func(c *cli.Context) error {
-			source := c.String("source")
-			if source != "" {
-				r, err := regexp.Compile(`^oci:|dir:|file:`)
-				if err != nil {
-					return nil
-				}
-				if !r.MatchString(source) {
-					return fmt.Errorf("source %s does not match any of oci:, dir: or file: ", source)
-				}
+			if err := validateSource(c.String("source")); err != nil {
+				return err
 			}
 
 			return checkRoot()
@@ -385,13 +378,23 @@ This command is meant to be used from the boot GRUB menu, but can be also starte
 			&cli.BoolFlag{
 				Name: "shell",
 			},
+			&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`",
+			},
 		},
 		Usage: "Starts interactive installation",
 		Before: func(c *cli.Context) error {
+			if err := validateSource(c.String("source")); err != nil {
+				return err
+			}
+
 			return checkRoot()
 		},
 		Action: func(c *cli.Context) error {
-			return agent.InteractiveInstall(c.Bool("debug"), c.Bool("shell"))
+			source := c.String("source")
+
+			return agent.InteractiveInstall(c.Bool("debug"), c.Bool("shell"), source)
 		},
 	},
 	{
@@ -410,8 +413,16 @@ This command is meant to be used from the boot GRUB menu, but can be also starte
 			&cli.BoolFlag{
 				Name: "reboot",
 			},
+			&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`",
+			},
 		},
 		Before: func(c *cli.Context) error {
+			if err := validateSource(c.String("source")); err != nil {
+				return err
+			}
+
 			return checkRoot()
 		},
 		Action: func(c *cli.Context) error {
@@ -420,7 +431,9 @@ This command is meant to be used from the boot GRUB menu, but can be also starte
 			}
 			config := c.Args().First()
 
-			return agent.ManualInstall(config, c.String("device"), c.Bool("reboot"), c.Bool("poweroff"), c.Bool("strict-validation"))
+			source := c.String("source")
+
+			return agent.ManualInstall(config, source, c.String("device"), c.Bool("reboot"), c.Bool("poweroff"), c.Bool("strict-validation"))
 		},
 	},
 	{
@@ -436,10 +449,22 @@ See also https://kairos.io/docs/installation/qrcode/ for documentation.
 This command is meant to be used from the boot GRUB menu, but can be started manually`,
 		Aliases: []string{"i"},
 		Before: func(c *cli.Context) error {
+			if err := validateSource(c.String("source")); err != nil {
+				return err
+			}
+
 			return checkRoot()
 		},
+		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) error {
-			return agent.Install(configScanDir...)
+			source := c.String("source")
+
+			return agent.Install(source, configScanDir...)
 		},
 	},
 	{
@@ -686,3 +711,19 @@ func checkRoot() error {
 
 	return nil
 }
+
+func validateSource(source string) error {
+	if source == "" {
+		return nil
+	}
+
+	r, err := regexp.Compile(`^oci:|dir:|file:`)
+	if err != nil {
+		return err
+	}
+	if !r.MatchString(source) {
+		return fmt.Errorf("source %s does not match any of oci:, dir: or file: ", source)
+	}
+
+	return nil
+}

From 3482e88aa387a41469b3e4fa694ff0a76a284f7c Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Wed, 27 Sep 2023 17:38:55 +0300
Subject: [PATCH 02/23] [WIP] Try to pass command line options for upgrade as
 kairos config

because it's actually configuration and it's better than using viper to
pass things around.

The `upgrade:` block seems to be ignored early in the process so even if
the user specifies an image there, the latest release will be looked up.

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/upgrade.go | 62 ++++++++++++++++++++++++++++++++++-----
 pkg/config/spec.go        | 62 +++++++++++++++++++++++----------------
 pkg/types/v1/common.go    | 36 +++++++++++------------
 3 files changed, 108 insertions(+), 52 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 8231a34..7bb4507 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -5,8 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"sort"
-
-	"github.com/spf13/viper"
+	"strings"
 
 	hook "github.com/kairos-io/kairos-agent/v2/internal/agent/hooks"
 
@@ -55,6 +54,10 @@ func Upgrade(
 	version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) error {
 	bus.Manager.Initialize()
 
+	// TODO: Before we check for empy source,
+	// shouldn't we read the `upgrade:` block from the config and check if something
+	// is defined there?
+
 	if version == "" && source == "" {
 		fmt.Println("Searching for releases")
 		if preReleases {
@@ -93,13 +96,11 @@ func Upgrade(
 		}
 	}
 
-	// Set this here with viper help so we can use it while creating the upgrade spec
-	// And its properly set since creation without having to modify it later
-	// This should be binded somehow but the current cli doesnt allow us to bind flags to values
-	viper.Set("upgradeSource", img)
-	viper.Set("upgradeRecovery", upgradeRecovery)
+	upgradeConf := generateUpgradeConf(img, upgradeRecovery)
 
-	c, err := config.Scan(collector.Directories(dirs...), collector.StrictValidation(strictValidations))
+	c, err := config.Scan(collector.Directories(dirs...),
+		collector.Readers(strings.NewReader(upgradeConf)),
+		collector.StrictValidation(strictValidations))
 	if err != nil {
 		return err
 	}
@@ -162,3 +163,48 @@ func determineUpgradeImage(version string) (string, error) {
 
 	return fmt.Sprintf("%s:%s", registry, version), nil
 }
+
+// generateUpgradeConf creates a kairos configuration for upgrade to be
+// added to the rest of the configurations.
+func generateUpgradeConf(source string, upgradeRecovery bool) string {
+	conf := ""
+
+	if source == "" {
+		return conf
+	}
+
+	// TODO: Do the same for recovery upgrade
+	conf = fmt.Sprintf(`
+upgrade:
+	recovery-system:
+		uri: %s`, source)
+
+	// source := viper.GetString("upgradeSource")
+	// recoveryUpgrade := viper.GetBool("upgradeRecovery")
+	// if source != "" {
+	// 	imgSource, err := v1.NewSrcFromURI(source)
+	// 	// TODO: Don't hide the error here!
+	// 	if err == nil {
+	// 		if recoveryUpgrade {
+	// 			spec.RecoveryUpgrade = recoveryUpgrade
+	// 			spec.Recovery.Source = imgSource
+	// 		} else {
+	// 			spec.Active.Source = imgSource
+	// 		}
+	// 		size, err := GetSourceSize(cfg, imgSource)
+	// 		if err != nil {
+	// 			cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+	// 		} else {
+	// 			cfg.Logger.Infof("Setting image size to %dMb", size)
+	// 			// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
+	// 			if recoveryUpgrade {
+	// 				spec.Recovery.Size = uint(size)
+	// 			} else {
+	// 				spec.Active.Size = uint(size)
+	// 			}
+	// 		}
+	// 	}
+	//}
+
+	return conf
+}
diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index cfce3e2..b6f5b52 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -27,7 +27,7 @@ import (
 	"github.com/kairos-io/kairos-agent/v2/internal/common"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	"github.com/kairos-io/kairos-agent/v2/pkg/utils/partitions"
 	"github.com/mitchellh/mapstructure"
 	"github.com/sanity-io/litter"
@@ -265,31 +265,32 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 		State:      installState,
 	}
 
-	source := viper.GetString("upgradeSource")
-	recoveryUpgrade := viper.GetBool("upgradeRecovery")
-	if source != "" {
-		imgSource, err := v1.NewSrcFromURI(source)
-		if err == nil {
-			if recoveryUpgrade {
-				spec.RecoveryUpgrade = recoveryUpgrade
-				spec.Recovery.Source = imgSource
-			} else {
-				spec.Active.Source = imgSource
-			}
-			size, err := GetSourceSize(cfg, imgSource)
-			if err != nil {
-				cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
-			} else {
-				cfg.Logger.Infof("Setting image size to %dMb", size)
-				// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
-				if recoveryUpgrade {
-					spec.Recovery.Size = uint(size)
-				} else {
-					spec.Active.Size = uint(size)
-				}
-			}
-		}
-	}
+	// source := viper.GetString("upgradeSource")
+	// recoveryUpgrade := viper.GetBool("upgradeRecovery")
+	// if source != "" {
+	// 	imgSource, err := v1.NewSrcFromURI(source)
+	// 	// TODO: Don't hide the error here!
+	// 	if err == nil {
+	// 		if recoveryUpgrade {
+	// 			spec.RecoveryUpgrade = recoveryUpgrade
+	// 			spec.Recovery.Source = imgSource
+	// 		} else {
+	// 			spec.Active.Source = imgSource
+	// 		}
+	// 		size, err := GetSourceSize(cfg, imgSource)
+	// 		if err != nil {
+	// 			cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+	// 		} else {
+	// 			cfg.Logger.Infof("Setting image size to %dMb", size)
+	// 			// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
+	// 			if recoveryUpgrade {
+	// 				spec.Recovery.Size = uint(size)
+	// 			} else {
+	// 				spec.Active.Size = uint(size)
+	// 			}
+	// 		}
+	// 	}
+	//}
 
 	return spec, nil
 }
@@ -543,6 +544,8 @@ func ReadUpgradeSpecFromConfig(c *Config) (*v1.UpgradeSpec, error) {
 
 // ReadSpecFromCloudConfig returns a v1.Spec for the given spec
 func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
+	fmt.Printf("litter.Sdump(r) before = %+v\n", litter.Sdump(r))
+
 	var sp v1.Spec
 	var err error
 
@@ -560,6 +563,8 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 		return nil, fmt.Errorf("failed initializing spec: %v", err)
 	}
 
+	fmt.Printf("litter.Sdump(sp) before = %+v\n", litter.Sdump(sp))
+
 	// Load the config into viper from the raw cloud config string
 	ccString, err := r.String()
 	if err != nil {
@@ -576,6 +581,11 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 	if err != nil {
 		r.Logger.Warnf("error unmarshalling %s Spec: %s", spec, err)
 	}
+
+	fmt.Println("---------------------------------")
+	fmt.Printf("litter.Sdump(sp) after = %+v\n", litter.Sdump(sp))
+	fmt.Println("---------------------------------")
+
 	err = sp.Sanitize()
 	if err != nil {
 		r.Logger.Warnf("Error sanitizing the % spec: %s", spec, err)
diff --git a/pkg/types/v1/common.go b/pkg/types/v1/common.go
index eeb9302..5f75a87 100644
--- a/pkg/types/v1/common.go
+++ b/pkg/types/v1/common.go
@@ -35,31 +35,31 @@ const (
 
 // ImageSource represents the source from where an image is created for easy identification
 type ImageSource struct {
-	source  string
-	srcType string
+	Source  string
+	SrcType string
 }
 
 func (i ImageSource) Value() string {
-	return i.source
+	return i.Source
 }
 
 func (i ImageSource) IsDocker() bool {
-	return i.srcType == oci
+	return i.SrcType == oci
 }
 
 func (i ImageSource) IsDir() bool {
-	return i.srcType == dir
+	return i.SrcType == dir
 }
 
 func (i ImageSource) IsFile() bool {
-	return i.srcType == file
+	return i.SrcType == file
 }
 
 func (i ImageSource) IsEmpty() bool {
-	if i.srcType == "" {
+	if i.SrcType == "" {
 		return true
 	}
-	if i.source == "" {
+	if i.Source == "" {
 		return true
 	}
 	return false
@@ -69,7 +69,7 @@ func (i ImageSource) String() string {
 	if i.IsEmpty() {
 		return ""
 	}
-	return fmt.Sprintf("%s://%s", i.srcType, i.source)
+	return fmt.Sprintf("%s://%s", i.SrcType, i.Source)
 }
 
 func (i ImageSource) MarshalYAML() (interface{}, error) {
@@ -103,11 +103,11 @@ func (i *ImageSource) updateFromURI(uri string) error {
 	case oci, docker:
 		return i.parseImageReference(value)
 	case dir:
-		i.srcType = dir
-		i.source = value
+		i.SrcType = dir
+		i.Source = value
 	case file:
-		i.srcType = file
-		i.source = value
+		i.SrcType = file
+		i.Source = value
 	default:
 		return i.parseImageReference(uri)
 	}
@@ -121,8 +121,8 @@ func (i *ImageSource) parseImageReference(ref string) error {
 	} else if reference.IsNameOnly(n) {
 		ref += ":latest"
 	}
-	i.srcType = oci
-	i.source = ref
+	i.SrcType = oci
+	i.Source = ref
 	return nil
 }
 
@@ -137,13 +137,13 @@ func NewEmptySrc() *ImageSource {
 }
 
 func NewDockerSrc(src string) *ImageSource {
-	return &ImageSource{source: src, srcType: oci}
+	return &ImageSource{Source: src, SrcType: oci}
 }
 
 func NewFileSrc(src string) *ImageSource {
-	return &ImageSource{source: src, srcType: file}
+	return &ImageSource{Source: src, SrcType: file}
 }
 
 func NewDirSrc(src string) *ImageSource {
-	return &ImageSource{source: src, srcType: dir}
+	return &ImageSource{Source: src, SrcType: dir}
 }

From 60815d060ea2d05c174c70ba3aa05e3fff756d5c Mon Sep 17 00:00:00 2001
From: Itxaka Serrano Garcia <itxaka.garcia@spectrocloud.com>
Date: Thu, 28 Sep 2023 12:27:27 +0300
Subject: [PATCH 03/23] Refactor

Signed-off-by: Dimitris Karakasilis <dimitris@spectrocloud.com>
---
 internal/agent/upgrade.go | 152 +++++++++++++++++++-------------------
 1 file changed, 75 insertions(+), 77 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 7bb4507..2de16c1 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -8,12 +8,14 @@ import (
 	"strings"
 
 	hook "github.com/kairos-io/kairos-agent/v2/internal/agent/hooks"
+	"github.com/sanity-io/litter"
 
 	"github.com/Masterminds/semver/v3"
 	"github.com/kairos-io/kairos-agent/v2/internal/bus"
 	"github.com/kairos-io/kairos-agent/v2/pkg/action"
 	"github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/github"
+	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
 	events "github.com/kairos-io/kairos-sdk/bus"
 	"github.com/kairos-io/kairos-sdk/collector"
 	"github.com/kairos-io/kairos-sdk/utils"
@@ -54,49 +56,7 @@ func Upgrade(
 	version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) error {
 	bus.Manager.Initialize()
 
-	// TODO: Before we check for empy source,
-	// shouldn't we read the `upgrade:` block from the config and check if something
-	// is defined there?
-
-	if version == "" && source == "" {
-		fmt.Println("Searching for releases")
-		if preReleases {
-			fmt.Println("Including pre-releases")
-		}
-		releases := ListReleases(preReleases)
-
-		if len(releases) == 0 {
-			return fmt.Errorf("no releases found")
-		}
-
-		// Using Original here because the parsing removes the v as its a semver. But it stores the original full version there
-		version = releases[0].Original()
-
-		if utils.Version() == version && !force {
-			fmt.Printf("version %s already installed. use --force to force upgrade\n", version)
-			return nil
-		}
-		msg := fmt.Sprintf("Latest release is %s\nAre you sure you want to upgrade to this release? (y/n)", version)
-		reply, err := promptBool(events.YAMLPrompt{Prompt: msg, Default: "y"})
-		if err != nil {
-			return err
-		}
-		if reply == "false" {
-			return nil
-		}
-	}
-
-	img := source
-	var err error
-	if img == "" {
-		img, err = determineUpgradeImage(version)
-		if err != nil {
-			fmt.Println(err.Error())
-			return err
-		}
-	}
-
-	upgradeConf := generateUpgradeConf(img, upgradeRecovery)
+	upgradeConf := generateUpgradeConf(source, upgradeRecovery)
 
 	c, err := config.Scan(collector.Directories(dirs...),
 		collector.Readers(strings.NewReader(upgradeConf)),
@@ -113,6 +73,14 @@ func Upgrade(
 		return err
 	}
 
+	err = handleEmptySource(upgradeSpec, version, preReleases, force)
+	if err != nil {
+		return err
+	}
+
+	// TODO: Remove this line
+	litter.Dump(upgradeSpec)
+
 	// Sanitize
 	err = upgradeSpec.Sanitize()
 	if err != nil {
@@ -139,7 +107,7 @@ func Upgrade(
 
 // determineUpgradeImage asks the provider plugin for an image or constructs
 // it using version and data from /etc/os-release
-func determineUpgradeImage(version string) (string, error) {
+func determineUpgradeImage(version string) (*v1.ImageSource, error) {
 	var img string
 	bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
 		img = r.Data
@@ -149,19 +117,19 @@ func determineUpgradeImage(version string) (string, error) {
 		Version: version,
 	})
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	if img != "" {
-		return img, nil
+		return nil, nil
 	}
 
 	registry, err := utils.OSRelease("IMAGE_REPO")
 	if err != nil {
-		return "", fmt.Errorf("can't find IMAGE_REPO key under /etc/os-release %w", err)
+		return nil, fmt.Errorf("can't find IMAGE_REPO key under /etc/os-release %w", err)
 	}
 
-	return fmt.Sprintf("%s:%s", registry, version), nil
+	return v1.NewSrcFromURI(fmt.Sprintf("%s:%s", registry, version))
 }
 
 // generateUpgradeConf creates a kairos configuration for upgrade to be
@@ -173,38 +141,68 @@ func generateUpgradeConf(source string, upgradeRecovery bool) string {
 		return conf
 	}
 
-	// TODO: Do the same for recovery upgrade
 	conf = fmt.Sprintf(`
 upgrade:
+	recovery: %t
 	recovery-system:
-		uri: %s`, source)
-
-	// source := viper.GetString("upgradeSource")
-	// recoveryUpgrade := viper.GetBool("upgradeRecovery")
-	// if source != "" {
-	// 	imgSource, err := v1.NewSrcFromURI(source)
-	// 	// TODO: Don't hide the error here!
-	// 	if err == nil {
-	// 		if recoveryUpgrade {
-	// 			spec.RecoveryUpgrade = recoveryUpgrade
-	// 			spec.Recovery.Source = imgSource
-	// 		} else {
-	// 			spec.Active.Source = imgSource
-	// 		}
-	// 		size, err := GetSourceSize(cfg, imgSource)
-	// 		if err != nil {
-	// 			cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
-	// 		} else {
-	// 			cfg.Logger.Infof("Setting image size to %dMb", size)
-	// 			// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
-	// 			if recoveryUpgrade {
-	// 				spec.Recovery.Size = uint(size)
-	// 			} else {
-	// 				spec.Active.Size = uint(size)
-	// 			}
-	// 		}
-	// 	}
-	//}
+		uri: %s`, upgradeRecovery, source)
 
 	return conf
 }
+
+func handleEmptySource(spec *v1.UpgradeSpec, version string, preReleases, force bool) error {
+	var err error
+	if spec.RecoveryUpgrade {
+		if spec.Recovery.Source.IsEmpty() {
+			spec.Recovery.Source, err = getLatestOrConstructSource(version, preReleases, force)
+		}
+	} else {
+		if spec.Active.Source.IsEmpty() {
+			spec.Active.Source, err = getLatestOrConstructSource(version, preReleases, force)
+		}
+	}
+
+	return err
+}
+
+func getLatestOrConstructSource(version string, preReleases, force bool) (*v1.ImageSource, error) {
+	var err error
+	if version == "" {
+		version, err = findLatestVersion(preReleases, force)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return determineUpgradeImage(version)
+}
+
+func findLatestVersion(preReleases, force bool) (string, error) {
+	fmt.Println("Searching for releases")
+	if preReleases {
+		fmt.Println("Including pre-releases")
+	}
+	releases := ListReleases(preReleases)
+
+	if len(releases) == 0 {
+		return "", fmt.Errorf("no releases found")
+	}
+
+	// Using Original here because the parsing removes the v as its a semver. But it stores the original full version there
+	version := releases[0].Original()
+
+	if utils.Version() == version && !force {
+		return "", fmt.Errorf("version %s already installed. use --force to force upgrade\n", version)
+	}
+
+	msg := fmt.Sprintf("Latest release is %s\nAre you sure you want to upgrade to this release? (y/n)", version)
+	reply, err := promptBool(events.YAMLPrompt{Prompt: msg, Default: "y"})
+	if err != nil {
+		return "", err
+	}
+	if reply == "false" {
+		return "", fmt.Errorf("cancelled by the user")
+	}
+
+	return version, nil
+}

From 8dd7577148b1ebb2439206cf55743fa76832ff1d Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 13:04:03 +0300
Subject: [PATCH 04/23] WIP

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/reset.go   |  5 +++++
 internal/agent/upgrade.go | 27 +++++++++++++++++++++------
 pkg/config/spec.go        |  4 ----
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/internal/agent/reset.go b/internal/agent/reset.go
index bbc9b0c..e6aba47 100644
--- a/internal/agent/reset.go
+++ b/internal/agent/reset.go
@@ -102,6 +102,11 @@ func Reset(reboot, unattended bool, dir ...string) error {
 		resetSpec.Reboot = reboot
 	}
 
+	err = resetSpec.Sanitize()
+	if err != nil {
+		return err
+	}
+
 	resetAction := action.NewResetAction(c, resetSpec)
 	if err := resetAction.Run(); err != nil {
 		fmt.Println(err)
diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 2de16c1..3b4003e 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -53,7 +53,7 @@ func ListReleases(includePrereleases bool) semver.Collection {
 }
 
 func Upgrade(
-	version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) error {
+	version, source string, force, strictValidations bool, dirs []string, preReleases bool, upgradeRecovery bool) error {
 	bus.Manager.Initialize()
 
 	upgradeConf := generateUpgradeConf(source, upgradeRecovery)
@@ -141,11 +141,26 @@ func generateUpgradeConf(source string, upgradeRecovery bool) string {
 		return conf
 	}
 
-	conf = fmt.Sprintf(`
-upgrade:
-	recovery: %t
-	recovery-system:
-		uri: %s`, upgradeRecovery, source)
+	upgrade := struct {
+		Recovery: Bool `json:"recovery,omitempy"`,
+		RecoverySystem: struct{ `json:"recovery-system,omitempty"`
+		  URI: string `json:"uri,omitempy"`
+		},
+		System: struct{ `json:"system,omitempy"`
+			URI: string
+		}
+	}{
+	}
+
+
+	conf = fmt.Sprintf(`upgrade:
+  recovery-system:
+    uri: %s
+`, source)
+
+	if upgradeRecovery {
+		conf += `  recovery: true`
+	}
 
 	return conf
 }
diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index b6f5b52..cda731d 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -586,10 +586,6 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 	fmt.Printf("litter.Sdump(sp) after = %+v\n", litter.Sdump(sp))
 	fmt.Println("---------------------------------")
 
-	err = sp.Sanitize()
-	if err != nil {
-		r.Logger.Warnf("Error sanitizing the % spec: %s", spec, err)
-	}
 	r.Logger.Debugf("Loaded %s spec: %s", spec, litter.Sdump(sp))
 	return sp, nil
 }

From afe08c5ca5b665ef2abecaca8e1594424ad94d7e Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 14:19:48 +0300
Subject: [PATCH 05/23] Generate additional config only if needed

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/upgrade.go | 54 +++++++++++++++++++--------------------
 1 file changed, 27 insertions(+), 27 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 3b4003e..d19a21d 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -56,7 +56,10 @@ func Upgrade(
 	version, source string, force, strictValidations bool, dirs []string, preReleases bool, upgradeRecovery bool) error {
 	bus.Manager.Initialize()
 
-	upgradeConf := generateUpgradeConf(source, upgradeRecovery)
+	upgradeConf, err := generateUpgradeConf(source, upgradeRecovery)
+	if err != nil {
+		return err
+	}
 
 	c, err := config.Scan(collector.Directories(dirs...),
 		collector.Readers(strings.NewReader(upgradeConf)),
@@ -132,37 +135,34 @@ func determineUpgradeImage(version string) (*v1.ImageSource, error) {
 	return v1.NewSrcFromURI(fmt.Sprintf("%s:%s", registry, version))
 }
 
-// generateUpgradeConf creates a kairos configuration for upgrade to be
-// added to the rest of the configurations.
-func generateUpgradeConf(source string, upgradeRecovery bool) string {
-	conf := ""
-
-	if source == "" {
-		return conf
+// generateUpgradeConf creates a kairos configuration for `--source` and `--recovery`
+// command line arguments. It will be added to the rest of the configurations.
+func generateUpgradeConf(source string, upgradeRecovery bool) (string, error) {
+	upgrade := map[string](map[string]interface{}){
+		"upgrade": {},
 	}
 
-	upgrade := struct {
-		Recovery: Bool `json:"recovery,omitempy"`,
-		RecoverySystem: struct{ `json:"recovery-system,omitempty"`
-		  URI: string `json:"uri,omitempy"`
-		},
-		System: struct{ `json:"system,omitempy"`
-			URI: string
-		}
-	}{
-	}
-
-
-	conf = fmt.Sprintf(`upgrade:
-  recovery-system:
-    uri: %s
-`, source)
-
 	if upgradeRecovery {
-		conf += `  recovery: true`
+		upgrade["upgrade"]["recovery"] = "true"
 	}
 
-	return conf
+	// Set uri both for active and recovery because we don't know what we are
+	// actually upgrading. The "upgradeRecovery" is just the command line argument.
+	// The user might have set it to "true" in the kairos config. Since we don't
+	// have access to that yet, we just set both uri values which shouldn't matter
+	// anyway, the right one will be used later in the process.
+	if source != "" {
+		upgrade["upgrade"]["recovery-system"] = map[string]string{
+			"uri": source,
+		}
+		upgrade["upgrade"]["system"] = map[string]string{
+			"uri": source,
+		}
+	}
+
+	d, err := json.Marshal(upgrade)
+
+	return string(d), err
 }
 
 func handleEmptySource(spec *v1.UpgradeSpec, version string, preReleases, force bool) error {

From 968812dfbb5286a0f618ce10def78d310bb40d33 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 14:50:14 +0300
Subject: [PATCH 06/23] Exctract the generation of the upgrade spec to a
 function

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/upgrade.go | 71 +++++++++++++++++++++------------------
 1 file changed, 38 insertions(+), 33 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index d19a21d..4c6d5db 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -8,12 +8,11 @@ import (
 	"strings"
 
 	hook "github.com/kairos-io/kairos-agent/v2/internal/agent/hooks"
-	"github.com/sanity-io/litter"
 
 	"github.com/Masterminds/semver/v3"
 	"github.com/kairos-io/kairos-agent/v2/internal/bus"
 	"github.com/kairos-io/kairos-agent/v2/pkg/action"
-	"github.com/kairos-io/kairos-agent/v2/pkg/config"
+	config "github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/github"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
 	events "github.com/kairos-io/kairos-sdk/bus"
@@ -56,36 +55,7 @@ func Upgrade(
 	version, source string, force, strictValidations bool, dirs []string, preReleases bool, upgradeRecovery bool) error {
 	bus.Manager.Initialize()
 
-	upgradeConf, err := generateUpgradeConf(source, upgradeRecovery)
-	if err != nil {
-		return err
-	}
-
-	c, err := config.Scan(collector.Directories(dirs...),
-		collector.Readers(strings.NewReader(upgradeConf)),
-		collector.StrictValidation(strictValidations))
-	if err != nil {
-		return err
-	}
-
-	utils.SetEnv(c.Env)
-
-	// Load the upgrade Config from the system
-	upgradeSpec, err := config.ReadUpgradeSpecFromConfig(c)
-	if err != nil {
-		return err
-	}
-
-	err = handleEmptySource(upgradeSpec, version, preReleases, force)
-	if err != nil {
-		return err
-	}
-
-	// TODO: Remove this line
-	litter.Dump(upgradeSpec)
-
-	// Sanitize
-	err = upgradeSpec.Sanitize()
+	upgradeSpec, c, err := generateUpgradeSpec(version, source, force, strictValidations, dirs, preReleases, upgradeRecovery)
 	if err != nil {
 		return err
 	}
@@ -207,7 +177,7 @@ func findLatestVersion(preReleases, force bool) (string, error) {
 	version := releases[0].Original()
 
 	if utils.Version() == version && !force {
-		return "", fmt.Errorf("version %s already installed. use --force to force upgrade\n", version)
+		return "", fmt.Errorf("version %s already installed. use --force to force upgrade", version)
 	}
 
 	msg := fmt.Sprintf("Latest release is %s\nAre you sure you want to upgrade to this release? (y/n)", version)
@@ -221,3 +191,38 @@ func findLatestVersion(preReleases, force bool) (string, error) {
 
 	return version, nil
 }
+
+func generateUpgradeSpec(version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) (*v1.UpgradeSpec, *config.Config, error) {
+	upgradeConf, err := generateUpgradeConf(source, upgradeRecovery)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	c, err := config.Scan(collector.Directories(dirs...),
+		collector.Readers(strings.NewReader(upgradeConf)),
+		collector.StrictValidation(strictValidations))
+	if err != nil {
+		return nil, nil, err
+	}
+
+	utils.SetEnv(c.Env)
+
+	// Load the upgrade Config from the system
+	upgradeSpec, err := config.ReadUpgradeSpecFromConfig(c)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = handleEmptySource(upgradeSpec, version, preReleases, force)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// Sanitize
+	err = upgradeSpec.Sanitize()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return upgradeSpec, c, nil
+}

From c58efaa623a9c9b44b15b45898ea670dd03fe436 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 15:08:16 +0300
Subject: [PATCH 07/23] Add test and remove debugging logs

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/install_test.go |  9 ++++---
 internal/agent/upgrade.go      | 19 +++++++-------
 internal/agent/upgrade_test.go | 47 ++++++++++++++++++++++++++++++++++
 pkg/config/spec.go             |  8 ------
 4 files changed, 61 insertions(+), 22 deletions(-)
 create mode 100644 internal/agent/upgrade_test.go

diff --git a/internal/agent/install_test.go b/internal/agent/install_test.go
index a7321f6..b5fdd79 100644
--- a/internal/agent/install_test.go
+++ b/internal/agent/install_test.go
@@ -3,13 +3,14 @@ package agent
 import (
 	"context"
 	"fmt"
-	"github.com/jaypipes/ghw/pkg/block"
-	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	"os"
 	"path/filepath"
 
+	"github.com/jaypipes/ghw/pkg/block"
+	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
+
 	"github.com/kairos-io/kairos-agent/v2/pkg/config"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	v1mock "github.com/kairos-io/kairos-agent/v2/tests/mocks"
 	"github.com/twpayne/go-vfs"
 	"github.com/twpayne/go-vfs/vfst"
@@ -196,7 +197,7 @@ var _ = Describe("RunInstall", func() {
 
 	It("runs the install", func() {
 		Skip("Not ready yet")
-		err = RunInstall(options)
+		err = RunInstall(options, "") // TODO
 		Expect(err).ToNot(HaveOccurred())
 	})
 })
diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 4c6d5db..7c0907a 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -60,6 +60,11 @@ func Upgrade(
 		return err
 	}
 
+	err = upgradeSpec.Sanitize()
+	if err != nil {
+		return err
+	}
+
 	upgradeAction := action.NewUpgradeAction(c, upgradeSpec)
 
 	err = upgradeAction.Run()
@@ -105,9 +110,9 @@ func determineUpgradeImage(version string) (*v1.ImageSource, error) {
 	return v1.NewSrcFromURI(fmt.Sprintf("%s:%s", registry, version))
 }
 
-// generateUpgradeConf creates a kairos configuration for `--source` and `--recovery`
+// generateConfForCLIArgs creates a kairos configuration for `--source` and `--recovery`
 // command line arguments. It will be added to the rest of the configurations.
-func generateUpgradeConf(source string, upgradeRecovery bool) (string, error) {
+func generateConfForCLIArgs(source string, upgradeRecovery bool) (string, error) {
 	upgrade := map[string](map[string]interface{}){
 		"upgrade": {},
 	}
@@ -193,13 +198,13 @@ func findLatestVersion(preReleases, force bool) (string, error) {
 }
 
 func generateUpgradeSpec(version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) (*v1.UpgradeSpec, *config.Config, error) {
-	upgradeConf, err := generateUpgradeConf(source, upgradeRecovery)
+	cliConf, err := generateConfForCLIArgs(source, upgradeRecovery)
 	if err != nil {
 		return nil, nil, err
 	}
 
 	c, err := config.Scan(collector.Directories(dirs...),
-		collector.Readers(strings.NewReader(upgradeConf)),
+		collector.Readers(strings.NewReader(cliConf)),
 		collector.StrictValidation(strictValidations))
 	if err != nil {
 		return nil, nil, err
@@ -218,11 +223,5 @@ func generateUpgradeSpec(version, source string, force, strictValidations bool,
 		return nil, nil, err
 	}
 
-	// Sanitize
-	err = upgradeSpec.Sanitize()
-	if err != nil {
-		return nil, nil, err
-	}
-
 	return upgradeSpec, c, nil
 }
diff --git a/internal/agent/upgrade_test.go b/internal/agent/upgrade_test.go
new file mode 100644
index 0000000..ffa4da8
--- /dev/null
+++ b/internal/agent/upgrade_test.go
@@ -0,0 +1,47 @@
+package agent
+
+import (
+	"fmt"
+	"os"
+	"path"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("generateUpgradeSpec", func() {
+	When("there are command line arguments", func() {
+		var configDir string
+		var upgradeRecovery bool
+
+		BeforeEach(func() {
+			upgradeRecovery = false
+			configDir, err := os.MkdirTemp("", "upgrade-test")
+			Expect(err).ToNot(HaveOccurred())
+
+			config := fmt.Sprintf(`upgrade:
+  recovery: %t
+  recovery-system:
+    uri: oci://image-in-conf
+`, upgradeRecovery)
+
+			configFilePath := path.Join(configDir, "config.yaml")
+			err = os.WriteFile(configFilePath, []byte(config), os.ModePerm)
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		AfterEach(func() {
+			os.RemoveAll(configDir)
+		})
+
+		It("overrides kairos config yaml values", func() {
+			spec, _, err := generateUpgradeSpec("", "oci:myimage", false, false,
+				[]string{}, false, !upgradeRecovery)
+
+			Expect(err).ToNot(HaveOccurred())
+			Expect(spec.Active.Source.String()).To(Equal("oci://myimage:latest"))
+			Expect(spec.Recovery.Source.String()).To(Equal("oci://myimage:latest"))
+			Expect(spec.RecoveryUpgrade).To(Equal(!upgradeRecovery))
+		})
+	})
+})
diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index cda731d..261c723 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -544,8 +544,6 @@ func ReadUpgradeSpecFromConfig(c *Config) (*v1.UpgradeSpec, error) {
 
 // ReadSpecFromCloudConfig returns a v1.Spec for the given spec
 func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
-	fmt.Printf("litter.Sdump(r) before = %+v\n", litter.Sdump(r))
-
 	var sp v1.Spec
 	var err error
 
@@ -563,8 +561,6 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 		return nil, fmt.Errorf("failed initializing spec: %v", err)
 	}
 
-	fmt.Printf("litter.Sdump(sp) before = %+v\n", litter.Sdump(sp))
-
 	// Load the config into viper from the raw cloud config string
 	ccString, err := r.String()
 	if err != nil {
@@ -582,10 +578,6 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 		r.Logger.Warnf("error unmarshalling %s Spec: %s", spec, err)
 	}
 
-	fmt.Println("---------------------------------")
-	fmt.Printf("litter.Sdump(sp) after = %+v\n", litter.Sdump(sp))
-	fmt.Println("---------------------------------")
-
 	r.Logger.Debugf("Loaded %s spec: %s", spec, litter.Sdump(sp))
 	return sp, nil
 }

From 15a926ad3b82b68c6ac89f29da3f0004ee509293 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 15:55:14 +0300
Subject: [PATCH 08/23] Handle `--source` in
 install/manual-install/interactive-install the same way

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/install.go             | 58 +++++++++++++--------------
 internal/agent/interactive_install.go | 10 +++--
 internal/agent/upgrade.go             |  6 +--
 3 files changed, 37 insertions(+), 37 deletions(-)

diff --git a/internal/agent/install.go b/internal/agent/install.go
index d2c99d0..3c23108 100644
--- a/internal/agent/install.go
+++ b/internal/agent/install.go
@@ -62,7 +62,12 @@ func ManualInstall(c, sourceImg, device string, reboot, poweroff, strictValidati
 		return err
 	}
 
-	cc, err := config.Scan(collector.Directories(ConfigSource), collector.MergeBootLine, collector.StrictValidation(strictValidations), collector.NoLogs)
+	cliConf := generateInstallConfForCLIArgs(sourceImg)
+
+	cc, err := config.Scan(collector.Directories(ConfigSource),
+		collector.Readers(strings.NewReader(cliConf)),
+		collector.MergeBootLine,
+		collector.StrictValidation(strictValidations), collector.NoLogs)
 	if err != nil {
 		return err
 	}
@@ -80,7 +85,7 @@ func ManualInstall(c, sourceImg, device string, reboot, poweroff, strictValidati
 		// Override from flags!
 		cc.Install.Device = device
 	}
-	return RunInstall(cc, sourceImg)
+	return RunInstall(cc)
 }
 
 func Install(sourceImg string, dir ...string) error {
@@ -116,11 +121,14 @@ func Install(sourceImg string, dir ...string) error {
 
 	ensureDataSourceReady()
 
-	// Reads config, and if present and offline is defined,
-	// runs the installation
-	cc, err = config.Scan(collector.Directories(dir...), collector.MergeBootLine)
+	cliConf := generateInstallConfForCLIArgs(sourceImg)
+
+	// Reads config, and if present and offline is defined, runs the installation
+	cc, err = config.Scan(collector.Directories(dir...),
+		collector.Readers(strings.NewReader(cliConf)),
+		collector.MergeBootLine)
 	if err == nil && cc.Install != nil && cc.Install.Auto {
-		err = RunInstall(cc, sourceImg)
+		err = RunInstall(cc)
 		if err != nil {
 			return err
 		}
@@ -184,7 +192,7 @@ func Install(sourceImg string, dir ...string) error {
 	pterm.Info.Println("Starting installation")
 
 	cc.Logger.Debugf("Runinstall with cc: %s\n", litter.Sdump(cc))
-	if err := RunInstall(cc, sourceImg); err != nil {
+	if err := RunInstall(cc); err != nil {
 		return err
 	}
 
@@ -213,7 +221,7 @@ func Install(sourceImg string, dir ...string) error {
 	return nil
 }
 
-func RunInstall(c *config.Config, sourceImg string) error {
+func RunInstall(c *config.Config) error {
 	utils.SetEnv(c.Env)
 	utils.SetEnv(c.Install.Env)
 
@@ -250,28 +258,7 @@ func RunInstall(c *config.Config, sourceImg string) error {
 	// Set our cloud-init to the file we just created
 	installSpec.CloudInit = append(installSpec.CloudInit, f.Name())
 	// Get the source of the installation if we are overriding it
-	if sourceImg != "" {
-		imgSource, err := v1.NewSrcFromURI(sourceImg)
-		if err != nil {
-			return err
-		}
-		installSpec.Active.Source = imgSource
-
-		// TODO: Why only setting active source above? What about size?
-		// TODO: These 2 blocks are identical, DRY them.
-
-		// size, err := GetSourceSize(cfg, imgSource)
-		// if err != nil {
-		// 	c.Logger.Warnf("Failed to infer size for images: %s", err.Error())
-		// }
-
-		// installSpec.Active.Source = imgSource
-		// installSpec.Passive.Source = imgSource
-		// installSpec.Recovery.Source = imgSource
-		// installSpec.Active.Size = uint(size)
-		// installSpec.Passive.Size = uint(size)
-		// installSpec.Recovery.Size = uint(size)
-	} else if c.Install.Image != "" {
+	if c.Install.Image != "" {
 		imgSource, err := v1.NewSrcFromURI(c.Install.Image)
 		if err != nil {
 			return err
@@ -352,3 +339,14 @@ func prepareConfiguration(ctx context.Context, source string) (string, error) {
 
 	return f.Name(), nil
 }
+
+func generateInstallConfForCLIArgs(source string) string {
+	if source == "" {
+		return ""
+	}
+
+	return fmt.Sprintf(`install:
+system:
+uri: %s
+`, source)
+}
diff --git a/internal/agent/interactive_install.go b/internal/agent/interactive_install.go
index 8b65b94..67c6fc7 100644
--- a/internal/agent/interactive_install.go
+++ b/internal/agent/interactive_install.go
@@ -282,9 +282,11 @@ func InteractiveInstall(debug, spawnShell bool, sourceImg string) error {
 		if err != nil {
 			fmt.Printf("could not write event cloud init: %s\n", err.Error())
 		}
-		// override cc with our new config object from the scan, so it's updated for the RunInstall function
-		// TODO: Alternative solution: pass a reader here (the new feature) and add the image source
-		cc, _ = config.Scan(collector.Directories(tmpdir), collector.MergeBootLine, collector.NoLogs)
+
+		cliConf := generateInstallConfForCLIArgs(sourceImg)
+		cc, _ = config.Scan(collector.Directories(tmpdir),
+			collector.Readers(strings.NewReader(cliConf)),
+			collector.MergeBootLine, collector.NoLogs)
 	}
 
 	pterm.Info.Println("Starting installation")
@@ -292,7 +294,7 @@ func InteractiveInstall(debug, spawnShell bool, sourceImg string) error {
 	ccString, _ := cc.String()
 	pterm.Info.Println(ccString)
 
-	err = RunInstall(cc, sourceImg)
+	err = RunInstall(cc)
 	if err != nil {
 		pterm.Error.Println(err.Error())
 	}
diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 7c0907a..e95eaf8 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -110,9 +110,9 @@ func determineUpgradeImage(version string) (*v1.ImageSource, error) {
 	return v1.NewSrcFromURI(fmt.Sprintf("%s:%s", registry, version))
 }
 
-// generateConfForCLIArgs creates a kairos configuration for `--source` and `--recovery`
+// generateUpgradeConfForCLIArgs creates a kairos configuration for `--source` and `--recovery`
 // command line arguments. It will be added to the rest of the configurations.
-func generateConfForCLIArgs(source string, upgradeRecovery bool) (string, error) {
+func generateUpgradeConfForCLIArgs(source string, upgradeRecovery bool) (string, error) {
 	upgrade := map[string](map[string]interface{}){
 		"upgrade": {},
 	}
@@ -198,7 +198,7 @@ func findLatestVersion(preReleases, force bool) (string, error) {
 }
 
 func generateUpgradeSpec(version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) (*v1.UpgradeSpec, *config.Config, error) {
-	cliConf, err := generateConfForCLIArgs(source, upgradeRecovery)
+	cliConf, err := generateUpgradeConfForCLIArgs(source, upgradeRecovery)
 	if err != nil {
 		return nil, nil, err
 	}

From 48a3f5704863beb70965d35b06751ee4b8e9415d Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 15:57:48 +0300
Subject: [PATCH 09/23] Rename var

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/install.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/internal/agent/install.go b/internal/agent/install.go
index 3c23108..f8831d4 100644
--- a/internal/agent/install.go
+++ b/internal/agent/install.go
@@ -57,14 +57,14 @@ func ManualInstall(c, sourceImg, device string, reboot, poweroff, strictValidati
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	ConfigSource, err := prepareConfiguration(ctx, c)
+	configSource, err := prepareConfiguration(ctx, c)
 	if err != nil {
 		return err
 	}
 
 	cliConf := generateInstallConfForCLIArgs(sourceImg)
 
-	cc, err := config.Scan(collector.Directories(ConfigSource),
+	cc, err := config.Scan(collector.Directories(configSource),
 		collector.Readers(strings.NewReader(cliConf)),
 		collector.MergeBootLine,
 		collector.StrictValidation(strictValidations), collector.NoLogs)

From 48c445756e906d19c0ddfa87e556b8daf97e3418 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 15:59:25 +0300
Subject: [PATCH 10/23] Fix conf identation

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/install.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/internal/agent/install.go b/internal/agent/install.go
index f8831d4..d707489 100644
--- a/internal/agent/install.go
+++ b/internal/agent/install.go
@@ -346,7 +346,7 @@ func generateInstallConfForCLIArgs(source string) string {
 	}
 
 	return fmt.Sprintf(`install:
-system:
-uri: %s
+  system:
+    uri: %s
 `, source)
 }

From f2ec19d53d5be299717dc7d11640ebcc0c6fc603 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 16:44:10 +0300
Subject: [PATCH 11/23] Fix syntax and don't hide error

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/install_test.go | 2 +-
 pkg/config/spec.go             | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/internal/agent/install_test.go b/internal/agent/install_test.go
index b5fdd79..98da9ac 100644
--- a/internal/agent/install_test.go
+++ b/internal/agent/install_test.go
@@ -197,7 +197,7 @@ var _ = Describe("RunInstall", func() {
 
 	It("runs the install", func() {
 		Skip("Not ready yet")
-		err = RunInstall(options, "") // TODO
+		err = RunInstall(options)
 		Expect(err).ToNot(HaveOccurred())
 	})
 })
diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index 261c723..d5170a2 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -199,7 +199,7 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 
 		squashedRec, err := hasSquashedRecovery(cfg, ep.Recovery)
 		if err != nil {
-			return nil, fmt.Errorf("failed checking for squashed recovery")
+			return nil, fmt.Errorf("failed checking for squashed recovery: %w", err)
 		}
 
 		if squashedRec {

From b68b07f278e2ceda654d6169c2e77d4fe0108b30 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 17:32:31 +0300
Subject: [PATCH 12/23] Remove unecessary type definition

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/upgrade.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index e95eaf8..b081092 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -52,7 +52,7 @@ func ListReleases(includePrereleases bool) semver.Collection {
 }
 
 func Upgrade(
-	version, source string, force, strictValidations bool, dirs []string, preReleases bool, upgradeRecovery bool) error {
+	version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) error {
 	bus.Manager.Initialize()
 
 	upgradeSpec, c, err := generateUpgradeSpec(version, source, force, strictValidations, dirs, preReleases, upgradeRecovery)

From a2d4df589b469fe7e1a14817bf950a63a3c11337 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Thu, 28 Sep 2023 17:34:52 +0300
Subject: [PATCH 13/23] Remove comments

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/config/spec.go | 27 ---------------------------
 1 file changed, 27 deletions(-)

diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index d5170a2..7c3acfa 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -265,33 +265,6 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 		State:      installState,
 	}
 
-	// source := viper.GetString("upgradeSource")
-	// recoveryUpgrade := viper.GetBool("upgradeRecovery")
-	// if source != "" {
-	// 	imgSource, err := v1.NewSrcFromURI(source)
-	// 	// TODO: Don't hide the error here!
-	// 	if err == nil {
-	// 		if recoveryUpgrade {
-	// 			spec.RecoveryUpgrade = recoveryUpgrade
-	// 			spec.Recovery.Source = imgSource
-	// 		} else {
-	// 			spec.Active.Source = imgSource
-	// 		}
-	// 		size, err := GetSourceSize(cfg, imgSource)
-	// 		if err != nil {
-	// 			cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
-	// 		} else {
-	// 			cfg.Logger.Infof("Setting image size to %dMb", size)
-	// 			// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
-	// 			if recoveryUpgrade {
-	// 				spec.Recovery.Size = uint(size)
-	// 			} else {
-	// 				spec.Active.Size = uint(size)
-	// 			}
-	// 		}
-	// 	}
-	//}
-
 	return spec, nil
 }
 

From f477776b579e5c25915d9c74c0579f5056ecc4f9 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Fri, 29 Sep 2023 10:12:12 +0300
Subject: [PATCH 14/23] Bring back the size calculation in Upgrade

it was removed by mistake

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/config/spec.go | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index 7c3acfa..e7b6b26 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -265,6 +265,11 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 		State:      installState,
 	}
 
+	err = setSourceSize(cfg, spec)
+	if err != nil {
+		cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+	}
+
 	return spec, nil
 }
 
@@ -693,3 +698,27 @@ func isMounted(config *Config, part *v1.Partition) (bool, error) {
 	}
 	return !notMnt, nil
 }
+
+func setSourceSize(cfg *Config, spec *v1.UpgradeSpec) error {
+	var size int64
+	var err error
+
+	if spec.RecoveryUpgrade {
+		size, err = GetSourceSize(cfg, spec.Recovery.Source)
+	} else {
+		size, err = GetSourceSize(cfg, spec.Active.Source)
+	}
+	if err != nil {
+		return err
+	}
+
+	cfg.Logger.Infof("Setting image size to %dMb", size)
+	// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
+	if spec.RecoveryUpgrade {
+		spec.Recovery.Size = uint(size)
+	} else {
+		spec.Active.Size = uint(size)
+	}
+
+	return nil
+}

From b13f1a9cc8859bfa5b60120649b703624b10b619 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Fri, 29 Sep 2023 11:38:29 +0300
Subject: [PATCH 15/23] Un-export fields again

because they are unmarshalled with a custom unmarshaller so they don't
need to be exported

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/types/v1/common.go | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/pkg/types/v1/common.go b/pkg/types/v1/common.go
index 5f75a87..eeb9302 100644
--- a/pkg/types/v1/common.go
+++ b/pkg/types/v1/common.go
@@ -35,31 +35,31 @@ const (
 
 // ImageSource represents the source from where an image is created for easy identification
 type ImageSource struct {
-	Source  string
-	SrcType string
+	source  string
+	srcType string
 }
 
 func (i ImageSource) Value() string {
-	return i.Source
+	return i.source
 }
 
 func (i ImageSource) IsDocker() bool {
-	return i.SrcType == oci
+	return i.srcType == oci
 }
 
 func (i ImageSource) IsDir() bool {
-	return i.SrcType == dir
+	return i.srcType == dir
 }
 
 func (i ImageSource) IsFile() bool {
-	return i.SrcType == file
+	return i.srcType == file
 }
 
 func (i ImageSource) IsEmpty() bool {
-	if i.SrcType == "" {
+	if i.srcType == "" {
 		return true
 	}
-	if i.Source == "" {
+	if i.source == "" {
 		return true
 	}
 	return false
@@ -69,7 +69,7 @@ func (i ImageSource) String() string {
 	if i.IsEmpty() {
 		return ""
 	}
-	return fmt.Sprintf("%s://%s", i.SrcType, i.Source)
+	return fmt.Sprintf("%s://%s", i.srcType, i.source)
 }
 
 func (i ImageSource) MarshalYAML() (interface{}, error) {
@@ -103,11 +103,11 @@ func (i *ImageSource) updateFromURI(uri string) error {
 	case oci, docker:
 		return i.parseImageReference(value)
 	case dir:
-		i.SrcType = dir
-		i.Source = value
+		i.srcType = dir
+		i.source = value
 	case file:
-		i.SrcType = file
-		i.Source = value
+		i.srcType = file
+		i.source = value
 	default:
 		return i.parseImageReference(uri)
 	}
@@ -121,8 +121,8 @@ func (i *ImageSource) parseImageReference(ref string) error {
 	} else if reference.IsNameOnly(n) {
 		ref += ":latest"
 	}
-	i.SrcType = oci
-	i.Source = ref
+	i.srcType = oci
+	i.source = ref
 	return nil
 }
 
@@ -137,13 +137,13 @@ func NewEmptySrc() *ImageSource {
 }
 
 func NewDockerSrc(src string) *ImageSource {
-	return &ImageSource{Source: src, SrcType: oci}
+	return &ImageSource{source: src, srcType: oci}
 }
 
 func NewFileSrc(src string) *ImageSource {
-	return &ImageSource{Source: src, SrcType: file}
+	return &ImageSource{source: src, srcType: file}
 }
 
 func NewDirSrc(src string) *ImageSource {
-	return &ImageSource{Source: src, SrcType: dir}
+	return &ImageSource{source: src, srcType: dir}
 }

From 5b9d043aa0986ea31a3765932d85a5412fda49ac Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Fri, 29 Sep 2023 11:44:53 +0300
Subject: [PATCH 16/23] Remove non-working test

because it needs to stub the world and currently there is no way of
doing so (we've even skipped the install_test.go for similar reasons).
---
 internal/agent/upgrade_test.go | 47 ----------------------------------
 1 file changed, 47 deletions(-)
 delete mode 100644 internal/agent/upgrade_test.go

diff --git a/internal/agent/upgrade_test.go b/internal/agent/upgrade_test.go
deleted file mode 100644
index ffa4da8..0000000
--- a/internal/agent/upgrade_test.go
+++ /dev/null
@@ -1,47 +0,0 @@
-package agent
-
-import (
-	"fmt"
-	"os"
-	"path"
-
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
-)
-
-var _ = Describe("generateUpgradeSpec", func() {
-	When("there are command line arguments", func() {
-		var configDir string
-		var upgradeRecovery bool
-
-		BeforeEach(func() {
-			upgradeRecovery = false
-			configDir, err := os.MkdirTemp("", "upgrade-test")
-			Expect(err).ToNot(HaveOccurred())
-
-			config := fmt.Sprintf(`upgrade:
-  recovery: %t
-  recovery-system:
-    uri: oci://image-in-conf
-`, upgradeRecovery)
-
-			configFilePath := path.Join(configDir, "config.yaml")
-			err = os.WriteFile(configFilePath, []byte(config), os.ModePerm)
-			Expect(err).ToNot(HaveOccurred())
-		})
-
-		AfterEach(func() {
-			os.RemoveAll(configDir)
-		})
-
-		It("overrides kairos config yaml values", func() {
-			spec, _, err := generateUpgradeSpec("", "oci:myimage", false, false,
-				[]string{}, false, !upgradeRecovery)
-
-			Expect(err).ToNot(HaveOccurred())
-			Expect(spec.Active.Source.String()).To(Equal("oci://myimage:latest"))
-			Expect(spec.Recovery.Source.String()).To(Equal("oci://myimage:latest"))
-			Expect(spec.RecoveryUpgrade).To(Equal(!upgradeRecovery))
-		})
-	})
-})

From b675730e8c8f7200ce9fd9a9305116cf7273169b Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Fri, 29 Sep 2023 12:50:34 +0300
Subject: [PATCH 17/23] Rename variables to better reflect what they are (PR
 comment)

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/install.go             | 14 +++++++-------
 internal/agent/interactive_install.go |  6 +++---
 internal/agent/upgrade.go             |  4 ++--
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/internal/agent/install.go b/internal/agent/install.go
index d707489..0420a23 100644
--- a/internal/agent/install.go
+++ b/internal/agent/install.go
@@ -53,7 +53,7 @@ func displayInfo(agentConfig *Config) {
 	}
 }
 
-func ManualInstall(c, sourceImg, device string, reboot, poweroff, strictValidations bool) error {
+func ManualInstall(c, sourceImgURL, device string, reboot, poweroff, strictValidations bool) error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -62,7 +62,7 @@ func ManualInstall(c, sourceImg, device string, reboot, poweroff, strictValidati
 		return err
 	}
 
-	cliConf := generateInstallConfForCLIArgs(sourceImg)
+	cliConf := generateInstallConfForCLIArgs(sourceImgURL)
 
 	cc, err := config.Scan(collector.Directories(configSource),
 		collector.Readers(strings.NewReader(cliConf)),
@@ -88,7 +88,7 @@ func ManualInstall(c, sourceImg, device string, reboot, poweroff, strictValidati
 	return RunInstall(cc)
 }
 
-func Install(sourceImg string, dir ...string) error {
+func Install(sourceImgURL string, dir ...string) error {
 	var cc *config.Config
 	var err error
 
@@ -121,7 +121,7 @@ func Install(sourceImg string, dir ...string) error {
 
 	ensureDataSourceReady()
 
-	cliConf := generateInstallConfForCLIArgs(sourceImg)
+	cliConf := generateInstallConfForCLIArgs(sourceImgURL)
 
 	// Reads config, and if present and offline is defined, runs the installation
 	cc, err = config.Scan(collector.Directories(dir...),
@@ -340,13 +340,13 @@ func prepareConfiguration(ctx context.Context, source string) (string, error) {
 	return f.Name(), nil
 }
 
-func generateInstallConfForCLIArgs(source string) string {
-	if source == "" {
+func generateInstallConfForCLIArgs(sourceImageURL string) string {
+	if sourceImageURL == "" {
 		return ""
 	}
 
 	return fmt.Sprintf(`install:
   system:
     uri: %s
-`, source)
+`, sourceImageURL)
 }
diff --git a/internal/agent/interactive_install.go b/internal/agent/interactive_install.go
index 67c6fc7..e24ca6e 100644
--- a/internal/agent/interactive_install.go
+++ b/internal/agent/interactive_install.go
@@ -129,7 +129,7 @@ func detectDevice() string {
 	return preferedDevice
 }
 
-func InteractiveInstall(debug, spawnShell bool, sourceImg string) error {
+func InteractiveInstall(debug, spawnShell bool, sourceImgURL string) error {
 	var sshUsers []string
 	bus.Manager.Initialize()
 
@@ -229,7 +229,7 @@ func InteractiveInstall(debug, spawnShell bool, sourceImg string) error {
 	}
 
 	if !isYes(allGood) {
-		return InteractiveInstall(debug, spawnShell, sourceImg)
+		return InteractiveInstall(debug, spawnShell, sourceImgURL)
 	}
 
 	usersToSet := map[string]schema.User{}
@@ -283,7 +283,7 @@ func InteractiveInstall(debug, spawnShell bool, sourceImg string) error {
 			fmt.Printf("could not write event cloud init: %s\n", err.Error())
 		}
 
-		cliConf := generateInstallConfForCLIArgs(sourceImg)
+		cliConf := generateInstallConfForCLIArgs(sourceImgURL)
 		cc, _ = config.Scan(collector.Directories(tmpdir),
 			collector.Readers(strings.NewReader(cliConf)),
 			collector.MergeBootLine, collector.NoLogs)
diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index b081092..3c79cb7 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -197,8 +197,8 @@ func findLatestVersion(preReleases, force bool) (string, error) {
 	return version, nil
 }
 
-func generateUpgradeSpec(version, source string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) (*v1.UpgradeSpec, *config.Config, error) {
-	cliConf, err := generateUpgradeConfForCLIArgs(source, upgradeRecovery)
+func generateUpgradeSpec(version, sourceImageURL string, force, strictValidations bool, dirs []string, preReleases, upgradeRecovery bool) (*v1.UpgradeSpec, *config.Config, error) {
+	cliConf, err := generateUpgradeConfForCLIArgs(sourceImageURL, upgradeRecovery)
 	if err != nil {
 		return nil, nil, err
 	}

From 85196fbc59be97b61ea0c4e3000698cba1760521 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Fri, 29 Sep 2023 15:49:36 +0300
Subject: [PATCH 18/23] Calculate sources sizes after unmarshalling the full
 spec

otherwise sources defined in the kairos config won't be there when we
calculate the size

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/config/spec.go | 99 ++++++++++++++++++++++++++++------------------
 1 file changed, 61 insertions(+), 38 deletions(-)

diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index e7b6b26..6d7f49f 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -36,7 +36,7 @@ import (
 )
 
 // NewInstallSpec returns an InstallSpec struct all based on defaults and basic host checks (e.g. EFI vs BIOS)
-func NewInstallSpec(cfg *Config) *v1.InstallSpec {
+func NewInstallSpec(cfg *Config) (*v1.InstallSpec, error) {
 	var firmware string
 	var recoveryImg, activeImg, passiveImg v1.Image
 
@@ -98,6 +98,11 @@ func NewInstallSpec(cfg *Config) *v1.InstallSpec {
 		Passive:   passiveImg,
 	}
 
+	err := unmarshallFullSpec(cfg, "install", spec)
+	if err != nil {
+		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
+	}
+
 	// Get the actual source size to calculate the image size and partitions size
 	size, err := GetSourceSize(cfg, spec.Active.Source)
 	if err != nil {
@@ -112,7 +117,7 @@ func NewInstallSpec(cfg *Config) *v1.InstallSpec {
 	// Calculate the partitions afterwards so they use the image sizes for the final partition sizes
 	spec.Partitions = NewInstallElementalPartitions(spec)
 
-	return spec
+	return spec, nil
 }
 
 func NewInstallElementalPartitions(spec *v1.InstallSpec) v1.ElementalPartitions {
@@ -265,14 +270,42 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 		State:      installState,
 	}
 
-	err = setSourceSize(cfg, spec)
+	err = unmarshallFullSpec(cfg, "upgrade", spec)
 	if err != nil {
-		cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
 	}
 
+	setUpgradeSourceSize(cfg, spec)
+
 	return spec, nil
 }
 
+func setUpgradeSourceSize(cfg *Config, spec *v1.UpgradeSpec) error {
+	var size int64
+	var err error
+
+	var targetSpec *v1.Image
+	if spec.RecoveryUpgrade {
+		targetSpec = &(spec.Recovery)
+	} else {
+		targetSpec = &(spec.Active)
+	}
+
+	if targetSpec.Source.IsEmpty() {
+		return nil
+	}
+
+	size, err = GetSourceSize(cfg, targetSpec.Source)
+	if err != nil {
+		return err
+	}
+
+	cfg.Logger.Infof("Setting image size to %dMb", size)
+	targetSpec.Size = uint(size)
+
+	return nil
+}
+
 // NewResetSpec returns a ResetSpec struct all based on defaults and current host state
 func NewResetSpec(cfg *Config) (*v1.ResetSpec, error) {
 	var imgSource *v1.ImageSource
@@ -396,6 +429,11 @@ func NewResetSpec(cfg *Config) (*v1.ResetSpec, error) {
 		State: installState,
 	}
 
+	err = unmarshallFullSpec(cfg, "reset", spec)
+	if err != nil {
+		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
+	}
+
 	// Get the actual source size to calculate the image size and partitions size
 	size, err := GetSourceSize(cfg, spec.Active.Source)
 	if err != nil {
@@ -527,7 +565,7 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 
 	switch spec {
 	case "install":
-		sp = NewInstallSpec(r)
+		sp, err = NewInstallSpec(r)
 	case "upgrade":
 		sp, err = NewUpgradeSpec(r)
 	case "reset":
@@ -539,23 +577,6 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 		return nil, fmt.Errorf("failed initializing spec: %v", err)
 	}
 
-	// Load the config into viper from the raw cloud config string
-	ccString, err := r.String()
-	if err != nil {
-		return nil, fmt.Errorf("failed initializing spec: %v", err)
-	}
-	viper.SetConfigType("yaml")
-	viper.ReadConfig(strings.NewReader(ccString))
-	vp := viper.Sub(spec)
-	if vp == nil {
-		vp = viper.New()
-	}
-
-	err = vp.Unmarshal(sp, setDecoder, decodeHook)
-	if err != nil {
-		r.Logger.Warnf("error unmarshalling %s Spec: %s", spec, err)
-	}
-
 	r.Logger.Debugf("Loaded %s spec: %s", spec, litter.Sdump(sp))
 	return sp, nil
 }
@@ -699,25 +720,27 @@ func isMounted(config *Config, part *v1.Partition) (bool, error) {
 	return !notMnt, nil
 }
 
-func setSourceSize(cfg *Config, spec *v1.UpgradeSpec) error {
-	var size int64
-	var err error
-
-	if spec.RecoveryUpgrade {
-		size, err = GetSourceSize(cfg, spec.Recovery.Source)
-	} else {
-		size, err = GetSourceSize(cfg, spec.Active.Source)
-	}
+func unmarshallFullSpec(r *Config, subkey string, sp v1.Spec) error {
+	// Load the config into viper from the raw cloud config string
+	ccString, err := r.String()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed initializing spec: %w", err)
+	}
+	viper.SetConfigType("yaml")
+	viper.ReadConfig(strings.NewReader(ccString))
+	vp := viper.Sub(subkey)
+	if vp == nil {
+		vp = viper.New()
 	}
 
-	cfg.Logger.Infof("Setting image size to %dMb", size)
-	// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
-	if spec.RecoveryUpgrade {
-		spec.Recovery.Size = uint(size)
-	} else {
-		spec.Active.Size = uint(size)
+	err = vp.Unmarshal(sp, setDecoder, decodeHook)
+	if err != nil {
+		return fmt.Errorf("error unmarshalling %s Spec: %w", subkey, err)
+	}
+
+	err = sp.Sanitize()
+	if err != nil {
+		return fmt.Errorf("sanitizing the % spec: %w", subkey, err)
 	}
 
 	return nil

From a10390e0a80b6c2d7da5ad01f904b8bc11358fa2 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Fri, 29 Sep 2023 17:11:51 +0300
Subject: [PATCH 19/23] DRY the definition of "--source" flag

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 main.go | 25 +++++++++----------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/main.go b/main.go
index 0649b75..56e809e 100644
--- a/main.go
+++ b/main.go
@@ -54,6 +54,11 @@ func ReleasesToOutput(rels semver.Collection, output string) []string {
 	}
 }
 
+var sourceFlag = 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`",
+}
+
 var cmds = []*cli.Command{
 	{
 		Name: "upgrade",
@@ -66,10 +71,7 @@ var cmds = []*cli.Command{
 				Name:  "image",
 				Usage: "[DEPRECATED] Specify a full image reference, e.g.: quay.io/some/image:tag",
 			},
-			&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`",
-			},
+			&sourceFlag,
 			&cli.BoolFlag{Name: "pre", Usage: "Include pre-releases (rc, beta, alpha)"},
 			&cli.BoolFlag{Name: "recovery", Usage: "Upgrade recovery"},
 		},
@@ -378,10 +380,7 @@ This command is meant to be used from the boot GRUB menu, but can be also starte
 			&cli.BoolFlag{
 				Name: "shell",
 			},
-			&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`",
-			},
+			&sourceFlag,
 		},
 		Usage: "Starts interactive installation",
 		Before: func(c *cli.Context) error {
@@ -413,10 +412,7 @@ This command is meant to be used from the boot GRUB menu, but can be also starte
 			&cli.BoolFlag{
 				Name: "reboot",
 			},
-			&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`",
-			},
+			&sourceFlag,
 		},
 		Before: func(c *cli.Context) error {
 			if err := validateSource(c.String("source")); err != nil {
@@ -456,10 +452,7 @@ This command is meant to be used from the boot GRUB menu, but can be started man
 			return checkRoot()
 		},
 		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`",
-			},
+			&sourceFlag,
 		},
 		Action: func(c *cli.Context) error {
 			source := c.String("source")

From d1e84c186cd77d8daac3ca971e20814899c4b5a2 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Mon, 2 Oct 2023 10:43:18 +0300
Subject: [PATCH 20/23] Fix minor syntax and linting issues

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/action/install_test.go      | 13 +++++-----
 pkg/config/spec_test.go         | 16 +++++++-----
 pkg/elemental/elemental_test.go | 46 +++++++++++++++++----------------
 3 files changed, 41 insertions(+), 34 deletions(-)

diff --git a/pkg/action/install_test.go b/pkg/action/install_test.go
index 69c9748..39bac2f 100644
--- a/pkg/action/install_test.go
+++ b/pkg/action/install_test.go
@@ -18,14 +18,14 @@ package action_test
 
 import (
 	"bytes"
-	"errors"
 	"fmt"
-	agentConfig "github.com/kairos-io/kairos-agent/v2/pkg/config"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
-	"github.com/kairos-io/kairos-sdk/collector"
 	"path/filepath"
 	"regexp"
 
+	agentConfig "github.com/kairos-io/kairos-agent/v2/pkg/config"
+	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+	"github.com/kairos-io/kairos-sdk/collector"
+
 	"github.com/jaypipes/ghw/pkg/block"
 	"github.com/kairos-io/kairos-agent/v2/pkg/action"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
@@ -111,7 +111,7 @@ var _ = Describe("Install action tests", func() {
 			runner.SideEffect = func(cmd string, args ...string) ([]byte, error) {
 				regexCmd := regexp.MustCompile(cmdFail)
 				if cmdFail != "" && regexCmd.MatchString(cmd) {
-					return []byte{}, errors.New(fmt.Sprintf("failed on %s", cmd))
+					return []byte{}, fmt.Errorf("failed on %s", cmd)
 				}
 				switch cmd {
 				case "parted":
@@ -152,7 +152,8 @@ var _ = Describe("Install action tests", func() {
 			err = fsutils.MkdirAll(fs, constants.IsoBaseTree, constants.DirPerm)
 			Expect(err).To(BeNil())
 
-			spec = agentConfig.NewInstallSpec(config)
+			spec, err = agentConfig.NewInstallSpec(config)
+			Expect(err).ToNot(HaveOccurred())
 			spec.Active.Size = 16
 
 			grubCfg := filepath.Join(spec.Active.MountPoint, constants.GrubConf)
diff --git a/pkg/config/spec_test.go b/pkg/config/spec_test.go
index 9717891..589fe90 100644
--- a/pkg/config/spec_test.go
+++ b/pkg/config/spec_test.go
@@ -18,11 +18,14 @@ package config_test
 
 import (
 	"fmt"
+	"os"
+	"path/filepath"
+
 	"github.com/jaypipes/ghw/pkg/block"
 	config "github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	v1mock "github.com/kairos-io/kairos-agent/v2/tests/mocks"
 	"github.com/kairos-io/kairos-sdk/collector"
 	. "github.com/onsi/ginkgo/v2"
@@ -31,8 +34,6 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/twpayne/go-vfs/vfst"
 	"k8s.io/mount-utils"
-	"os"
-	"path/filepath"
 )
 
 var _ = Describe("Types", Label("types", "config"), func() {
@@ -154,7 +155,8 @@ var _ = Describe("Types", Label("types", "config"), func() {
 				_, err = fs.Create(recoveryImgFile)
 				Expect(err).ShouldNot(HaveOccurred())
 
-				spec := config.NewInstallSpec(c)
+				spec, err := config.NewInstallSpec(c)
+				Expect(err).ToNot(HaveOccurred())
 				Expect(spec.Firmware).To(Equal(v1.EFI))
 				Expect(spec.Active.Source.Value()).To(Equal(constants.IsoBaseTree))
 				Expect(spec.Recovery.Source.Value()).To(Equal(recoveryImgFile))
@@ -175,7 +177,8 @@ var _ = Describe("Types", Label("types", "config"), func() {
 				_, err = fs.Create(constants.IsoBaseTree)
 				Expect(err).ShouldNot(HaveOccurred())
 
-				spec := config.NewInstallSpec(c)
+				spec, err := config.NewInstallSpec(c)
+				Expect(err).ToNot(HaveOccurred())
 				Expect(spec.Firmware).To(Equal(v1.BIOS))
 				Expect(spec.Active.Source.Value()).To(Equal(constants.IsoBaseTree))
 				Expect(spec.Recovery.Source.Value()).To(Equal(spec.Active.File))
@@ -190,7 +193,8 @@ var _ = Describe("Types", Label("types", "config"), func() {
 				Expect(spec.Partitions.BIOS).NotTo(BeNil())
 			})
 			It("sets installation defaults without being on installation media", Label("install"), func() {
-				spec := config.NewInstallSpec(c)
+				spec, err := config.NewInstallSpec(c)
+				Expect(err).ToNot(HaveOccurred())
 				Expect(spec.Firmware).To(Equal(v1.BIOS))
 				Expect(spec.Active.Source.IsEmpty()).To(BeTrue())
 				Expect(spec.Recovery.Source.Value()).To(Equal(spec.Active.File))
diff --git a/pkg/elemental/elemental_test.go b/pkg/elemental/elemental_test.go
index 5f86148..474365e 100644
--- a/pkg/elemental/elemental_test.go
+++ b/pkg/elemental/elemental_test.go
@@ -19,15 +19,15 @@ package elemental_test
 import (
 	"errors"
 	"fmt"
-	agentConfig "github.com/kairos-io/kairos-agent/v2/pkg/config"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	"os"
 	"path/filepath"
 	"testing"
 
+	agentConfig "github.com/kairos-io/kairos-agent/v2/pkg/config"
+	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+
 	"github.com/jaypipes/ghw/pkg/block"
 
-	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	cnst "github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	"github.com/kairos-io/kairos-agent/v2/pkg/elemental"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
@@ -340,13 +340,15 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 		var printOut string
 		var failPart bool
 		var install *v1.InstallSpec
+		var err error
 
 		BeforeEach(func() {
 			cInit = &v1mock.FakeCloudInitRunner{ExecStages: []string{}, Error: false}
 			config.CloudInitRunner = cInit
 			config.Install.Device = "/some/device"
 			el = elemental.NewElemental(config)
-			install = agentConfig.NewInstallSpec(config)
+			install, err = agentConfig.NewInstallSpec(config)
+			Expect(err).ToNot(HaveOccurred())
 			install.Target = "/some/device"
 
 			err := fsutils.MkdirAll(fs, "/some", cnst.DirPerm)
@@ -498,7 +500,7 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 			cmdFail = ""
 			el = elemental.NewElemental(config)
 			img = &v1.Image{
-				FS:         constants.LinuxImgFs,
+				FS:         cnst.LinuxImgFs,
 				Size:       16,
 				Source:     v1.NewDirSrc(sourceDir),
 				MountPoint: destDir,
@@ -524,7 +526,7 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 			Expect(el.DeployImage(img, false)).To(BeNil())
 		})
 		It("Deploys an squashfs image from a directory", func() {
-			img.FS = constants.SquashFs
+			img.FS = cnst.SquashFs
 			Expect(el.DeployImage(img, true)).To(BeNil())
 			Expect(runner.MatchMilestones([][]string{
 				{
@@ -569,7 +571,7 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 		})
 		It("Fails creating the squashfs filesystem", func() {
 			cmdFail = "mksquashfs"
-			img.FS = constants.SquashFs
+			img.FS = cnst.SquashFs
 			_, err := el.DeployImage(img, true)
 			Expect(err).NotTo(BeNil())
 			Expect(runner.MatchMilestones([][]string{
@@ -609,7 +611,7 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 			src := ""
 			dest := ""
 			runner.SideEffect = func(cmd string, args ...string) ([]byte, error) {
-				if cmd == constants.Rsync {
+				if cmd == cnst.Rsync {
 					rsyncCount += 1
 					src = args[len(args)-2]
 					dest = args[len(args)-1]
@@ -647,7 +649,7 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 			Expect(err).To(BeNil())
 			_, err = e.DumpSource(destFile, v1.NewFileSrc(sourceImg))
 			Expect(err).To(BeNil())
-			Expect(runner.IncludesCmds([][]string{{constants.Rsync}}))
+			Expect(runner.IncludesCmds([][]string{{cnst.Rsync}}))
 		})
 		It("Fails to copy, source file is not present", func() {
 			_, err := e.DumpSource("whatever", v1.NewFileSrc("/source.img"))
@@ -687,7 +689,7 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 		var relabelCmd []string
 		BeforeEach(func() {
 			// to mock the existance of setfiles command on non selinux hosts
-			err := fsutils.MkdirAll(fs, "/usr/sbin", constants.DirPerm)
+			err := fsutils.MkdirAll(fs, "/usr/sbin", cnst.DirPerm)
 			Expect(err).ShouldNot(HaveOccurred())
 			sbin, err := fs.RawPath("/usr/sbin")
 			Expect(err).ShouldNot(HaveOccurred())
@@ -700,23 +702,23 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 			Expect(err).ShouldNot(HaveOccurred())
 
 			// to mock SELinux policy files
-			policyFile = filepath.Join(constants.SELinuxTargetedPolicyPath, "policy.31")
-			err = fsutils.MkdirAll(fs, filepath.Dir(constants.SELinuxTargetedContextFile), constants.DirPerm)
+			policyFile = filepath.Join(cnst.SELinuxTargetedPolicyPath, "policy.31")
+			err = fsutils.MkdirAll(fs, filepath.Dir(cnst.SELinuxTargetedContextFile), cnst.DirPerm)
 			Expect(err).ShouldNot(HaveOccurred())
-			_, err = fs.Create(constants.SELinuxTargetedContextFile)
+			_, err = fs.Create(cnst.SELinuxTargetedContextFile)
 			Expect(err).ShouldNot(HaveOccurred())
-			err = fsutils.MkdirAll(fs, constants.SELinuxTargetedPolicyPath, constants.DirPerm)
+			err = fsutils.MkdirAll(fs, cnst.SELinuxTargetedPolicyPath, cnst.DirPerm)
 			Expect(err).ShouldNot(HaveOccurred())
 			_, err = fs.Create(policyFile)
 			Expect(err).ShouldNot(HaveOccurred())
 
 			relabelCmd = []string{
 				"setfiles", "-c", policyFile, "-e", "/dev", "-e", "/proc", "-e", "/sys",
-				"-F", constants.SELinuxTargetedContextFile, "/",
+				"-F", cnst.SELinuxTargetedContextFile, "/",
 			}
 		})
 		It("does nothing if the context file is not found", func() {
-			err := fs.Remove(constants.SELinuxTargetedContextFile)
+			err := fs.Remove(cnst.SELinuxTargetedContextFile)
 			Expect(err).ShouldNot(HaveOccurred())
 
 			c := elemental.NewElemental(config)
@@ -753,13 +755,13 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 			Expect(runner.CmdsMatch([][]string{relabelCmd})).To(BeNil())
 		})
 		It("relabels the given root-tree path", func() {
-			contextFile := filepath.Join("/root", constants.SELinuxTargetedContextFile)
-			err := fsutils.MkdirAll(fs, filepath.Dir(contextFile), constants.DirPerm)
+			contextFile := filepath.Join("/root", cnst.SELinuxTargetedContextFile)
+			err := fsutils.MkdirAll(fs, filepath.Dir(contextFile), cnst.DirPerm)
 			Expect(err).ShouldNot(HaveOccurred())
 			_, err = fs.Create(contextFile)
 			Expect(err).ShouldNot(HaveOccurred())
 			policyFile = filepath.Join("/root", policyFile)
-			err = fsutils.MkdirAll(fs, filepath.Join("/root", constants.SELinuxTargetedPolicyPath), constants.DirPerm)
+			err = fsutils.MkdirAll(fs, filepath.Join("/root", cnst.SELinuxTargetedPolicyPath), cnst.DirPerm)
 			Expect(err).ShouldNot(HaveOccurred())
 			_, err = fs.Create(policyFile)
 			Expect(err).ShouldNot(HaveOccurred())
@@ -888,9 +890,9 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 			Expect(runner.CmdsMatch([][]string{{"grub2-editenv"}})).NotTo(BeNil())
 		})
 		It("loads /etc/os-release on empty default entry", func() {
-			err := fsutils.MkdirAll(config.Fs, "/imgMountPoint/etc", constants.DirPerm)
+			err := fsutils.MkdirAll(config.Fs, "/imgMountPoint/etc", cnst.DirPerm)
 			Expect(err).ShouldNot(HaveOccurred())
-			err = config.Fs.WriteFile("/imgMountPoint/etc/os-release", []byte("GRUB_ENTRY_NAME=test"), constants.FilePerm)
+			err = config.Fs.WriteFile("/imgMountPoint/etc/os-release", []byte("GRUB_ENTRY_NAME=test"), cnst.FilePerm)
 			Expect(err).ShouldNot(HaveOccurred())
 
 			el := elemental.NewElemental(config)
@@ -909,7 +911,7 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 	})
 	Describe("FindKernelInitrd", Label("find"), func() {
 		BeforeEach(func() {
-			err := fsutils.MkdirAll(fs, "/path/boot", constants.DirPerm)
+			err := fsutils.MkdirAll(fs, "/path/boot", cnst.DirPerm)
 			Expect(err).ShouldNot(HaveOccurred())
 		})
 		It("finds kernel and initrd files", func() {

From 80f83ba676850a039eb2d6407be3f074be76ea52 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Mon, 2 Oct 2023 10:57:11 +0300
Subject: [PATCH 21/23] Run sanitization outside the spec building functions

To limit responsibility or those functions and to keep the tests passing
(because they assumed no sanitization)

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/action/upgrade_test.go |  5 +++--
 pkg/config/spec.go         | 10 +++++-----
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/pkg/action/upgrade_test.go b/pkg/action/upgrade_test.go
index 7ebc13e..a01fd71 100644
--- a/pkg/action/upgrade_test.go
+++ b/pkg/action/upgrade_test.go
@@ -19,10 +19,11 @@ package action_test
 import (
 	"bytes"
 	"fmt"
-	agentConfig "github.com/kairos-io/kairos-agent/v2/pkg/config"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	"path/filepath"
 
+	agentConfig "github.com/kairos-io/kairos-agent/v2/pkg/config"
+	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+
 	"github.com/jaypipes/ghw/pkg/block"
 	"github.com/kairos-io/kairos-agent/v2/pkg/action"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index 6d7f49f..b8af902 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -577,6 +577,11 @@ func ReadSpecFromCloudConfig(r *Config, spec string) (v1.Spec, error) {
 		return nil, fmt.Errorf("failed initializing spec: %v", err)
 	}
 
+	err = sp.Sanitize()
+	if err != nil {
+		return sp, fmt.Errorf("sanitizing the %s spec: %w", spec, err)
+	}
+
 	r.Logger.Debugf("Loaded %s spec: %s", spec, litter.Sdump(sp))
 	return sp, nil
 }
@@ -738,10 +743,5 @@ func unmarshallFullSpec(r *Config, subkey string, sp v1.Spec) error {
 		return fmt.Errorf("error unmarshalling %s Spec: %w", subkey, err)
 	}
 
-	err = sp.Sanitize()
-	if err != nil {
-		return fmt.Errorf("sanitizing the % spec: %w", subkey, err)
-	}
-
 	return nil
 }

From e0dfc79ed5b53c9213da09f7afe99c73955870f3 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Mon, 2 Oct 2023 12:28:33 +0300
Subject: [PATCH 22/23] Unmarshal yaml after auto size calculation and fix
 tests

because we want the user's preferences to be applied last, thus
overriding whatever calculations we may do automatically.

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/config/spec.go      | 25 +++++++++++++------------
 pkg/config/spec_test.go | 14 +++++++++++++-
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index b8af902..629faa1 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -60,6 +60,7 @@ func NewInstallSpec(cfg *Config) (*v1.InstallSpec, error) {
 	activeImg.File = filepath.Join(constants.StateDir, "cOS", constants.ActiveImgFile)
 	activeImg.FS = constants.LinuxImgFs
 	activeImg.MountPoint = constants.ActiveDir
+
 	if isoRootExists {
 		activeImg.Source = v1.NewDirSrc(constants.IsoBaseTree)
 	} else {
@@ -98,11 +99,6 @@ func NewInstallSpec(cfg *Config) (*v1.InstallSpec, error) {
 		Passive:   passiveImg,
 	}
 
-	err := unmarshallFullSpec(cfg, "install", spec)
-	if err != nil {
-		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
-	}
-
 	// Get the actual source size to calculate the image size and partitions size
 	size, err := GetSourceSize(cfg, spec.Active.Source)
 	if err != nil {
@@ -114,6 +110,11 @@ func NewInstallSpec(cfg *Config) (*v1.InstallSpec, error) {
 		spec.Recovery.Size = uint(size)
 	}
 
+	err = unmarshallFullSpec(cfg, "install", spec)
+	if err != nil {
+		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
+	}
+
 	// Calculate the partitions afterwards so they use the image sizes for the final partition sizes
 	spec.Partitions = NewInstallElementalPartitions(spec)
 
@@ -270,13 +271,13 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 		State:      installState,
 	}
 
+	setUpgradeSourceSize(cfg, spec)
+
 	err = unmarshallFullSpec(cfg, "upgrade", spec)
 	if err != nil {
 		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
 	}
 
-	setUpgradeSourceSize(cfg, spec)
-
 	return spec, nil
 }
 
@@ -429,11 +430,6 @@ func NewResetSpec(cfg *Config) (*v1.ResetSpec, error) {
 		State: installState,
 	}
 
-	err = unmarshallFullSpec(cfg, "reset", spec)
-	if err != nil {
-		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
-	}
-
 	// Get the actual source size to calculate the image size and partitions size
 	size, err := GetSourceSize(cfg, spec.Active.Source)
 	if err != nil {
@@ -444,6 +440,11 @@ func NewResetSpec(cfg *Config) (*v1.ResetSpec, error) {
 		spec.Passive.Size = uint(size)
 	}
 
+	err = unmarshallFullSpec(cfg, "reset", spec)
+	if err != nil {
+		return nil, fmt.Errorf("failed unmarshalling the full spec: %w", err)
+	}
+
 	return spec, nil
 }
 
diff --git a/pkg/config/spec_test.go b/pkg/config/spec_test.go
index 589fe90..c3af284 100644
--- a/pkg/config/spec_test.go
+++ b/pkg/config/spec_test.go
@@ -463,6 +463,8 @@ upgrade:
   recovery: true
   system:
     uri: docker:test/image:latest
+  recovery-system:
+    uri: docker:test/image:latest
 cloud-init-paths:
 - /what
 `)
@@ -502,6 +504,12 @@ cloud-init-paths:
 				ghwTest = v1mock.GhwMock{}
 				ghwTest.AddDisk(mainDisk)
 				ghwTest.CreateDevices()
+
+				fs, cleanup, err = vfst.NewTestFS(nil)
+				err = fsutils.MkdirAll(fs, filepath.Dir(constants.IsoBaseTree), constants.DirPerm)
+				Expect(err).ShouldNot(HaveOccurred())
+				_, err = fs.Create(constants.IsoBaseTree)
+				Expect(err).ShouldNot(HaveOccurred())
 			})
 
 			AfterEach(func() {
@@ -509,7 +517,11 @@ cloud-init-paths:
 				ghwTest.Clean()
 			})
 			It("Reads properly the cloud config for install", func() {
-				cfg, err := config.Scan(collector.Directories([]string{dir}...), collector.NoLogs)
+				cfg, err := config.Scan(collector.Directories([]string{dir}...),
+					collector.NoLogs,
+				)
+				cfg.Fs = fs
+
 				Expect(err).ToNot(HaveOccurred())
 				// Once we got the cfg override the fs to our test fs
 				cfg.Runner = runner

From 48fc6180cf430938915460688bb26867917343e8 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Mon, 2 Oct 2023 12:47:04 +0300
Subject: [PATCH 23/23] Return image from bus when one exists

it was accidentally returning "nil"

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 internal/agent/upgrade.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index 3c79cb7..5402cac 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -99,7 +99,7 @@ func determineUpgradeImage(version string) (*v1.ImageSource, error) {
 	}
 
 	if img != "" {
-		return nil, nil
+		return v1.NewSrcFromURI(img)
 	}
 
 	registry, err := utils.OSRelease("IMAGE_REPO")