From ba10f3e0f7cae1af3cc0503243aa8c98920c35e3 Mon Sep 17 00:00:00 2001
From: Federico Di Pierro <nierro92@gmail.com>
Date: Mon, 6 Nov 2023 15:40:18 +0100
Subject: [PATCH] new(pkg/driver,cmd/driver): allow users to override
 kernelrelease and kernelversion.

Moreover, added support for a couple of curl related flags in falco-driver-loader.

Signed-off-by: Federico Di Pierro <nierro92@gmail.com>
---
 cmd/driver/prepare/prepare.go     | 52 +++++++++++++----
 cmd/driver/select/select.go       |  4 +-
 pkg/driver/distro/amzn.go         |  7 +--
 pkg/driver/distro/bottlerocket.go |  7 +--
 pkg/driver/distro/debian.go       |  2 +-
 pkg/driver/distro/distro.go       | 95 +++++++++++++++++++++++++++----
 pkg/driver/distro/flatcar.go      |  7 +--
 pkg/driver/distro/generic.go      | 69 +---------------------
 pkg/driver/distro/minikube.go     |  3 +-
 pkg/driver/distro/talos.go        |  7 +--
 pkg/driver/distro/ubuntu.go       | 23 ++++----
 pkg/driver/kernel/kernel.go       | 24 ++++----
 12 files changed, 170 insertions(+), 130 deletions(-)

diff --git a/cmd/driver/prepare/prepare.go b/cmd/driver/prepare/prepare.go
index 5472b3b5..386b7d9b 100644
--- a/cmd/driver/prepare/prepare.go
+++ b/cmd/driver/prepare/prepare.go
@@ -16,8 +16,11 @@
 package driverprepare
 
 import (
+	"crypto/tls"
 	"errors"
 	"fmt"
+	"net/http"
+	"time"
 
 	"github.com/spf13/cobra"
 	"golang.org/x/net/context"
@@ -29,14 +32,22 @@ import (
 	"github.com/falcosecurity/falcoctl/pkg/output"
 )
 
+type driverDownloadOptions struct {
+	DriverRepos      []string
+	InsecureDownload bool
+	HttpTimeout      time.Duration
+}
+
 // DriverPrepareOptions contains the options for the driver prepare command.
 type driverPrepareOptions struct {
 	*options.Common
 	*options.Driver
-	Download      bool
-	Build         bool
-	DriverVersion string
-	DriverRepos   []string
+	Download            bool
+	Build               bool
+	DriverVersion       string
+	DriverKernelRelease string
+	DriverKernelVersion string
+	driverDownloadOptions
 }
 
 // NewDriverPrepareCmd returns the driver prepare command.
@@ -63,7 +74,11 @@ func NewDriverPrepareCmd(ctx context.Context, opt *options.Common) *cobra.Comman
 	cmd.Flags().BoolVar(&o.Download, "download", true, "Whether to enable download of drivers")
 	cmd.Flags().BoolVar(&o.Build, "build", true, "Whether to enable build of drivers")
 	cmd.Flags().StringVar(&o.DriverVersion, "driver-version", "", "Driver version to be built")
+	cmd.Flags().StringVar(&o.DriverKernelRelease, "driver-kernelrelease", "", "Specify the kernel release for which to download/build the driver in the same format used by 'uname -r' (e.g. '6.1.0-10-cloud-amd64')")
+	cmd.Flags().StringVar(&o.DriverKernelVersion, "driver-kernelversion", "", "Specify the kernel version for which to download/build the driver in the same format used by 'uname -v' (e.g. '#1 SMP PREEMPT_DYNAMIC Debian 6.1.38-2 (2023-07-27)')")
 	cmd.Flags().StringSliceVar(&o.DriverRepos, "driver-repo", []string{driverdistro.DefaultFalcoRepo}, "Specify different URL(s) where to look for prebuilt drivers")
+	cmd.Flags().BoolVar(&o.InsecureDownload, "http-insecure", false, "Whether you want to allow insecure downloads or not")
+	cmd.Flags().DurationVar(&o.HttpTimeout, "http-timeout", 60*time.Second, "Timeout for each http try")
 
 	if err := cmd.MarkFlagRequired("driver-version"); err != nil {
 		output.ExitOnErr(o.Printer, fmt.Errorf("unable to mark flag \"driver-version\" as required"))
@@ -71,6 +86,14 @@ func NewDriverPrepareCmd(ctx context.Context, opt *options.Common) *cobra.Comman
 	return cmd
 }
 
+func setDefaultHTTPClientOpts(downloadOptions driverDownloadOptions) {
+	// Skip insecure verify
+	if downloadOptions.InsecureDownload {
+		http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
+	}
+	http.DefaultClient.Timeout = downloadOptions.HttpTimeout
+}
+
 // RunDriverPrepare implements the driver prepare command.
 func (o *driverPrepareOptions) RunDriverPrepare(_ context.Context, _ []string) error {
 	driver, err := config.Driverer()
@@ -83,11 +106,18 @@ func (o *driverPrepareOptions) RunDriverPrepare(_ context.Context, _ []string) e
 		return err
 	}
 
+	if o.DriverKernelRelease != "" {
+		info.KernelRelease = o.DriverKernelRelease
+	}
+	if o.DriverKernelVersion != "" {
+		info.KernelVersion = driverkernel.FormatVersion(o.DriverKernelVersion)
+	}
+
 	o.Printer.Logger.Info("Running falcoctl driver prepare", o.Printer.Logger.Args(
 		"driver version", o.DriverVersion,
 		"driver type", driver.Type,
 		"driver name", o.Name,
-		"arch", info.Architecture,
+		"arch", driverdistro.GetArchitecture(),
 		"kernel release", info.KernelRelease,
 		"kernel version", info.KernelVersion))
 
@@ -96,7 +126,7 @@ func (o *driverPrepareOptions) RunDriverPrepare(_ context.Context, _ []string) e
 		return nil
 	}
 
-	d, err := driverdistro.DiscoverDistro(o.Printer, o.HostRoot)
+	d, err := driverdistro.DiscoverDistro(info, o.HostRoot)
 	if err != nil {
 		if errors.Is(err, driverdistro.ErrUnsupported) && o.Build {
 			o.Download = false
@@ -112,12 +142,12 @@ func (o *driverPrepareOptions) RunDriverPrepare(_ context.Context, _ []string) e
 		return err
 	}
 
+	setDefaultHTTPClientOpts(o.driverDownloadOptions)
+
 	var dest string
 	if o.Download {
-		// Fixup kernel information before attempting to download
-		downloadInfo := d.FixupKernel(info)
-
-		dest, err = d.Download(downloadInfo, o.Driver, driver.Type, o.DriverVersion, o.DriverRepos)
+		driverFileName := driverdistro.DriverToFilename(d, info, o.Driver, driver.Type)
+		dest, err = driverdistro.Download(d, o.Printer, driverFileName, o.DriverVersion, o.DriverRepos)
 		if err == nil {
 			// We don't care about errors at this stage
 			_ = driver.Type.Load(o.Printer, dest, false)
@@ -129,7 +159,7 @@ func (o *driverPrepareOptions) RunDriverPrepare(_ context.Context, _ []string) e
 	}
 
 	if o.Build {
-		dest, err = d.Build()
+		dest, err = driverdistro.Build(d, o.Printer)
 		if err == nil {
 			// We don't care about errors
 			_ = driver.Type.Load(o.Printer, dest, false)
diff --git a/cmd/driver/select/select.go b/cmd/driver/select/select.go
index a3a5f006..bd3b8644 100644
--- a/cmd/driver/select/select.go
+++ b/cmd/driver/select/select.go
@@ -79,11 +79,11 @@ func (o *driverSelectOptions) RunDriverSelect(_ context.Context, args []string)
 	}
 
 	o.Printer.Logger.Info("Running falcoctl driver select", o.Printer.Logger.Args(
-		"arch", info.Architecture,
+		"arch", driverdistro.GetArchitecture(),
 		"kernel release", info.KernelRelease,
 		"kernel version", info.KernelVersion))
 
-	d, err := driverdistro.DiscoverDistro(o.Printer, o.HostRoot)
+	d, err := driverdistro.DiscoverDistro(info, o.HostRoot)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/driver/distro/amzn.go b/pkg/driver/distro/amzn.go
index 69a268d7..f0d80dd3 100644
--- a/pkg/driver/distro/amzn.go
+++ b/pkg/driver/distro/amzn.go
@@ -17,10 +17,9 @@ package driverdistro
 
 import (
 	"fmt"
+	driverkernel "github.com/falcosecurity/falcoctl/pkg/driver/kernel"
 
 	"gopkg.in/ini.v1"
-
-	"github.com/falcosecurity/falcoctl/pkg/output"
 )
 
 func init() {
@@ -31,7 +30,7 @@ type amzn struct {
 	*generic
 }
 
-func (a *amzn) init(printer *output.Printer, _ string, cfg *ini.File) error {
+func (a *amzn) init(i driverkernel.Info, _ string, cfg *ini.File) error {
 	idKey := cfg.Section("").Key("VERSION_ID")
 	if idKey == nil {
 		// OS-release without `VERSION_ID` (can it happen?)
@@ -49,5 +48,5 @@ func (a *amzn) init(printer *output.Printer, _ string, cfg *ini.File) error {
 	default:
 		newID = "amazonlinux"
 	}
-	return a.generic.init(printer, newID, cfg)
+	return a.generic.init(i, newID, cfg)
 }
diff --git a/pkg/driver/distro/bottlerocket.go b/pkg/driver/distro/bottlerocket.go
index b1427c31..1bae13ef 100644
--- a/pkg/driver/distro/bottlerocket.go
+++ b/pkg/driver/distro/bottlerocket.go
@@ -22,7 +22,6 @@ import (
 	"gopkg.in/ini.v1"
 
 	driverkernel "github.com/falcosecurity/falcoctl/pkg/driver/kernel"
-	"github.com/falcosecurity/falcoctl/pkg/output"
 )
 
 func init() {
@@ -35,7 +34,7 @@ type bottlerocket struct {
 	versionID string
 }
 
-func (b *bottlerocket) init(printer *output.Printer, id string, cfg *ini.File) error {
+func (b *bottlerocket) init(i driverkernel.Info, id string, cfg *ini.File) error {
 	idKey := cfg.Section("").Key("VERSION_ID")
 	if idKey == nil {
 		// OS-release without `VERSION_ID` (can it happen?)
@@ -49,10 +48,10 @@ func (b *bottlerocket) init(printer *output.Printer, id string, cfg *ini.File) e
 	}
 	b.variantID = strings.Split(idKey.String(), "-")[0]
 
-	return b.generic.init(printer, id, cfg)
+	return b.generic.init(i, id, cfg)
 }
 
-func (b *bottlerocket) FixupKernel(i driverkernel.Info) driverkernel.Info {
+func (b *bottlerocket) fixupKernel(i driverkernel.Info) driverkernel.Info {
 	i.KernelVersion = fmt.Sprintf("1_%s-%s", b.versionID, b.variantID)
 	return i
 }
diff --git a/pkg/driver/distro/debian.go b/pkg/driver/distro/debian.go
index be869f9e..c48c4187 100644
--- a/pkg/driver/distro/debian.go
+++ b/pkg/driver/distro/debian.go
@@ -43,7 +43,7 @@ func (d *debian) check(hostRoot string) bool {
 	return false
 }
 
-func (d *debian) FixupKernel(i driverkernel.Info) driverkernel.Info {
+func (d *debian) fixupKernel(i driverkernel.Info) driverkernel.Info {
 	// Workaround: debian kernelreleases might now be actual kernel running;
 	// instead, they might be the Debian kernel package
 	// providing the compatible kernel ABI
diff --git a/pkg/driver/distro/distro.go b/pkg/driver/distro/distro.go
index 9d4068af..4e01504d 100644
--- a/pkg/driver/distro/distro.go
+++ b/pkg/driver/distro/distro.go
@@ -17,13 +17,18 @@ package driverdistro
 
 import (
 	"fmt"
+	"github.com/docker/docker/pkg/homedir"
+	"github.com/falcosecurity/falcoctl/pkg/options"
+	"io"
+	"net/http"
+	"os"
+	"runtime"
 	"strings"
 
 	"gopkg.in/ini.v1"
 
 	driverkernel "github.com/falcosecurity/falcoctl/pkg/driver/kernel"
 	drivertype "github.com/falcosecurity/falcoctl/pkg/driver/type"
-	"github.com/falcosecurity/falcoctl/pkg/options"
 	"github.com/falcosecurity/falcoctl/pkg/output"
 )
 
@@ -40,11 +45,9 @@ var ErrUnsupported = fmt.Errorf("failed to determine distro")
 // Distro is the common interface used by distro-specific implementations.
 // Most of the distro-specific only partially override the default `generic` implementation.
 type Distro interface {
-	init(printer *output.Printer, id string, cfg *ini.File) error // private
+	init(i driverkernel.Info, id string, cfg *ini.File) error // private
 	GetTargetID(i driverkernel.Info) string
-	FixupKernel(i driverkernel.Info) driverkernel.Info
-	Download(i driverkernel.Info, opts *options.Driver, driverType drivertype.DriverType, driverVer string, repos []string) (string, error)
-	Build() (string, error)
+	fixupKernel(i driverkernel.Info) driverkernel.Info // private
 	PreferredDriver(i driverkernel.Info) drivertype.DriverType
 }
 
@@ -54,8 +57,8 @@ type checker interface {
 
 // DiscoverDistro tries to fetch the correct Distro by looking at /etc/os-release or
 // by cycling on all supported distros and checking them one by one.
-func DiscoverDistro(printer *output.Printer, hostRoot string) (Distro, error) {
-	distro, err := getOSReleaseDistro(printer, hostRoot)
+func DiscoverDistro(i driverkernel.Info, hostRoot string) (Distro, error) {
+	distro, err := getOSReleaseDistro(i, hostRoot)
 	if err == nil {
 		return distro, nil
 	}
@@ -64,20 +67,20 @@ func DiscoverDistro(printer *output.Printer, hostRoot string) (Distro, error) {
 	for id, d := range distros {
 		dd, ok := d.(checker)
 		if ok && dd.check(hostRoot) {
-			err = d.init(printer, id, nil)
+			err = d.init(i, id, nil)
 			return d, err
 		}
 	}
 
 	// Return a generic distro to try the build
 	distro = &generic{}
-	if err = distro.init(printer, "undetermined", nil); err != nil {
+	if err = distro.init(i, "undetermined", nil); err != nil {
 		return nil, err
 	}
 	return distro, ErrUnsupported
 }
 
-func getOSReleaseDistro(printer *output.Printer, hostRoot string) (Distro, error) {
+func getOSReleaseDistro(i driverkernel.Info, hostRoot string) (Distro, error) {
 	cfg, err := ini.Load(hostRoot + "/etc/os-release")
 	if err != nil {
 		return nil, err
@@ -100,8 +103,78 @@ func getOSReleaseDistro(printer *output.Printer, hostRoot string) (Distro, error
 	if !exist {
 		distro = &generic{}
 	}
-	if err = distro.init(printer, id, cfg); err != nil {
+	if err = distro.init(i, id, cfg); err != nil {
 		return nil, err
 	}
 	return distro, nil
 }
+
+func toURL(repo, driverVer, fileName string) string {
+	return fmt.Sprintf("%s/%s/%s/%s", repo, driverVer, GetArchitecture(), fileName)
+}
+
+func toLocalPath(driverVer, fileName string) string {
+	return fmt.Sprintf("%s/.falco/%s/%s/%s", homedir.Get(), driverVer, GetArchitecture(), fileName)
+}
+
+// DriverToFilename returns the filename for a given distro's driver.
+func DriverToFilename(d Distro, i driverkernel.Info, opts *options.Driver, driverType drivertype.DriverType) string {
+	// Fixup kernel information before attempting to download
+	fixedInfo := d.fixupKernel(i)
+	return fmt.Sprintf("%s_%s_%s_%s%s", opts.Name, d.GetTargetID(i), fixedInfo.KernelRelease, fixedInfo.KernelVersion, driverType.Extension())
+}
+
+// GetArchitecture returns current arch as non-deb name (ie: x86_64/aarch64)
+func GetArchitecture() string {
+	switch runtime.GOARCH {
+	case "amd64":
+		return "x86_64"
+	case "arm64":
+		return "aarch64"
+	default:
+		return "uknown"
+	}
+}
+
+func Build(_ Distro, _ *output.Printer) (string, error) {
+	// TODO compile
+	// flatcar specific: flatcar_relocate_tools
+	// cos specific...
+	// minikube specific...
+	return "", fmt.Errorf("build unimplemented")
+}
+
+func Download(_ Distro, printer *output.Printer, driverFileName, driverVer string, repos []string) (string, error) {
+	// Skip if existent
+	destination := toLocalPath(driverVer, driverFileName)
+	f, err := os.Open(destination) //nolint:gosec // false positive
+	if err == nil {
+		_ = f.Close()
+		printer.Logger.Info("Skipping download, driver already present.", printer.Logger.Args("path", destination))
+		return destination, nil
+	}
+
+	// Try to download from any specified repository,
+	// stopping at first successful http GET.
+	for _, repo := range repos {
+		url := toURL(repo, driverVer, driverFileName)
+		printer.Logger.Info("Trying to download a driver", printer.Logger.Args("url", url))
+
+		resp, err := http.Get(url) //nolint:gosec // false positive
+		if err != nil || resp.StatusCode != 200 {
+			continue
+		}
+
+		out, err := os.Create(destination) //nolint:gosec // false positive
+		if err != nil {
+			return destination, err
+		}
+		_, err = io.Copy(out, resp.Body)
+		if err != nil {
+			return destination, err
+		}
+		_ = resp.Body.Close()
+		return destination, nil
+	}
+	return destination, fmt.Errorf("unable to find a prebuilt driver")
+}
diff --git a/pkg/driver/distro/flatcar.go b/pkg/driver/distro/flatcar.go
index 3e13327b..43e5c5eb 100644
--- a/pkg/driver/distro/flatcar.go
+++ b/pkg/driver/distro/flatcar.go
@@ -21,7 +21,6 @@ import (
 	"gopkg.in/ini.v1"
 
 	driverkernel "github.com/falcosecurity/falcoctl/pkg/driver/kernel"
-	"github.com/falcosecurity/falcoctl/pkg/output"
 )
 
 func init() {
@@ -33,17 +32,17 @@ type flatcar struct {
 	versionID string
 }
 
-func (f *flatcar) init(printer *output.Printer, id string, cfg *ini.File) error {
+func (f *flatcar) init(i driverkernel.Info, id string, cfg *ini.File) error {
 	idKey := cfg.Section("").Key("VERSION_ID")
 	if idKey == nil {
 		// OS-release without `VERSION_ID` (can it happen?)
 		return fmt.Errorf("no VERSION_ID present for flatcar")
 	}
 	f.versionID = idKey.String()
-	return f.generic.init(printer, id, cfg)
+	return f.generic.init(i, id, cfg)
 }
 
-func (f *flatcar) FixupKernel(i driverkernel.Info) driverkernel.Info {
+func (f *flatcar) fixupKernel(i driverkernel.Info) driverkernel.Info {
 	i.KernelRelease = f.versionID
 	return i
 }
diff --git a/pkg/driver/distro/generic.go b/pkg/driver/distro/generic.go
index b29560d9..624705cd 100644
--- a/pkg/driver/distro/generic.go
+++ b/pkg/driver/distro/generic.go
@@ -16,28 +16,17 @@
 package driverdistro
 
 import (
-	"fmt"
-	"io"
-	"net/http"
-	"os"
-
-	"github.com/docker/docker/pkg/homedir"
-	"gopkg.in/ini.v1"
-
 	driverkernel "github.com/falcosecurity/falcoctl/pkg/driver/kernel"
 	drivertype "github.com/falcosecurity/falcoctl/pkg/driver/type"
-	"github.com/falcosecurity/falcoctl/pkg/options"
-	"github.com/falcosecurity/falcoctl/pkg/output"
+	"gopkg.in/ini.v1"
 )
 
 type generic struct {
 	targetID string
-	printer  *output.Printer
 }
 
-func (g *generic) init(printer *output.Printer, id string, _ *ini.File) error {
+func (g *generic) init(_ driverkernel.Info, id string, _ *ini.File) error {
 	g.targetID = id
-	g.printer = printer
 	return nil
 }
 
@@ -45,63 +34,11 @@ func (g *generic) GetTargetID(_ driverkernel.Info) string {
 	return g.targetID
 }
 
-func (g *generic) FixupKernel(i driverkernel.Info) driverkernel.Info {
+func (g *generic) fixupKernel(i driverkernel.Info) driverkernel.Info {
 	return i
 }
 
-func (g *generic) Build() (string, error) {
-	// TODO compile
-	// flatcar specific: flatcar_relocate_tools
-	// cos specific...
-	// minikube specific...
-	return "", fmt.Errorf("build unimplemented")
-}
-
-func (g *generic) Download(i driverkernel.Info, opts *options.Driver, driverType drivertype.DriverType, driverVer string, repos []string) (string, error) {
-	// Skip if existent
-	driverFileName := g.toFilename(i, opts, driverType)
-	destination := fmt.Sprintf("%s/.falco/%s/%s/%s", homedir.Get(), driverVer, i.Architecture, driverFileName)
-	f, err := os.Open(destination) //nolint:gosec // false positive
-	if err == nil {
-		_ = f.Close()
-		g.printer.Logger.Info("Skipping download, driver already present.", g.printer.Logger.Args("path", destination))
-		return destination, nil
-	}
-
-	// Try to download from any specified repository,
-	// stopping at first successful http GET.
-	for _, repo := range repos {
-		url := g.toURL(i.Architecture, repo, driverVer, driverFileName)
-		g.printer.Logger.Info("Trying to download a driver", g.printer.Logger.Args("url", url))
-
-		resp, err := http.Get(url) //nolint:gosec
-		if err != nil || resp.StatusCode != 200 {
-			continue
-		}
-
-		out, err := os.Create(destination) //nolint:gosec // false positive
-		if err != nil {
-			return destination, err
-		}
-		_, err = io.Copy(out, resp.Body)
-		if err != nil {
-			return destination, err
-		}
-		_ = resp.Body.Close()
-		return destination, nil
-	}
-	return destination, fmt.Errorf("unable to find a prebuilt %s driver", opts.Name)
-}
-
 func (g *generic) PreferredDriver(_ driverkernel.Info) drivertype.DriverType {
 	// We don't support automagic driver selection logic at this stage
 	return nil
 }
-
-func (g *generic) toFilename(i driverkernel.Info, opts *options.Driver, driverType drivertype.DriverType) string {
-	return fmt.Sprintf("%s_%s_%s_%s%s", opts.Name, g.GetTargetID(i), i.KernelRelease, i.KernelVersion, driverType.Extension())
-}
-
-func (g *generic) toURL(arch, repo, driverVer, fileName string) string {
-	return fmt.Sprintf("%s/%s/%s/%s", repo, driverVer, arch, fileName)
-}
diff --git a/pkg/driver/distro/minikube.go b/pkg/driver/distro/minikube.go
index 99e72c47..891ea199 100644
--- a/pkg/driver/distro/minikube.go
+++ b/pkg/driver/distro/minikube.go
@@ -53,7 +53,6 @@ func (m *minikube) check(hostRoot string) bool {
 		}
 		matches := minikubeVersionRegex.FindStringSubmatch(string(bytes))
 		if len(matches) == 0 {
-			m.printer.Logger.Warn("Unable to extract minikube version.", m.printer.Logger.Args("path", hostRoot+"/etc/VERSION"))
 			return false
 		}
 		m.version = matches[1]
@@ -62,7 +61,7 @@ func (m *minikube) check(hostRoot string) bool {
 	return false
 }
 
-func (m *minikube) FixupKernel(i driverkernel.Info) driverkernel.Info {
+func (m *minikube) fixupKernel(i driverkernel.Info) driverkernel.Info {
 	i.KernelVersion = fmt.Sprintf("1_%s", m.version)
 	return i
 }
diff --git a/pkg/driver/distro/talos.go b/pkg/driver/distro/talos.go
index 8409b893..03e016f3 100644
--- a/pkg/driver/distro/talos.go
+++ b/pkg/driver/distro/talos.go
@@ -21,7 +21,6 @@ import (
 	"gopkg.in/ini.v1"
 
 	driverkernel "github.com/falcosecurity/falcoctl/pkg/driver/kernel"
-	"github.com/falcosecurity/falcoctl/pkg/output"
 )
 
 func init() {
@@ -33,7 +32,7 @@ type talos struct {
 	versionID string
 }
 
-func (t *talos) init(printer *output.Printer, id string, cfg *ini.File) error {
+func (t *talos) init(i driverkernel.Info, id string, cfg *ini.File) error {
 	idKey := cfg.Section("").Key("VERSION_ID")
 	if idKey == nil {
 		// OS-release without `VERSION_ID` (can it happen?)
@@ -41,10 +40,10 @@ func (t *talos) init(printer *output.Printer, id string, cfg *ini.File) error {
 	}
 	t.versionID = idKey.String()
 
-	return t.generic.init(printer, id, cfg)
+	return t.generic.init(i, id, cfg)
 }
 
-func (t *talos) FixupKernel(i driverkernel.Info) driverkernel.Info {
+func (t *talos) fixupKernel(i driverkernel.Info) driverkernel.Info {
 	i.KernelVersion = fmt.Sprintf("1_%s", t.versionID)
 	return i
 }
diff --git a/pkg/driver/distro/ubuntu.go b/pkg/driver/distro/ubuntu.go
index cdf1ab4b..28caf124 100644
--- a/pkg/driver/distro/ubuntu.go
+++ b/pkg/driver/distro/ubuntu.go
@@ -16,6 +16,7 @@
 package driverdistro
 
 import (
+	"gopkg.in/ini.v1"
 	"regexp"
 	"strings"
 
@@ -31,29 +32,29 @@ type ubuntu struct {
 }
 
 var ubuntuTargetIDRegex = regexp.MustCompile("-([a-zA-Z]+)(-.*)?$")
-var ubuntuKernelVersionRegex = regexp.MustCompile("^#\\d+\\~[^-]*-Ubuntu .*$")
+var ubuntuKernelVersionRegex = regexp.MustCompile("^\\d+\\~[^-]*-Ubuntu$")
 
-func (u *ubuntu) GetTargetID(i driverkernel.Info) string {
+func (u *ubuntu) init(i driverkernel.Info, _ string, f *ini.File) error {
 	// # Extract the flavor from the kernelrelease
-	//		# Examples:
-	//		# 5.0.0-1028-aws-5.0 -> ubuntu-aws
-	//		# 5.15.0-1009-aws -> ubuntu-aws
+	// # Examples:
+	//   # 5.0.0-1028-aws-5.0 -> ubuntu-aws
+	//   # 5.15.0-1009-aws -> ubuntu-aws
 	if ubuntuTargetIDRegex.MatchString(i.KernelRelease) {
-		return "ubuntu-" + ubuntuTargetIDRegex.FindStringSubmatch(i.KernelRelease)[1]
+		flavor := ubuntuTargetIDRegex.FindStringSubmatch(i.KernelRelease)[1]
+		return u.generic.init(i, "ubuntu-"+flavor, f)
 	}
-	return "ubuntu-generic"
+	return u.generic.init(i, "ubuntu-generic", f)
 }
 
-func (u *ubuntu) FixupKernel(i driverkernel.Info) driverkernel.Info {
+func (u *ubuntu) fixupKernel(i driverkernel.Info) driverkernel.Info {
 	// In the case that the kernelversion isn't just a number
 	// we keep also the remaining part excluding `-Ubuntu`.
 	// E.g.:
 	// from the following `uname -v` result
 	// `#26~22.04.1-Ubuntu SMP Mon Apr 24 01:58:15 UTC 2023`
 	// we obtain the kernelversion`26~22.04.1`
-	if ubuntuKernelVersionRegex.MatchString(i.FullKernelVersion) {
-		fullKV := strings.Trim(i.FullKernelVersion, "#")
-		i.KernelVersion = strings.Split(fullKV, "-Ubuntu")[0]
+	if ubuntuKernelVersionRegex.MatchString(i.KernelVersion) {
+		i.KernelVersion = strings.Split(i.KernelVersion, "-Ubuntu")[0]
 	}
 	return i
 }
diff --git a/pkg/driver/kernel/kernel.go b/pkg/driver/kernel/kernel.go
index 5ab32560..885bb225 100644
--- a/pkg/driver/kernel/kernel.go
+++ b/pkg/driver/kernel/kernel.go
@@ -24,10 +24,8 @@ import (
 
 // Info is the struct that holds all the kernel related information that are needed.
 type Info struct {
-	Architecture      string
-	KernelRelease     string
-	KernelVersion     string
-	FullKernelVersion string
+	KernelRelease string
+	KernelVersion string
 }
 
 // FetchInfo returns information about currently running kernel.
@@ -39,13 +37,19 @@ func FetchInfo() (Info, error) {
 
 	// Take eg: "#1 SMP PREEMPT_DYNAMIC Tue, 10 Oct 2023 21:10:21 +0000" and return "1".
 	kv := string(bytes.Trim(u.Version[:], "\x00"))
-	kv = strings.Trim(kv, "#")
-	kv = strings.Split(kv, " ")[0]
 
 	return Info{
-		Architecture:      string(bytes.Trim(u.Machine[:], "\x00")),
-		KernelRelease:     string(bytes.Trim(u.Release[:], "\x00")),
-		FullKernelVersion: string(bytes.Trim(u.Version[:], "\x00")),
-		KernelVersion:     kv,
+		KernelRelease: string(bytes.Trim(u.Release[:], "\x00")),
+		KernelVersion: FormatVersion(kv),
 	}, nil
 }
+
+// FormatVersion takes a kernelversion string (as taken from `uname -v`
+// and extracts the first part of the string.
+// Eg: '#26~22.04.1-Ubuntu SMP Mon Apr 24 01:58:15 UTC 2023' -> '26~22.04.1-Ubuntu'
+func FormatVersion(kv string) string {
+	// Take eg: "#1 SMP PREEMPT_DYNAMIC Tue, 10 Oct 2023 21:10:21 +0000" and return "1".
+	kv = strings.Trim(kv, "#")
+	kv = strings.Split(kv, " ")[0]
+	return kv
+}