From 7acb4d1bb2c5b100bc54d34002a8ab8e68260a6a Mon Sep 17 00:00:00 2001
From: daemon1024 <barun1024@gmail.com>
Date: Thu, 22 Jun 2023 23:11:32 +0530
Subject: [PATCH] feat(install): wait for KubeArmor to create probe file before
 probing

We keep trying to probe information from KubeArmor till we get info out or context times out

Signed-off-by: daemon1024 <barun1024@gmail.com>
---
 install/install.go | 52 +++++++++++++++++++++++++++++++---------------
 probe/probe.go     |  6 ++----
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/install/install.go b/install/install.go
index bbec012b..8bff197b 100644
--- a/install/install.go
+++ b/install/install.go
@@ -105,7 +105,7 @@ func printAnimation(msg string, flag bool) int {
 	if flag {
 		progress++
 	}
-	printBar("    KubeArmor Installing ", 16)
+	printBar("    KubeArmor Installing ", 17)
 	return 0
 }
 
@@ -141,23 +141,41 @@ func checkPods(c *k8s.Client, o Options) {
 			break
 		}
 	}
-	probeData, err := probe.ProbeRunningKubeArmorNodes(c, probe.Options{
-		Namespace: o.Namespace,
-	})
-	if err != nil || len(probeData) == 0 {
-		return
-	}
-	enforcing := true
-	for _, k := range probeData {
-		if k.ActiveLSM == "" || !k.ContainerSecurity {
-			enforcing = false
-			break
+	fmt.Print("\n🔧  Verifying KubeArmor functionality (this may take upto a minute) ...")
+	ctx, cancel := context.WithTimeout(context.Background(), 40*time.Second)
+	defer cancel()
+
+	for {
+		select {
+		case <-time.After(1 * time.Second):
+		case <-ctx.Done():
+			fmt.Print("⚠️  Failed verifying KubeArmor functionality ...")
+			return
 		}
-	}
-	if enforcing {
-		fmt.Print(color.New(color.FgWhite, color.Bold).Sprint("\n\t🛡️  Your Cluster is Armored Up Now! \n"))
-	} else {
-		color.Yellow("\n\t⚠️  KubeArmor is running in Audit mode, only Observability will be available and Policy Enforcement won't work. \n")
+		probeData, err := probe.ProbeRunningKubeArmorNodes(c, probe.Options{
+			Namespace: o.Namespace,
+		})
+		if err != nil || len(probeData) == 0 {
+			fmt.Printf("\r🔧  Verifying KubeArmor functionality (this may take upto a minute) ... %s", cursor[cursorcount])
+			cursorcount++
+			if cursorcount == 4 {
+				cursorcount = 0
+			}
+			continue
+		}
+		enforcing := true
+		for _, k := range probeData {
+			if k.ActiveLSM == "" || !k.ContainerSecurity {
+				enforcing = false
+				break
+			}
+		}
+		if enforcing {
+			fmt.Print(color.New(color.FgWhite, color.Bold).Sprint("\n\n\t🛡️  Your Cluster is Armored Up! \n"))
+		} else {
+			color.Yellow("\n\n\t⚠️  KubeArmor is running in Audit mode, only Observability will be available and Policy Enforcement won't work. \n")
+		}
+		break
 	}
 
 }
diff --git a/probe/probe.go b/probe/probe.go
index 32f3a65c..a2ecbd8d 100644
--- a/probe/probe.go
+++ b/probe/probe.go
@@ -489,9 +489,9 @@ func readDataFromKubeArmor(c *k8s.Client, o Options, nodeName string) (KubeArmor
 		VersionedParams(&corev1.PodExecOptions{
 			Container: pods.Items[0].Spec.Containers[0].Name,
 			Command:   cmdArr,
-			Stdin:     true,
+			Stdin:     false,
 			Stdout:    true,
-			Stderr:    true,
+			Stderr:    false,
 			TTY:       false,
 		}, scheme.ParameterCodec)
 	exec, err := remotecommand.NewSPDYExecutor(c.Config, "POST", req.URL())
@@ -501,9 +501,7 @@ func readDataFromKubeArmor(c *k8s.Client, o Options, nodeName string) (KubeArmor
 	go func() {
 		defer outStream.Close()
 		err = exec.StreamWithContext(context.TODO(), remotecommand.StreamOptions{
-			Stdin:  os.Stdin,
 			Stdout: outStream,
-			Stderr: os.Stderr,
 			Tty:    false,
 		})
 	}()