Add attack track controls to gitRegoStore

Signed-off-by: YiscahLevySilas1 <[email protected]>

gitregostore/datastructures.go

@@ -18,6 +18,7 @@ type GitRegoStore struct {
 	DefaultConfigInputsLock            sync.RWMutex
 	rulesLock                          sync.RWMutex
 	controlsLock                       sync.RWMutex
+	attackTrackControlsLock            sync.RWMutex
 	attackTracksLock                   sync.RWMutex
 	systemPostureExceptionPoliciesLock sync.RWMutex
 	controlRelationsLock               sync.RWMutex
@@ -39,6 +40,7 @@ type GitRegoStore struct {
 	AttackTracks                       []v1alpha1.AttackTrack
 	Frameworks                         []opapolicy.Framework
 	Controls                           []opapolicy.Control
+	AttackTrackControls                []opapolicy.Control
 	Rules                              []opapolicy.PolicyRule
 	SystemPostureExceptionPolicies     []armotypes.PostureExceptionPolicy
 	FrequencyPullFromGitMinutes        int

gitregostore/gitstoremethods.go

@@ -203,6 +203,28 @@ func (gs *GitRegoStore) GetOPAControls() ([]opapolicy.Control, error) {
 	return controlsList, nil
 }
 
+func (gs *GitRegoStore) GetOPAAttackTrackControls() ([]opapolicy.Control, error) {
+	gs.attackTrackControlsLock.RLock()
+	defer gs.attackTrackControlsLock.RUnlock()
+
+	if gs.AttackTrackControls == nil {
+		return nil, fmt.Errorf("no controls found in GitRegoStore")
+	}
+
+	attackTrackControlsList := make([]opapolicy.Control, 0, len(gs.AttackTrackControls))
+	for _, controlToPin := range gs.AttackTrackControls {
+		control := controlToPin
+
+		if err := gs.fillRulesAndRulesIDsInControl(&control); err != nil {
+			return nil, err
+		}
+
+		attackTrackControlsList = append(attackTrackControlsList, control)
+	}
+
+	return attackTrackControlsList, nil
+}
+
 func (gs *GitRegoStore) GetOPAControlsNamesList() ([]string, error) {
 	gs.controlsLock.RLock()
 	defer gs.controlsLock.RUnlock()

gitregostore/gitstoremethods_test.go

@@ -93,6 +93,16 @@ func gs_tests(t *testing.T, gs *GitRegoStore) {
 		)
 	})
 
+	t.Run("should retrieve OPA Attack Track controls", func(t *testing.T) {
+		t.Parallel()
+
+		controls, err := gs.GetOPAAttackTrackControls()
+		assert.NoError(t, err)
+		assert.NotEmptyf(t, controls,
+			"failed to get attack track controls %v", err,
+		)
+	})
+
 	t.Run("should retrieve OPA control by id or name", func(t *testing.T) {
 		t.Parallel()
 

gitregostore/gitstoreutils.go

@@ -195,6 +195,7 @@ func (gs *GitRegoStore) setAttackTracks(respStr string) error {
 	return nil
 }
 
+// Set controls set the controls list and attackTrackControls in gitRegoStore
 func (gs *GitRegoStore) setControls(respStr string) error {
 	controls := []opapolicy.Control{}
 	if err := JSONDecoder(respStr).Decode(&controls); err != nil {
@@ -204,6 +205,24 @@ func (gs *GitRegoStore) setControls(respStr string) error {
 	defer gs.controlsLock.Unlock()
 
 	gs.Controls = controls
+	gs.setAttackTracksControls()
+	return nil
+}
+
+// GetAttackTracksControls sets controls that are related to attack tracks
+func (gs *GitRegoStore) setAttackTracksControls() error {
+	allAttackTrackControls := []opapolicy.Control{}
+
+	for i, control := range gs.Controls {
+		controlCategories := control.GetAllAttackTrackCategories()
+		if controlCategories != nil && len(controlCategories) > 0 {
+			allAttackTrackControls = append(allAttackTrackControls, gs.Controls[i])
+		}
+	}
+	gs.attackTrackControlsLock.Lock()
+	defer gs.attackTrackControlsLock.Unlock()
+	gs.AttackTrackControls = allAttackTrackControls
+
 	return nil
 }
 

gitregostore/gitstoreutils_test.go

@@ -317,3 +317,67 @@ func TestHasNumbers(t *testing.T) {
 		}
 	}
 }
+
+func TestSetControls(t *testing.T) {
+	store := &GitRegoStore{}
+	// Successful test case
+	respStr := `[{"name": "control1"}, {"name": "control2"}]`
+	err := store.setControls(respStr)
+	if err != nil {
+		t.Errorf("Error setting controls: %v", err)
+	}
+	if len(store.Controls) != 2 {
+		t.Errorf("Controls not added to store")
+	}
+	if len(store.AttackTrackControls) != 0 {
+		t.Errorf("Attack track controls not added to store")
+	}
+
+	// Error test case
+	respStr = `invalid JSON`
+	expectedErr := errors.New("invalid character 'i' looking for beginning of value")
+	err = store.setControls(respStr)
+	if err == nil || err.Error() != expectedErr.Error() {
+		t.Errorf("Expected error '%v', but got: %v", expectedErr, err)
+	}
+	if len(store.Controls) != 2 {
+		t.Errorf("Expected 2 controls, but got: %d", len(store.Controls))
+	}
+
+	//
+	respStr = `[{"name":"TEST","attributes":{"armoBuiltin":true,"controlTypeTags":["security","compliance"],"attackTracks":[{"attackTrack": "container","categories": ["Execution","Initial access"]},{"attackTrack": "network","categories": ["Eavesdropping","Spoofing"]}]},"description":"","remediation":"","rulesNames":["CVE-2022-0185"],"id":"C-0079","long_description":"","test":"","controlID":"C-0079","baseScore":4,"example":""}]`
+	err = store.setControls(respStr)
+	if err != nil {
+		t.Errorf("Error setting controls: %v", err)
+	}
+	if len(store.Controls) != 1 {
+		t.Errorf("Controls not added to store")
+	}
+	if len(store.AttackTrackControls) != 1 {
+		t.Errorf("Attack track controls not added to store")
+	}
+}
+
+func TestSetAttackTracks(t *testing.T) {
+	store := &GitRegoStore{}
+	// Successful test case
+	respStr := `[{"name": "attack_track1"}, {"name": "attack_track2"}]`
+	err := store.setAttackTracks(respStr)
+	if err != nil {
+		t.Errorf("Error setting attack tracks: %v", err)
+	}
+	if len(store.AttackTracks) != 2 {
+		t.Errorf("Attack tracks added to store")
+	}
+
+	// Error test case
+	respStr = `invalid JSON`
+	expectedErr := errors.New("invalid character 'i' looking for beginning of value")
+	err = store.setAttackTracks(respStr)
+	if err == nil || err.Error() != expectedErr.Error() {
+		t.Errorf("Expected error '%v', but got: %v", expectedErr, err)
+	}
+	if len(store.AttackTracks) != 2 {
+		t.Errorf("Expected 2 attack tracks, but got: %d", len(store.AttackTracks))
+	}
+}