From f2007ec98da3708d4ed98ea08e4271e4e0e23c0e Mon Sep 17 00:00:00 2001 From: YiscahLevySilas1 <80635572+YiscahLevySilas1@users.noreply.github.com> Date: Thu, 10 Aug 2023 12:31:07 +0300 Subject: [PATCH] Add GetAllAttackTrackCategories func for control (#135) * add indicator of IsVulnerabilityStep Signed-off-by: YiscahLevySilas1 * name changes Signed-off-by: YiscahLevySilas1 * changes following review Signed-off-by: YiscahLevySilas1 * add GetAllAttackTrackCategories func for control Signed-off-by: YiscahLevySilas1 --------- Signed-off-by: YiscahLevySilas1 --- reporthandling/datastructuresmethods.go | 14 +++++++ reporthandling/datastructuresmethods_test.go | 44 ++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/reporthandling/datastructuresmethods.go b/reporthandling/datastructuresmethods.go index d6914029..a7094a9b 100644 --- a/reporthandling/datastructuresmethods.go +++ b/reporthandling/datastructuresmethods.go @@ -403,6 +403,20 @@ func (ruleResponse *RuleResponse) RemoveData(keepFields, keepMetadataFields []st } } +// GetAllAttackTracks returns all attack tracks of the control +// if the control has no attack tracks, it returns nil +func (control *Control) GetAllAttackTrackCategories() []AttackTrackCategories { + if v, exist := control.Attributes[ControlAttributeKeyAttackTracks]; exist { + var attackTrackToCategories []AttackTrackCategories + if err := mapstructure.Decode(v, &attackTrackToCategories); err == nil { + if len(attackTrackToCategories) > 0 { + return attackTrackToCategories + } + } + } + return nil +} + func (control *Control) GetAttackTrackCategories(attackTrackName string) []string { if v, exist := control.Attributes[ControlAttributeKeyAttackTracks]; exist { var attackTrackToCategories []AttackTrackCategories diff --git a/reporthandling/datastructuresmethods_test.go b/reporthandling/datastructuresmethods_test.go index 50f342d4..509c03ad 100644 --- a/reporthandling/datastructuresmethods_test.go +++ b/reporthandling/datastructuresmethods_test.go @@ -135,6 +135,50 @@ func TestControl_GetAttackTrackCategories(t *testing.T) { assert.Equal(t, []string{}, invalidControl2.GetAttackTrackCategories("container")) } +func TestControl_GetAllAttackTracks(t *testing.T) { + validControlJson := `{"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":""}` + var validControl Control + err := json.Unmarshal([]byte(validControlJson), &validControl) + assert.NoError(t, err) + expectedTracks := []AttackTrackCategories{ + { + AttackTrack: "container", + Categories: []string{"Execution", "Initial access"}, + }, + { + AttackTrack: "network", + Categories: []string{"Eavesdropping", "Spoofing"}, + }, + } + assert.ElementsMatch(t, expectedTracks, validControl.GetAllAttackTrackCategories()) + + invalidControlJson1 := `{"name":"TEST","attributes":{"armoBuiltin":true,"controlTypeTags":["security","compliance"],"attackTracks":{"container": "x"}},"description":"","remediation":"","rulesNames":["CVE-2022-0185"],"id":"C-0079","long_description":"","test":"","controlID":"C-0079","baseScore":4,"example":""}` + var invalidControl1 Control + err = json.Unmarshal([]byte(invalidControlJson1), &invalidControl1) + assert.NoError(t, err) + assert.Nil(t, invalidControl1.GetAllAttackTrackCategories()) + + invalidControlJson2 := `{"name":"TEST","attributes":{"armoBuiltin":true,"controlTypeTags":["security","compliance"],"attack":{"container": "x"}},"description":"","remediation":"","rulesNames":["CVE-2022-0185"],"id":"C-0079","long_description":"","test":"","controlID":"C-0079","baseScore":4,"example":""}` + var invalidControl2 Control + err = json.Unmarshal([]byte(invalidControlJson2), &invalidControl2) + assert.NoError(t, err) + assert.Nil(t, invalidControl2.GetAllAttackTrackCategories()) + + // Case: control that has no "attackTracks" field + missingAttackTrackJson := `{"name":"TEST","attributes":{"armoBuiltin":true,"controlTypeTags":["security","compliance"]},"description":"","remediation":"","rulesNames":["CVE-2022-0185"],"id":"C-0079","long_description":"","test":"","controlID":"C-0079","baseScore":4,"example":""}` + var missingAttackTrackControl Control + err = json.Unmarshal([]byte(missingAttackTrackJson), &missingAttackTrackControl) + assert.NoError(t, err) + assert.Nil(t, missingAttackTrackControl.GetAllAttackTrackCategories()) + + // Case: control that has "attackTracks" but it's an empty list + emptyAttackTrackJson := `{"name":"TEST","attributes":{"armoBuiltin":true,"controlTypeTags":["security","compliance"],"attackTracks":[]},"description":"","remediation":"","rulesNames":["CVE-2022-0185"],"id":"C-0079","long_description":"","test":"","controlID":"C-0079","baseScore":4,"example":""}` + var emptyAttackTrackControl Control + err = json.Unmarshal([]byte(emptyAttackTrackJson), &emptyAttackTrackControl) + assert.NoError(t, err) + assert.Nil(t, emptyAttackTrackControl.GetAllAttackTrackCategories()) +} + func TestControl_GetControlTypeTags(t *testing.T) { validControlJson := `{"name":"TEST","attributes":{"armoBuiltin":true,"controlTypeTags":["security","compliance"],"attackTracks":{"container":["Privilege escalation"]}},"description":"","remediation":"","rulesNames":["CVE-2022-0185"],"id":"C-0079","long_description":"","test":"","controlID":"C-0079","baseScore":4,"example":""}` var validControl Control