From 05903e34fffbe045f8f9c55bea658aedf339d107 Mon Sep 17 00:00:00 2001
From: David Wertenteil <dwertent@armosec.io>
Date: Thu, 22 Feb 2024 17:05:38 +0200
Subject: [PATCH] Handle unknown fw

Signed-off-by: David Wertenteil <dwertent@armosec.io>
---
 core/pkg/policyhandler/handlepullpolicies.go  |  4 ++--
 .../policyhandler/handlepullpoliciesutils.go  | 14 ++++++++++-
 .../handlepullpoliciesutils_test.go           | 24 ++++++++++++++++++-
 3 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/core/pkg/policyhandler/handlepullpolicies.go b/core/pkg/policyhandler/handlepullpolicies.go
index 051dadc8..0ffb70f2 100644
--- a/core/pkg/policyhandler/handlepullpolicies.go
+++ b/core/pkg/policyhandler/handlepullpolicies.go
@@ -150,7 +150,7 @@ func (policyHandler *PolicyHandler) downloadScanPolicies(ctx context.Context, po
 			logger.L().Debug("Downloading framework", helpers.String("framework", rule.Identifier))
 			receivedFramework, err := policyHandler.getters.PolicyGetter.GetFramework(rule.Identifier)
 			if err != nil {
-				return frameworks, policyDownloadError(err)
+				return frameworks, frameworkDownloadError(err, rule.Identifier)
 			}
 			if err := validateFramework(receivedFramework); err != nil {
 				return frameworks, err
@@ -171,7 +171,7 @@ func (policyHandler *PolicyHandler) downloadScanPolicies(ctx context.Context, po
 			logger.L().Debug("Downloading control", helpers.String("control", policy.Identifier))
 			receivedControl, err = policyHandler.getters.PolicyGetter.GetControl(policy.Identifier)
 			if err != nil {
-				return frameworks, policyDownloadError(err)
+				return frameworks, controlDownloadError(err, policy.Identifier)
 			}
 			if receivedControl != nil {
 				f.Controls = append(f.Controls, *receivedControl)
diff --git a/core/pkg/policyhandler/handlepullpoliciesutils.go b/core/pkg/policyhandler/handlepullpoliciesutils.go
index 693941da..89a905de 100644
--- a/core/pkg/policyhandler/handlepullpoliciesutils.go
+++ b/core/pkg/policyhandler/handlepullpoliciesutils.go
@@ -17,10 +17,22 @@ func getScanKind(policyIdentifier []cautils.PolicyIdentifier) apisv1.Notificatio
 	}
 	return "unknown"
 }
-func policyDownloadError(err error) error {
+func frameworkDownloadError(err error, fwName string) error {
 	if strings.Contains(err.Error(), "unsupported protocol scheme") {
 		err = fmt.Errorf("failed to download from GitHub release, try running with `--use-default` flag")
 	}
+	if strings.Contains(err.Error(), "not found") {
+		err = fmt.Errorf("framework '%s' not found, run `kubescape list frameworks` for available frameworks", fwName)
+	}
+	return err
+}
+func controlDownloadError(err error, controls string) error {
+	if strings.Contains(err.Error(), "unsupported protocol scheme") {
+		err = fmt.Errorf("failed to download from GitHub release, try running with `--use-default` flag")
+	}
+	if strings.Contains(err.Error(), "not found") {
+		err = fmt.Errorf("control '%s' not found, run `kubescape list controls` for available controls", controls)
+	}
 	return err
 }
 
diff --git a/core/pkg/policyhandler/handlepullpoliciesutils_test.go b/core/pkg/policyhandler/handlepullpoliciesutils_test.go
index 0791ed59..b3d0d3fe 100644
--- a/core/pkg/policyhandler/handlepullpoliciesutils_test.go
+++ b/core/pkg/policyhandler/handlepullpoliciesutils_test.go
@@ -89,6 +89,8 @@ func TestPolicyDownloadError(t *testing.T) {
 	tests := []struct {
 		err  error
 		want error
+		name string
+		kind string
 	}{
 		{
 			err:  errors.New("Some error"),
@@ -98,11 +100,31 @@ func TestPolicyDownloadError(t *testing.T) {
 			err:  errors.New("unsupported protocol scheme"),
 			want: fmt.Errorf("failed to download from GitHub release, try running with `--use-default` flag"),
 		},
+		{
+			err:  errors.New("framework 'cis' not found"),
+			want: fmt.Errorf("framework 'cis' not found, run `kubescape list frameworks` for available frameworks"),
+			name: "cis",
+			kind: "framework",
+		},
+		{
+			err:  errors.New("control 'c-0005' not found"),
+			want: fmt.Errorf("control 'c-0005' not found, run `kubescape list controls` for available controls"),
+			name: "c-0005",
+			kind: "control",
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run("", func(t *testing.T) {
-			assert.Equal(t, tt.want, policyDownloadError(tt.err))
+			switch tt.kind {
+			case "framework":
+				assert.Equal(t, tt.want, frameworkDownloadError(tt.err, tt.name))
+			case "control":
+				assert.Equal(t, tt.want, controlDownloadError(tt.err, tt.name))
+			default:
+				assert.Equal(t, tt.want, frameworkDownloadError(tt.err, tt.name))
+				assert.Equal(t, tt.want, controlDownloadError(tt.err, tt.name))
+			}
 		})
 	}
 }