Merge pull request #1276 from amirmalka/time-based-cached-policies

Time-based cached policies

core/cautils/strutils.go

@@ -2,6 +2,9 @@ package cautils
 
 import (
 	"fmt"
+	"os"
+	"sort"
+	"strconv"
 	"strings"
 )
 
@@ -42,3 +45,32 @@ func StringInSlice(strSlice []string, str string) int {
 	}
 	return ValueNotFound
 }
+
+func StringSlicesAreEqual(a, b []string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+
+	sort.Strings(a)
+	sort.Strings(b)
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}
+
+func ParseIntEnvVar(varName string, defaultValue int) (int, error) {
+	varValue, exists := os.LookupEnv(varName)
+	if !exists {
+		return defaultValue, nil
+	}
+
+	intValue, err := strconv.Atoi(varValue)
+	if err != nil {
+		return defaultValue, fmt.Errorf("failed to parse %s env var as int: %w", varName, err)
+	}
+
+	return intValue, nil
+}

core/cautils/strutils_test.go

@@ -2,8 +2,11 @@ package cautils
 
 import (
 	"fmt"
+	"os"
 	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestConvertLabelsToString(t *testing.T) {
@@ -33,3 +36,102 @@ func TestConvertStringToLabels(t *testing.T) {
 		t.Errorf("%s != %s", fmt.Sprintf("%v", rstrMap), fmt.Sprintf("%v", strMap))
 	}
 }
+
+func TestParseIntEnvVar(t *testing.T) {
+	testCases := []struct {
+		expectedErr  string
+		name         string
+		varName      string
+		varValue     string
+		defaultValue int
+		expected     int
+	}{
+		{
+			name:         "Variable does not exist",
+			varName:      "DOES_NOT_EXIST",
+			varValue:     "",
+			defaultValue: 123,
+			expected:     123,
+			expectedErr:  "",
+		},
+		{
+			name:         "Variable exists and is a valid integer",
+			varName:      "MY_VAR",
+			varValue:     "456",
+			defaultValue: 123,
+			expected:     456,
+			expectedErr:  "",
+		},
+		{
+			name:         "Variable exists but is not a valid integer",
+			varName:      "MY_VAR",
+			varValue:     "not_an_integer",
+			defaultValue: 123,
+			expected:     123,
+			expectedErr:  "failed to parse MY_VAR env var as int",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.varValue != "" {
+				os.Setenv(tc.varName, tc.varValue)
+			} else {
+				os.Unsetenv(tc.varName)
+			}
+
+			actual, err := ParseIntEnvVar(tc.varName, tc.defaultValue)
+			if tc.expectedErr != "" {
+				assert.NotNil(t, err)
+				assert.ErrorContains(t, err, tc.expectedErr)
+			} else {
+				assert.Nil(t, err)
+			}
+
+			assert.Equalf(t, tc.expected, actual, "unexpected result")
+		})
+	}
+}
+
+func TestStringSlicesAreEqual(t *testing.T) {
+	tt := []struct {
+		name string
+		a    []string
+		b    []string
+		want bool
+	}{
+		{
+			name: "equal unsorted slices",
+			a:    []string{"foo", "bar", "baz"},
+			b:    []string{"baz", "foo", "bar"},
+			want: true,
+		},
+		{
+			name: "equal sorted slices",
+			a:    []string{"bar", "baz", "foo"},
+			b:    []string{"bar", "baz", "foo"},
+			want: true,
+		},
+		{
+			name: "unequal slices",
+			a:    []string{"foo", "bar", "baz"},
+			b:    []string{"foo", "bar", "qux"},
+			want: false,
+		},
+		{
+			name: "different length slices",
+			a:    []string{"foo", "bar", "baz"},
+			b:    []string{"foo", "bar"},
+			want: false,
+		},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			got := StringSlicesAreEqual(tc.a, tc.b)
+			if got != tc.want {
+				t.Errorf("StringSlicesAreEqual(%v, %v) = %v; want %v", tc.a, tc.b, got, tc.want)
+			}
+		})
+	}
+}

core/core/scan.go

@@ -27,9 +27,9 @@ type componentInterfaces struct {
 	tenantConfig      cautils.ITenantConfig
 	resourceHandler   resourcehandler.IResourceHandler
 	report            reporter.IReport
-	outputPrinters    []printer.IPrinter
 	uiPrinter         printer.IPrinter
 	hostSensorHandler hostsensorutils.IHostSensor
+	outputPrinters    []printer.IPrinter
 }
 
 func getInterfaces(ctx context.Context, scanInfo *cautils.ScanInfo) componentInterfaces {
@@ -116,7 +116,6 @@ func getInterfaces(ctx context.Context, scanInfo *cautils.ScanInfo) componentInt
 
 func (ks *Kubescape) Scan(ctx context.Context, scanInfo *cautils.ScanInfo) (*resultshandling.ResultsHandler, error) {
 	ctxInit, spanInit := otel.Tracer("").Start(ctx, "initialization")
-
 	logger.L().Info("Kubescape scanner starting")
 
 	// ===================== Initialization =====================
@@ -151,15 +150,24 @@ func (ks *Kubescape) Scan(ctx context.Context, scanInfo *cautils.ScanInfo) (*res
 
 	resultsHandling := resultshandling.NewResultsHandler(interfaces.report, interfaces.outputPrinters, interfaces.uiPrinter)
 
-	// ===================== policies & resources =====================
-	ctxPolicies, spanPolicies := otel.Tracer("").Start(ctxInit, "policies & resources")
-	policyHandler := policyhandler.NewPolicyHandler(interfaces.resourceHandler)
-	scanData, err := policyHandler.CollectResources(ctxPolicies, scanInfo.PolicyIdentifier, scanInfo, cautils.NewProgressHandler(""))
+	// ===================== policies =====================
+	ctxPolicies, spanPolicies := otel.Tracer("").Start(ctxInit, "policies")
+	policyHandler := policyhandler.NewPolicyHandler()
+	scanData, err := policyHandler.CollectPolicies(ctxPolicies, scanInfo.PolicyIdentifier, scanInfo)
 	if err != nil {
 		spanInit.End()
 		return resultsHandling, err
 	}
 	spanPolicies.End()
+
+	// ===================== resources =====================
+	ctxResources, spanResources := otel.Tracer("").Start(ctxInit, "resources")
+	err = resourcehandler.CollectResources(ctxResources, interfaces.resourceHandler, scanInfo.PolicyIdentifier, scanData, cautils.NewProgressHandler(""))
+	if err != nil {
+		spanInit.End()
+		return resultsHandling, err
+	}
+	spanResources.End()
 	spanInit.End()
 
 	// ========================= opa testing =====================

core/pkg/opaprocessor/processorhandler.go

@@ -63,7 +63,7 @@ func (opap *OPAProcessor) ProcessRulesListener(ctx context.Context, progressList
 
 	ConvertFrameworksToSummaryDetails(&opap.Report.SummaryDetails, opap.Policies, opap.OPASessionObj.AllPolicies)
 
-	maxGoRoutines, err := parseIntEnvVar("RULE_PROCESSING_GOMAXPROCS", 2*runtime.NumCPU())
+	maxGoRoutines, err := cautils.ParseIntEnvVar("RULE_PROCESSING_GOMAXPROCS", 2*runtime.NumCPU())
 	if err != nil {
 		logger.L().Ctx(ctx).Warning(err.Error())
 	}

core/pkg/opaprocessor/utils.go

@@ -2,8 +2,6 @@ package opaprocessor
 
 import (
 	"fmt"
-	"os"
-	"strconv"
 
 	logger "github.com/kubescape/go-logger"
 	"github.com/kubescape/go-logger/helpers"
@@ -95,17 +93,3 @@ var cosignHasSignatureDefinition = func(bctx rego.BuiltinContext, a *ast.Term) (
 	}
 	return ast.BooleanTerm(has_signature(string(aStr))), nil
 }
-
-func parseIntEnvVar(varName string, defaultValue int) (int, error) {
-	varValue, exists := os.LookupEnv(varName)
-	if !exists {
-		return defaultValue, nil
-	}
-
-	intValue, err := strconv.Atoi(varValue)
-	if err != nil {
-		return defaultValue, fmt.Errorf("failed to parse %s env var as int: %w", varName, err)
-	}
-
-	return intValue, nil
-}

core/pkg/opaprocessor/utils_test.go

@@ -1,7 +1,6 @@
 package opaprocessor
 
 import (
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -29,59 +28,3 @@ func TestInitializeSummaryDetails(t *testing.T) {
 	assert.Equal(t, 2, len(summaryDetails.Frameworks))
 	assert.Equal(t, 3, len(summaryDetails.Controls))
 }
-
-func TestParseIntEnvVar(t *testing.T) {
-	testCases := []struct {
-		expectedErr  string
-		name         string
-		varName      string
-		varValue     string
-		defaultValue int
-		expected     int
-	}{
-		{
-			name:         "Variable does not exist",
-			varName:      "DOES_NOT_EXIST",
-			varValue:     "",
-			defaultValue: 123,
-			expected:     123,
-			expectedErr:  "",
-		},
-		{
-			name:         "Variable exists and is a valid integer",
-			varName:      "MY_VAR",
-			varValue:     "456",
-			defaultValue: 123,
-			expected:     456,
-			expectedErr:  "",
-		},
-		{
-			name:         "Variable exists but is not a valid integer",
-			varName:      "MY_VAR",
-			varValue:     "not_an_integer",
-			defaultValue: 123,
-			expected:     123,
-			expectedErr:  "failed to parse MY_VAR env var as int",
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			if tc.varValue != "" {
-				os.Setenv(tc.varName, tc.varValue)
-			} else {
-				os.Unsetenv(tc.varName)
-			}
-
-			actual, err := parseIntEnvVar(tc.varName, tc.defaultValue)
-			if tc.expectedErr != "" {
-				assert.NotNil(t, err)
-				assert.ErrorContains(t, err, tc.expectedErr)
-			} else {
-				assert.Nil(t, err)
-			}
-
-			assert.Equalf(t, tc.expected, actual, "unexpected result")
-		})
-	}
-}

core/pkg/policyhandler/cache.go

@@ -0,0 +1,73 @@
+package policyhandler
+
+import (
+	"sync"
+	"time"
+)
+
+// TimedCache provides functionality for managing a timed cache.
+// The timed cache holds a value for a specified time duration (TTL).
+// After the TTL has passed, the value is invalidated.
+//
+// The cache is thread safe.
+type TimedCache[T any] struct {
+	value      T
+	isSet      bool
+	ttl        time.Duration
+	expiration int64
+	mutex      sync.RWMutex
+}
+
+func NewTimedCache[T any](ttl time.Duration) *TimedCache[T] {
+	cache := &TimedCache[T]{
+		ttl:   ttl,
+		isSet: false,
+	}
+
+	// start the invalidate task only when the ttl is greater than 0 (cache is enabled)
+	if ttl > 0 {
+		go cache.invalidateTask()
+	}
+
+	return cache
+}
+
+func (c *TimedCache[T]) Set(value T) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	// cache is disabled
+	if c.ttl == 0 {
+		return
+	}
+
+	c.isSet = true
+	c.value = value
+	c.expiration = time.Now().Add(c.ttl).UnixNano()
+}
+
+func (c *TimedCache[T]) Get() (T, bool) {
+	c.mutex.RLock()
+	defer c.mutex.RUnlock()
+
+	if !c.isSet || time.Now().UnixNano() > c.expiration {
+		return c.value, false
+	}
+	return c.value, true
+}
+
+func (c *TimedCache[T]) invalidateTask() {
+	for {
+		<-time.After(c.ttl)
+		if time.Now().UnixNano() > c.expiration {
+			c.Invalidate()
+		}
+	}
+}
+
+func (c *TimedCache[T]) Invalidate() {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	c.isSet = false
+}

core/pkg/policyhandler/cache_test.go

@@ -0,0 +1,75 @@
+package policyhandler
+
+import (
+	"testing"
+	"time"
+)
+
+func TestTimedCache(t *testing.T) {
+	tests := []struct {
+		name string
+		// value ttl
+		ttl time.Duration
+		// value to set
+		value int
+		// time to wait before checking if value exists
+		wait time.Duration
+		// number of times to check if value exists (with wait in between)
+		checks int
+		// should the value exist in cache
+		exists bool
+		// expected cache value
+		wantVal int
+	}{
+		{
+			name:    "value exists before ttl",
+			ttl:     time.Second * 5,
+			value:   42,
+			wait:    time.Second * 1,
+			checks:  2,
+			exists:  true,
+			wantVal: 42,
+		},
+		{
+			name:   "value does not exist after ttl",
+			ttl:    time.Second * 3,
+			value:  55,
+			wait:   time.Second * 4,
+			checks: 1,
+			exists: false,
+		},
+		{
+			name:   "cache is disabled (ttl = 0) always returns false",
+			ttl:    0,
+			value:  55,
+			wait:   0,
+			checks: 1,
+			exists: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cache := NewTimedCache[int](tt.ttl)
+			cache.Set(tt.value)
+
+			for i := 0; i < tt.checks; i++ {
+				// Wait for the specified duration
+				time.Sleep(tt.wait)
+
+				// Get the value from the cache
+				value, exists := cache.Get()
+
+				// Check if value exists
+				if exists != tt.exists {
+					t.Errorf("Expected exists to be %v, got %v", tt.exists, exists)
+				}
+
+				// Check value
+				if exists && value != tt.wantVal {
+					t.Errorf("Expected value to be %d, got %d", tt.wantVal, value)
+				}
+			}
+		})
+	}
+}

core/pkg/policyhandler/handlepullpolicies.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/armosec/armoapi-go/armotypes"
 	logger "github.com/kubescape/go-logger"
 	"github.com/kubescape/go-logger/helpers"
 	"github.com/kubescape/kubescape/v2/core/cautils"
@@ -14,7 +15,55 @@ import (
 	"go.opentelemetry.io/otel"
 )
 
-func (policyHandler *PolicyHandler) getPolicies(ctx context.Context, policyIdentifier []cautils.PolicyIdentifier, policiesAndResources *cautils.OPASessionObj) error {
+const (
+	PoliciesCacheTtlEnvVar = "POLICIES_CACHE_TTL"
+)
+
+var policyHandlerInstance *PolicyHandler
+
+// PolicyHandler
+type PolicyHandler struct {
+	getters                 *cautils.Getters
+	cachedPolicyIdentifiers *TimedCache[[]string]
+	cachedFrameworks        *TimedCache[[]reporthandling.Framework]
+	cachedExceptions        *TimedCache[[]armotypes.PostureExceptionPolicy]
+	cachedControlInputs     *TimedCache[map[string][]string]
+}
+
+// NewPolicyHandler creates and returns an instance of the `PolicyHandler`. The function initializes the `PolicyHandler` only if it hasn't been previously created.
+// The PolicyHandler supports caching of downloaded policies and exceptions by setting the `POLICIES_CACHE_TTL` environment variable (default is no caching).
+func NewPolicyHandler() *PolicyHandler {
+	if policyHandlerInstance == nil {
+		cacheTtl := getPoliciesCacheTtl()
+		policyHandlerInstance = &PolicyHandler{
+			cachedPolicyIdentifiers: NewTimedCache[[]string](cacheTtl),
+			cachedFrameworks:        NewTimedCache[[]reporthandling.Framework](cacheTtl),
+			cachedExceptions:        NewTimedCache[[]armotypes.PostureExceptionPolicy](cacheTtl),
+			cachedControlInputs:     NewTimedCache[map[string][]string](cacheTtl),
+		}
+	}
+	return policyHandlerInstance
+}
+
+func (policyHandler *PolicyHandler) CollectPolicies(ctx context.Context, policyIdentifier []cautils.PolicyIdentifier, scanInfo *cautils.ScanInfo) (*cautils.OPASessionObj, error) {
+	opaSessionObj := cautils.NewOPASessionObj(ctx, nil, nil, scanInfo)
+
+	policyHandler.getters = &scanInfo.Getters
+
+	// get policies, exceptions and controls inputs
+	policies, exceptions, controlInputs, err := policyHandler.getPolicies(ctx, policyIdentifier)
+	if err != nil {
+		return opaSessionObj, err
+	}
+
+	opaSessionObj.Policies = policies
+	opaSessionObj.Exceptions = exceptions
+	opaSessionObj.RegoInputData.PostureControlInputs = controlInputs
+
+	return opaSessionObj, nil
+}
+
+func (policyHandler *PolicyHandler) getPolicies(ctx context.Context, policyIdentifier []cautils.PolicyIdentifier) (policies []reporthandling.Framework, exceptions []armotypes.PostureExceptionPolicy, controlInputs map[string][]string, err error) {
 	ctx, span := otel.Tracer("").Start(ctx, "policyHandler.getPolicies")
 	defer span.End()
 	logger.L().Info("Downloading/Loading policy definitions")
@@ -22,38 +71,57 @@ func (policyHandler *PolicyHandler) getPolicies(ctx context.Context, policyIdent
 	cautils.StartSpinner()
 	defer cautils.StopSpinner()
 
-	policies, err := policyHandler.getScanPolicies(ctx, policyIdentifier)
+	// get policies
+	policies, err = policyHandler.getScanPolicies(ctx, policyIdentifier)
 	if err != nil {
-		return err
+		return nil, nil, nil, err
 	}
 	if len(policies) == 0 {
-		return fmt.Errorf("failed to download policies: '%s'. Make sure the policy exist and you spelled it correctly. For more information, please feel free to contact ARMO team", strings.Join(policyIdentifierToSlice(policyIdentifier), ", "))
+		return nil, nil, nil, fmt.Errorf("failed to download policies: '%s'. Make sure the policy exist and you spelled it correctly. For more information, please feel free to contact ARMO team", strings.Join(policyIdentifierToSlice(policyIdentifier), ", "))
 	}
 
-	policiesAndResources.Policies = policies
-
 	// get exceptions
-	exceptionPolicies, err := policyHandler.getters.ExceptionsGetter.GetExceptions(cautils.ClusterName)
-	if err == nil {
-		policiesAndResources.Exceptions = exceptionPolicies
-	} else {
+	if exceptions, err = policyHandler.getExceptions(); err != nil {
 		logger.L().Ctx(ctx).Warning("failed to load exceptions", helpers.Error(err))
 	}
 
 	// get account configuration
-	controlsInputs, err := policyHandler.getters.ControlsInputsGetter.GetControlsInputs(cautils.ClusterName)
-	if err == nil {
-		policiesAndResources.RegoInputData.PostureControlInputs = controlsInputs
-	} else {
+	if controlInputs, err = policyHandler.getControlInputs(); err != nil {
 		logger.L().Ctx(ctx).Warning(err.Error())
 	}
-	cautils.StopSpinner()
 
+	cautils.StopSpinner()
 	logger.L().Success("Downloaded/Loaded policy")
-	return nil
+
+	return policies, exceptions, controlInputs, nil
 }
 
+// getScanPolicies - get policies from cache or downloads them. The function returns an error if the policies could not be downloaded.
 func (policyHandler *PolicyHandler) getScanPolicies(ctx context.Context, policyIdentifier []cautils.PolicyIdentifier) ([]reporthandling.Framework, error) {
+	policyIdentifiersSlice := policyIdentifierToSlice(policyIdentifier)
+	// check if policies are cached
+	if cachedPolicies, policiesExist := policyHandler.cachedFrameworks.Get(); policiesExist {
+		// check if the cached policies are the same as the requested policies, otherwise download the policies
+		if cachedIdentifiers, identifiersExist := policyHandler.cachedPolicyIdentifiers.Get(); identifiersExist && cautils.StringSlicesAreEqual(cachedIdentifiers, policyIdentifiersSlice) {
+			logger.L().Info("Using cached policies")
+			return cachedPolicies, nil
+		}
+
+		logger.L().Debug("Cached policies are not the same as the requested policies")
+		policyHandler.cachedPolicyIdentifiers.Invalidate()
+		policyHandler.cachedFrameworks.Invalidate()
+	}
+
+	policies, err := policyHandler.downloadScanPolicies(ctx, policyIdentifier)
+	if err == nil {
+		policyHandler.cachedFrameworks.Set(policies)
+		policyHandler.cachedPolicyIdentifiers.Set(policyIdentifiersSlice)
+	}
+
+	return policies, err
+}
+
+func (policyHandler *PolicyHandler) downloadScanPolicies(ctx context.Context, policyIdentifier []cautils.PolicyIdentifier) ([]reporthandling.Framework, error) {
 	frameworks := []reporthandling.Framework{}
 
 	switch getScanKind(policyIdentifier) {
@@ -102,10 +170,30 @@ func (policyHandler *PolicyHandler) getScanPolicies(ctx context.Context, policyI
 	return frameworks, nil
 }
 
-func policyIdentifierToSlice(rules []cautils.PolicyIdentifier) []string {
-	s := []string{}
-	for i := range rules {
-		s = append(s, fmt.Sprintf("%s: %s", rules[i].Kind, rules[i].Identifier))
+func (policyHandler *PolicyHandler) getExceptions() ([]armotypes.PostureExceptionPolicy, error) {
+	if cachedExceptions, exist := policyHandler.cachedExceptions.Get(); exist {
+		logger.L().Info("Using cached exceptions")
+		return cachedExceptions, nil
 	}
-	return s
+
+	exceptions, err := policyHandler.getters.ExceptionsGetter.GetExceptions(cautils.ClusterName)
+	if err == nil {
+		policyHandler.cachedExceptions.Set(exceptions)
+	}
+
+	return exceptions, err
+}
+
+func (policyHandler *PolicyHandler) getControlInputs() (map[string][]string, error) {
+	if cachedControlInputs, exist := policyHandler.cachedControlInputs.Get(); exist {
+		logger.L().Info("Using cached control inputs")
+		return cachedControlInputs, nil
+	}
+
+	controlInputs, err := policyHandler.getters.ControlsInputsGetter.GetControlsInputs(cautils.ClusterName)
+	if err == nil {
+		policyHandler.cachedControlInputs.Set(controlInputs)
+	}
+
+	return controlInputs, err
 }

core/pkg/policyhandler/handlepullpoliciesutils.go

@@ -3,6 +3,7 @@ package policyhandler
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	apisv1 "github.com/kubescape/opa-utils/httpserver/apis/v1"
 	"github.com/kubescape/opa-utils/reporthandling"
@@ -35,3 +36,20 @@ func validateFramework(framework *reporthandling.Framework) error {
 	}
 	return nil
 }
+
+// getPoliciesCacheTtl - get policies cache TTL from environment variable or return 0 if not set
+func getPoliciesCacheTtl() time.Duration {
+	if val, err := cautils.ParseIntEnvVar(PoliciesCacheTtlEnvVar, 0); err == nil {
+		return time.Duration(val) * time.Minute
+	}
+
+	return 0
+}
+
+func policyIdentifierToSlice(rules []cautils.PolicyIdentifier) []string {
+	s := []string{}
+	for i := range rules {
+		s = append(s, fmt.Sprintf("%s: %s", rules[i].Kind, rules[i].Identifier))
+	}
+	return s
+}

core/pkg/policyhandler/handlepullpoliciesutils_test.go

@@ -11,8 +11,8 @@ func Test_validateFramework(t *testing.T) {
 		framework *reporthandling.Framework
 	}
 	tests := []struct {
-		name    string
 		args    args
+		name    string
 		wantErr bool
 	}{
 		{

core/pkg/resourcehandler/handlepullresources_test.go

@@ -1,4 +1,4 @@
-package policyhandler
+package resourcehandler
 
 import (
 	"context"
@@ -8,7 +8,7 @@ import (
 
 	"github.com/kubescape/k8s-interface/k8sinterface"
 	"github.com/kubescape/kubescape/v2/core/cautils"
-	"github.com/kubescape/kubescape/v2/core/pkg/resourcehandler"
+
 	"github.com/kubescape/opa-utils/reporthandling/apis"
 	helpersv1 "github.com/kubescape/opa-utils/reporthandling/helpers/v1"
 	reporthandlingv2 "github.com/kubescape/opa-utils/reporthandling/v2"
@@ -22,7 +22,7 @@ import (
 )
 
 var (
-	//go:embed kubeconfig_mock.json
+	//go:embed testdata/kubeconfig_mock.json
 	kubeConfigMock string
 )
 
@@ -35,9 +35,9 @@ func getKubeConfigMock() *clientcmdapi.Config {
 }
 func Test_getCloudMetadata(t *testing.T) {
 	type args struct {
-		context       string
 		opaSessionObj *cautils.OPASessionObj
 		kubeConfig    *clientcmdapi.Config
+		context       string
 	}
 	kubeConfig := getKubeConfigMock()
 	tests := []struct {
@@ -221,7 +221,7 @@ func (*iResourceHandlerMock) GetClusterAPIServerInfo() *version.Info {
 
 // https://github.com/kubescape/kubescape/pull/1004
 // Cluster named .*eks.* config without a cloudconfig panics whereas we just want to scan a file
-func getResourceHandlerMock() *resourcehandler.K8sResourceHandler {
+func getResourceHandlerMock() *K8sResourceHandler {
 	client := fakeclientset.NewSimpleClientset()
 	fakeDiscovery := client.Discovery()
 
@@ -232,10 +232,10 @@ func getResourceHandlerMock() *resourcehandler.K8sResourceHandler {
 		Context:          context.Background(),
 	}
 
-	return resourcehandler.NewK8sResourceHandler(k8s, &resourcehandler.EmptySelector{}, nil, nil, nil)
+	return NewK8sResourceHandler(k8s, &EmptySelector{}, nil, nil, nil)
 }
-func Test_getResources(t *testing.T) {
-	policyHandler := &PolicyHandler{resourceHandler: getResourceHandlerMock()}
+func Test_CollectResources(t *testing.T) {
+	resourceHandler := getResourceHandlerMock()
 	objSession := &cautils.OPASessionObj{
 		Metadata: &reporthandlingv2.Metadata{
 			ScanMetadata: reporthandlingv2.ScanMetadata{
@@ -249,12 +249,12 @@ func Test_getResources(t *testing.T) {
 	policyIdentifier := []cautils.PolicyIdentifier{{}}
 
 	assert.NotPanics(t, func() {
-		policyHandler.getResources(context.TODO(), policyIdentifier, objSession, cautils.NewProgressHandler(""))
+		CollectResources(context.TODO(), resourceHandler, policyIdentifier, objSession, cautils.NewProgressHandler(""))
 	}, "Cluster named .*eks.* without a cloud config panics on cluster scan !")
 
 	assert.NotPanics(t, func() {
 		objSession.Metadata.ScanMetadata.ScanningTarget = reportv2.File
-		policyHandler.getResources(context.TODO(), policyIdentifier, objSession, cautils.NewProgressHandler(""))
+		CollectResources(context.TODO(), resourceHandler, policyIdentifier, objSession, cautils.NewProgressHandler(""))
 	}, "Cluster named .*eks.* without a cloud config panics on non-cluster scan !")
 
 }

core/pkg/resourcehandler/handlerpullresources.go

@@ -1,4 +1,4 @@
-package policyhandler
+package resourcehandler
 
 import (
 	"context"
@@ -7,70 +7,30 @@ import (
 
 	logger "github.com/kubescape/go-logger"
 	"github.com/kubescape/go-logger/helpers"
-	helpersv1 "github.com/kubescape/opa-utils/reporthandling/helpers/v1"
-	"go.opentelemetry.io/otel"
-	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
-
-	cloudsupportv1 "github.com/kubescape/k8s-interface/cloudsupport/v1"
-	"github.com/kubescape/kubescape/v2/core/pkg/opaprocessor"
-	reportv2 "github.com/kubescape/opa-utils/reporthandling/v2"
-
 	"github.com/kubescape/k8s-interface/cloudsupport"
+	cloudsupportv1 "github.com/kubescape/k8s-interface/cloudsupport/v1"
 	"github.com/kubescape/k8s-interface/k8sinterface"
 	"github.com/kubescape/kubescape/v2/core/cautils"
-	"github.com/kubescape/kubescape/v2/core/pkg/resourcehandler"
+	"github.com/kubescape/kubescape/v2/core/pkg/opaprocessor"
 	"github.com/kubescape/opa-utils/reporthandling/apis"
+	helpersv1 "github.com/kubescape/opa-utils/reporthandling/helpers/v1"
+	reportv2 "github.com/kubescape/opa-utils/reporthandling/v2"
+	"go.opentelemetry.io/otel"
+	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
 )
 
-// PolicyHandler -
-type PolicyHandler struct {
-	resourceHandler resourcehandler.IResourceHandler
-	// we are listening on this chan in opaprocessor/processorhandler.go/ProcessRulesListener func
-	getters *cautils.Getters
-}
-
-// CreatePolicyHandler Create ws-handler obj
-func NewPolicyHandler(resourceHandler resourcehandler.IResourceHandler) *PolicyHandler {
-	return &PolicyHandler{
-		resourceHandler: resourceHandler,
-	}
-}
-
-func (policyHandler *PolicyHandler) CollectResources(ctx context.Context, policyIdentifier []cautils.PolicyIdentifier, scanInfo *cautils.ScanInfo, progressListener opaprocessor.IJobProgressNotificationClient) (*cautils.OPASessionObj, error) {
-	opaSessionObj := cautils.NewOPASessionObj(ctx, nil, nil, scanInfo)
-
-	// validate notification
-	// TODO
-	policyHandler.getters = &scanInfo.Getters
-
-	// get policies
-	if err := policyHandler.getPolicies(ctx, policyIdentifier, opaSessionObj); err != nil {
-		return opaSessionObj, err
-	}
-
-	err := policyHandler.getResources(ctx, policyIdentifier, opaSessionObj, progressListener)
-	if err != nil {
-		return opaSessionObj, err
-	}
-	if (opaSessionObj.K8SResources == nil || len(*opaSessionObj.K8SResources) == 0) && (opaSessionObj.ArmoResource == nil || len(*opaSessionObj.ArmoResource) == 0) {
-		return opaSessionObj, fmt.Errorf("empty list of resources")
-	}
-
-	// update channel
-	return opaSessionObj, nil
-}
-
-func (policyHandler *PolicyHandler) getResources(ctx context.Context, policyIdentifier []cautils.PolicyIdentifier, opaSessionObj *cautils.OPASessionObj, progressListener opaprocessor.IJobProgressNotificationClient) error {
-	ctx, span := otel.Tracer("").Start(ctx, "policyHandler.getResources")
+// CollectResources uses the provided resource handler to collect resources and returns an updated OPASessionObj
+func CollectResources(ctx context.Context, rsrcHandler IResourceHandler, policyIdentifier []cautils.PolicyIdentifier, opaSessionObj *cautils.OPASessionObj, progressListener opaprocessor.IJobProgressNotificationClient) error {
+	ctx, span := otel.Tracer("").Start(ctx, "resourcehandler.CollectResources")
 	defer span.End()
-	opaSessionObj.Report.ClusterAPIServerInfo = policyHandler.resourceHandler.GetClusterAPIServerInfo(ctx)
+	opaSessionObj.Report.ClusterAPIServerInfo = rsrcHandler.GetClusterAPIServerInfo(ctx)
 
 	// set cloud metadata only when scanning a cluster
 	if opaSessionObj.Metadata.ScanMetadata.ScanningTarget == reportv2.Cluster {
 		setCloudMetadata(opaSessionObj)
 	}
 
-	resourcesMap, allResources, ksResources, err := policyHandler.resourceHandler.GetResources(ctx, opaSessionObj, &policyIdentifier[0].Designators, progressListener)
+	resourcesMap, allResources, ksResources, err := rsrcHandler.GetResources(ctx, opaSessionObj, &policyIdentifier[0].Designators, progressListener)
 	if err != nil {
 		return err
 	}
@@ -79,18 +39,13 @@ func (policyHandler *PolicyHandler) getResources(ctx context.Context, policyIden
 	opaSessionObj.AllResources = allResources
 	opaSessionObj.ArmoResource = ksResources
 
+	if (opaSessionObj.K8SResources == nil || len(*opaSessionObj.K8SResources) == 0) && (opaSessionObj.ArmoResource == nil || len(*opaSessionObj.ArmoResource) == 0) {
+		return fmt.Errorf("empty list of resources")
+	}
+
 	return nil
 }
 
-/* unused for now
-func getDesignator(policyIdentifier []cautils.PolicyIdentifier) *armotypes.PortalDesignator {
-	if len(policyIdentifier) > 0 {
-		return &policyIdentifier[0].Designators
-	}
-	return &armotypes.PortalDesignator{}
-}
-*/
-
 func setCloudMetadata(opaSessionObj *cautils.OPASessionObj) {
 	iCloudMetadata := getCloudMetadata(opaSessionObj, k8sinterface.GetConfig())
 	if iCloudMetadata == nil {

core/pkg/resultshandling/printer/v2/attacktracks.go

@@ -2,9 +2,7 @@ package printer
 
 import (
 	"fmt"
-	"os"
 	"sort"
-	"strconv"
 	"strings"
 
 	"github.com/fatih/color"
@@ -77,23 +75,14 @@ func (prettyPrinter *PrettyPrinter) printResourceAttackGraph(attackTrack v1alpha
 	fmt.Fprintln(prettyPrinter.writer, tree.Print())
 }
 
-func getNumericValueFromEnvVar(envVar string, defaultValue int) int {
-	value := os.Getenv(envVar)
-	if value != "" {
-		if value, err := strconv.Atoi(value); err == nil {
-			return value
-		}
-	}
-	return defaultValue
-}
 func (prettyPrinter *PrettyPrinter) printAttackTracks(opaSessionObj *cautils.OPASessionObj) {
 	if !prettyPrinter.printAttackTree || opaSessionObj.ResourceAttackTracks == nil {
 		return
 	}
 
 	// check if counters are set in env vars and use them, otherwise use default values
-	topResourceCount := getNumericValueFromEnvVar("ATTACK_TREE_TOP_RESOURCES", TOP_RESOURCE_COUNT)
-	topVectorCount := getNumericValueFromEnvVar("ATTACK_TREE_TOP_VECTORS", TOP_VECTOR_COUNT)
+	topResourceCount, _ := cautils.ParseIntEnvVar("ATTACK_TREE_TOP_RESOURCES", TOP_RESOURCE_COUNT)
+	topVectorCount, _ := cautils.ParseIntEnvVar("ATTACK_TREE_TOP_VECTORS", TOP_VECTOR_COUNT)
 
 	prioritizedResources := opaSessionObj.ResourcesPrioritized
 	resourceToAttackTrack := opaSessionObj.ResourceAttackTracks