From 6a338fc6adc8c85f25e51a8b05c650dafe5b17c4 Mon Sep 17 00:00:00 2001
From: Mikhail Klementev <blame@dumpstack.io>
Date: Thu, 18 May 2023 11:46:12 +0000
Subject: [PATCH] refactor: move ubutu install/cleanup to module

---
 distro/ubuntu/ubuntu.go | 20 ++++++++++++++++++++
 kernel/kernel.go        | 16 ++++++++++++----
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/distro/ubuntu/ubuntu.go b/distro/ubuntu/ubuntu.go
index d69b494..d24af95 100644
--- a/distro/ubuntu/ubuntu.go
+++ b/distro/ubuntu/ubuntu.go
@@ -82,3 +82,23 @@ func Match(km config.KernelMask) (pkgs []string, err error) {
 
 	return
 }
+
+func Install(km config.KernelMask, pkgname string, headers bool) (commands []string, err error) {
+
+	var headerspkg string
+	if headers {
+		headerspkg = strings.Replace(pkgname, "image", "headers", -1)
+	}
+
+	cmdf := func(f string, s ...interface{}) {
+		commands = append(commands, fmt.Sprintf(f, s...))
+	}
+
+	cmdf("apt-get install -y %s %s", pkgname, headerspkg)
+
+	return
+}
+
+func Cleanup(km config.KernelMask, pkgname string) {
+	return
+}
diff --git a/kernel/kernel.go b/kernel/kernel.go
index e45d9c9..8594a1d 100644
--- a/kernel/kernel.go
+++ b/kernel/kernel.go
@@ -229,12 +229,20 @@ func installKernel(sk config.KernelMask, pkgname string, force, headers bool) (e
 	// TODO install/cleanup kernel interface
 	switch sk.DistroType {
 	case config.Ubuntu:
-		var headerspkg string
-		if headers {
-			headerspkg = strings.Replace(pkgname, "image", "headers", -1)
+		var commands []string
+		commands, err = oraclelinux.Install(sk, pkgname, headers)
+		if err != nil {
+			return
 		}
+		defer func() {
+			if err != nil {
+				oraclelinux.Cleanup(sk, pkgname)
+			}
+		}()
 
-		cmd += fmt.Sprintf(" && apt-get install -y %s %s", pkgname, headerspkg)
+		for _, command := range commands {
+			cmd += fmt.Sprintf(" && %s", command)
+		}
 	case config.OracleLinux, config.CentOS:
 		var commands []string
 		commands, err = oraclelinux.Install(sk, pkgname, headers)