From 958ae9393dd7e4e30e459c8142a4314db0c42b14 Mon Sep 17 00:00:00 2001
From: mudler <mudler@localai.io>
Date: Tue, 27 Feb 2024 17:25:20 +0100
Subject: [PATCH] fix(after-install): run after-install hook after partitions
 are encrypted

---
 internal/agent/hooks/hook.go |  9 +++++++--
 pkg/action/install.go        | 12 +++++++++---
 pkg/uki/install.go           |  5 +++++
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/internal/agent/hooks/hook.go b/internal/agent/hooks/hook.go
index 1957eba1..d165a6a0 100644
--- a/internal/agent/hooks/hook.go
+++ b/internal/agent/hooks/hook.go
@@ -13,7 +13,6 @@ var AfterInstall = []Interface{
 	&GrubOptions{}, // Set custom GRUB options
 	&BundlePostInstall{},
 	&CustomMounts{},
-	&Kcrypt{},
 	&Lifecycle{}, // Handles poweroff/reboot by config options
 }
 
@@ -31,10 +30,16 @@ var FirstBoot = []Interface{
 }
 
 // AfterUkiInstall sets which Hooks to run after uki runs the install action
-var AfterUkiInstall = []Interface{
+var AfterUkiInstall = []Interface{}
+
+var UKIEncryptionHooks = []Interface{
 	&KcryptUKI{},
 }
 
+var EncryptionHooks = []Interface{
+	&Kcrypt{},
+}
+
 func Run(c config.Config, spec v1.Spec, hooks ...Interface) error {
 	for _, h := range hooks {
 		if err := h.Run(c, spec); err != nil {
diff --git a/pkg/action/install.go b/pkg/action/install.go
index 8ec631fa..31243dcd 100644
--- a/pkg/action/install.go
+++ b/pkg/action/install.go
@@ -18,13 +18,14 @@ package action
 
 import (
 	"fmt"
-	hook "github.com/kairos-io/kairos-agent/v2/internal/agent/hooks"
-	"github.com/kairos-io/kairos-agent/v2/pkg/config"
-	events "github.com/kairos-io/kairos-sdk/bus"
 	"path/filepath"
 	"strings"
 	"time"
 
+	hook "github.com/kairos-io/kairos-agent/v2/internal/agent/hooks"
+	"github.com/kairos-io/kairos-agent/v2/pkg/config"
+	events "github.com/kairos-io/kairos-sdk/bus"
+
 	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"
@@ -252,6 +253,11 @@ func (i InstallAction) Run() (err error) {
 		return err
 	}
 
+	err = hook.Run(*i.cfg, i.spec, hook.EncryptionHooks...)
+	if err != nil {
+		return err
+	}
+
 	err = i.installHook(cnst.AfterInstallHook, false)
 	if err != nil {
 		return err
diff --git a/pkg/uki/install.go b/pkg/uki/install.go
index 66a69f40..001304a0 100644
--- a/pkg/uki/install.go
+++ b/pkg/uki/install.go
@@ -144,6 +144,11 @@ func (i *InstallAction) Run() (err error) {
 		return fmt.Errorf("removing artifact set with role %s: %w", UnassignedArtifactRole, err)
 	}
 
+	err = hook.Run(*i.cfg, i.spec, hook.UKIEncryptionHooks...)
+	if err != nil {
+		return err
+	}
+
 	// after install hook happens after install (this is for compatibility with normal install, so users can reuse their configs)
 	err = Hook(i.cfg, constants.AfterInstallHook)
 	if err != nil {