From 69e0a8a4a2752b3484724f34d435db3ad2966a7d Mon Sep 17 00:00:00 2001
From: daemon1024 <barun1024@gmail.com>
Date: Thu, 25 May 2023 00:25:57 +0530
Subject: [PATCH] fix,core: restore container policy

We had a bug where we restored Backed Up container policy as host policy, this commit updates the restore logic to accomodate both container and host policies

Signed-off-by: daemon1024 <barun1024@gmail.com>
---
 KubeArmor/core/kubeArmor.go             |  2 +-
 KubeArmor/core/unorchestratedUpdates.go | 37 +++++++++++++++++++++----
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/KubeArmor/core/kubeArmor.go b/KubeArmor/core/kubeArmor.go
index b2ec103060..c89bc151be 100644
--- a/KubeArmor/core/kubeArmor.go
+++ b/KubeArmor/core/kubeArmor.go
@@ -660,7 +660,7 @@ func KubeArmor() {
 
 	if cfg.GlobalCfg.KVMAgent || !dm.K8sEnabled {
 		// Restore and apply all kubearmor host security policies
-		dm.restoreKubeArmorHostPolicies()
+		dm.restoreKubeArmorPolicies()
 	}
 
 	// == //
diff --git a/KubeArmor/core/unorchestratedUpdates.go b/KubeArmor/core/unorchestratedUpdates.go
index b70b824d14..9d4a678eea 100644
--- a/KubeArmor/core/unorchestratedUpdates.go
+++ b/KubeArmor/core/unorchestratedUpdates.go
@@ -576,7 +576,7 @@ func (dm *KubeArmorDaemon) backupKubeArmorContainerPolicy(policy tp.SecurityPoli
 	}
 }
 
-func (dm *KubeArmorDaemon) restoreKubeArmorHostPolicies() {
+func (dm *KubeArmorDaemon) restoreKubeArmorPolicies() {
 	if _, err := os.Stat(cfg.PolicyDir); err != nil {
 		kg.Warn("Policies dir not found for restoration")
 		return
@@ -586,15 +586,42 @@ func (dm *KubeArmorDaemon) restoreKubeArmorHostPolicies() {
 	if policyFiles, err := os.ReadDir(cfg.PolicyDir); err == nil {
 		for _, file := range policyFiles {
 			if data, err := os.ReadFile(cfg.PolicyDir + file.Name()); err == nil {
-				var hostPolicy tp.HostSecurityPolicy
-				if err := json.Unmarshal(data, &hostPolicy); err == nil {
-					dm.HostSecurityPolicies = append(dm.HostSecurityPolicies, hostPolicy)
+
+				var k struct {
+					Metadata map[string]string `json:"metadata"`
+				}
+
+				err := json.Unmarshal(data, &k)
+				if err != nil {
+					kg.Errf("Failed to unmarshal policy: %v", err)
+					continue
+				}
+
+				if _, ok := k.Metadata["namespaceName"]; ok { // ContainerPolicy contains namespaceName
+					var containerPolicy tp.K8sKubeArmorPolicy
+					if err := json.Unmarshal(data, &containerPolicy); err == nil {
+						containerPolicy.Metadata.Name = k.Metadata["policyName"]
+						dm.ParseAndUpdateContainerSecurityPolicy(tp.K8sKubeArmorPolicyEvent{
+							Type:   "ADDED",
+							Object: containerPolicy,
+						})
+					}
+
+				} else { // HostSecurityPolicy
+					var hostPolicy tp.HostSecurityPolicy
+					if err := json.Unmarshal(data, &hostPolicy); err == nil {
+						dm.HostSecurityPolicies = append(dm.HostSecurityPolicies, hostPolicy)
+					} else {
+						kg.Errf("Failed to unmarshal host policy: %v", err)
+					}
 				}
 			}
 		}
 
 		if len(policyFiles) != 0 {
-			dm.UpdateHostSecurityPolicies()
+			if len(dm.HostSecurityPolicies) != 0 {
+				dm.UpdateHostSecurityPolicies()
+			}
 		} else {
 			kg.Warn("No policies found for restoration")
 		}