Merge pull request #24 from armosec/add_container_to_attack_chains

add container name to attack chains vuls
armosec · Sep 7, 2023 · b43f066 · b43f066
2 parents 7a73855 + 8ba21ee
commit b43f066
Show file tree

Hide file tree

Showing 7 changed files with 107 additions and 91 deletions.
diff --git a/attackchains/attackchainhandler.go b/attackchains/attackchainhandler.go
@@ -67,7 +67,7 @@ func (h *AttackChainsEngine) detectSingleAttackChain(attackTrack v1alpha1.IAttac
 }
 
 // getAttackTrackControlsLookup returns a lookup of all the controls that are relevant to the attack tracks
-func (h *AttackChainsEngine) getAttackTrackControlsLookup(postureResourceSummary *armotypes.PostureResourceSummary, vul *cscanlib.CommonContainerScanSummaryResult) (v1alpha1.AttackTrackControlsLookup, error) {
+func (h *AttackChainsEngine) getAttackTrackControlsLookup(postureResourceSummary *armotypes.PostureResourceSummary, vuls []*cscanlib.CommonContainerScanSummaryResult) (v1alpha1.AttackTrackControlsLookup, error) {
 
 	relevantControls, err := h.getRelevantControls(postureResourceSummary)
 	if err != nil {
@@ -83,18 +83,17 @@ func (h *AttackChainsEngine) getAttackTrackControlsLookup(postureResourceSummary
 		return nil, err
 	}
 
-	// If the vulnarable image is relevant to the attack chain, add it as a control
-	vulRelevant, err := isVulnerableRelevantToAttackChain(vul)
-	if err != nil {
-		return nil, err
-	}
+	vulRelevants := make([]*cscanlib.CommonContainerScanSummaryResult, 0, len(vuls))
 
-	if vulRelevant {
+	for _, vul := range vuls {
+		if isVulnerableRelevantToAttackChain(vul) {
+			vulRelevants = append(vulRelevants, vul)
 
-		// Convert the vulnarable image to a control structure
-		volAsControl := convertVulToControl(vul, []string{securityFrameworkName}, attackTracks)
-		if volAsControl != nil {
-			relevantControls[volAsControl.ControlID] = volAsControl
+			// Convert the vulnarable image to a control structure
+			volAsControl := convertVulToControl(vul, []string{securityFrameworkName}, attackTracks)
+			if volAsControl != nil {
+				relevantControls[volAsControl.ControlID] = volAsControl
+			}
 		}
 	}
 
@@ -111,7 +110,7 @@ func (h *AttackChainsEngine) getAttackTrackControlsLookup(postureResourceSummary
 }
 
 // DetectAllAttackChains - Detects all the attack chains that are relevant to the postureResourceSummary
-func (h *AttackChainsEngine) DetectAllAttackChains(postureResourceSummary *armotypes.PostureResourceSummary, vul *cscanlib.CommonContainerScanSummaryResult) ([]v1alpha1.IAttackTrack, error) {
+func (h *AttackChainsEngine) DetectAllAttackChains(postureResourceSummary *armotypes.PostureResourceSummary, vul []*cscanlib.CommonContainerScanSummaryResult) ([]v1alpha1.IAttackTrack, error) {
 
 	attackChains := []v1alpha1.IAttackTrack{}
 
@@ -153,32 +152,6 @@ func (h *AttackChainsEngine) DetectAllAttackChains(postureResourceSummary *armot
 
 }
 
-func (h *AttackChainsEngine) DetectAllAttackChainsFromLists(postureResourceSummaries []*armotypes.PostureResourceSummary, vuls []*cscanlib.CommonContainerScanSummaryResult) ([]*armotypes.AttackChain, error) {
-	var attackChains []*armotypes.AttackChain
-
-	for i := range postureResourceSummaries {
-		for j := range vuls {
-			// validate that the vulnarable image matches to the postureResourceSummary}
-			// ignoring the error, if they don't match they won't create an attack chain
-			if validateWorkLoadMatch(postureResourceSummaries[i], vuls[j]) {
-				attackTracks, err := h.DetectAllAttackChains(postureResourceSummaries[i], vuls[j])
-				if err != nil {
-					return nil, err
-				}
-
-				if len(attackTracks) == 0 {
-					continue
-				}
-				prsAttributes := postureResourceSummaries[i].Designators.Attributes
-				currentAttackChains := ConvertAttackTracksToAttackChains(attackTracks, prsAttributes, postureResourceSummaries[i].ResourceID, postureResourceSummaries[i].ReportID)
-				attackChains = append(attackChains, currentAttackChains...)
-			}
-		}
-	}
-
-	return attackChains, nil
-}
-
 // GetAttackTrack - Returns all the attack tracks
 func (h *AttackChainsEngine) GetAttackTrack() ([]v1alpha1.IAttackTrack, error) {
 	if len(h.attackTracks) == 0 {

diff --git a/attackchains/attackchainhandler_mocks.go b/attackchains/attackchainhandler_mocks.go
@@ -190,9 +190,9 @@ func GetAttackTrackMocks() ([]v1alpha1.AttackTrack, error) {
 
 	attackTracks := []v1alpha1.AttackTrack{}
 
-	for _, attackTrackMock := range attackTracksMocks {
+	for i := range attackTracksMocks {
 		attackTrack := &v1alpha1.AttackTrack{}
-		err := json.Unmarshal([]byte(attackTrackMock), &attackTrack)
+		err := json.Unmarshal([]byte(attackTracksMocks[i]), &attackTrack)
 
 		if err != nil {
 			return nil, err

diff --git a/attackchains/attackchainhandler_test.go b/attackchains/attackchainhandler_test.go
@@ -4,15 +4,11 @@ import (
 	"testing"
 
 	"github.com/armosec/armoapi-go/containerscan"
+	cscanlib "github.com/armosec/armoapi-go/containerscan"
 	"github.com/kubescape/opa-utils/reporthandling/attacktrack/v1alpha1"
 	"github.com/stretchr/testify/assert"
 )
 
-type vulnerabilityTest struct {
-	HasRelevancyData bool
-	RelevantLabel    containerscan.RelevantLabel
-}
-
 var allControls = AllControlsMock()
 
 func TestDetectSingleAttackTrack(t *testing.T) {
@@ -53,18 +49,18 @@ func TestDetectSingleAttackTrack(t *testing.T) {
 		attackTrack     v1alpha1.IAttackTrack
 		FailedControls  []string
 		WarningControls []string
-		Vul             vulnerabilityTest
+		Vuls            []*cscanlib.CommonContainerScanSummaryResult
 		Expected        v1alpha1.IAttackTrack
 	}{
 		{
 			name:            "Attack chain exists with vulnarable image 'yes'",
 			attackTrack:     attackTrack1,
 			FailedControls:  []string{"control1", "control2"},
 			WarningControls: []string{"control3", "control4"},
-			Vul: vulnerabilityTest{
-				HasRelevancyData: true,
-				RelevantLabel:    "yes",
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(true, containerscan.RelevantLabelYes, Attributes),
 			},
+
 			Expected: AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
 				Name: "A",
 				SubSteps: []v1alpha1.AttackTrackStep{
@@ -85,10 +81,35 @@ func TestDetectSingleAttackTrack(t *testing.T) {
 			attackTrack:     attackTrack1,
 			FailedControls:  []string{"control1", "control2"},
 			WarningControls: []string{"control3", "control4"},
-			Vul: vulnerabilityTest{
-				HasRelevancyData: false,
-				RelevantLabel:    "",
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(false, containerscan.RelevantLabelYes, Attributes),
 			},
+
+			Expected: AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
+				Name: "A",
+				SubSteps: []v1alpha1.AttackTrackStep{
+					{
+						ChecksVulnerabilities: true,
+						Name:                  "vulnerableImageStepName",
+						SubSteps: []v1alpha1.AttackTrackStep{
+							{
+								Name: "C",
+							},
+						},
+					},
+				},
+			}),
+		},
+		{
+			name:            "Attack chain exists with 1 vulnarable image out of two",
+			attackTrack:     attackTrack1,
+			FailedControls:  []string{"control1", "control2"},
+			WarningControls: []string{"control3", "control4"},
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(true, containerscan.RelevantLabelNo, Attributes),
+				CommonContainerScanSummaryResultMock(false, containerscan.RelevantLabelYes, Attributes),
+			},
+
 			Expected: AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
 				Name: "A",
 				SubSteps: []v1alpha1.AttackTrackStep{
@@ -109,21 +130,21 @@ func TestDetectSingleAttackTrack(t *testing.T) {
 			attackTrack:     attackTrack1,
 			FailedControls:  []string{"control1", "control2"},
 			WarningControls: []string{"control3", "control4"},
-			Vul: vulnerabilityTest{
-				HasRelevancyData: true,
-				RelevantLabel:    "no",
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(true, containerscan.RelevantLabelNo, Attributes),
 			},
+
 			Expected: nil,
 		},
 		{
 			name:            "Attack Chain exists, no vulnarable image",
 			attackTrack:     attackTrack1,
 			FailedControls:  []string{"control1", "control2"},
 			WarningControls: []string{"control6"},
-			Vul: vulnerabilityTest{
-				HasRelevancyData: true,
-				RelevantLabel:    "no",
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(true, containerscan.RelevantLabelNo, Attributes),
 			},
+
 			Expected: AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
 				Name: "A",
 				SubSteps: []v1alpha1.AttackTrackStep{
@@ -143,10 +164,10 @@ func TestDetectSingleAttackTrack(t *testing.T) {
 			attackTrack:     attackTrack1,
 			FailedControls:  []string{"control1", "control2", "control6"},
 			WarningControls: []string{"control3", "control4"},
-			Vul: vulnerabilityTest{
-				HasRelevancyData: true,
-				RelevantLabel:    "yes",
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(true, containerscan.RelevantLabelYes, Attributes),
 			},
+
 			Expected: AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
 				Name: "A",
 				SubSteps: []v1alpha1.AttackTrackStep{
@@ -180,9 +201,7 @@ func TestDetectSingleAttackTrack(t *testing.T) {
 
 			postureResourcesSummary := PostureResourcesSummaryMock(Attributes, test.FailedControls, test.WarningControls)
 
-			commonContainerScanSummaryResult := CommonContainerScanSummaryResultMock(test.Vul.HasRelevancyData, test.Vul.RelevantLabel, Attributes)
-
-			controlsLookup, err := attackChainHandler.getAttackTrackControlsLookup(postureResourcesSummary, commonContainerScanSummaryResult)
+			controlsLookup, err := attackChainHandler.getAttackTrackControlsLookup(postureResourcesSummary, test.Vuls)
 			assert.NoError(t, err)
 			attackChain, err := attackChainHandler.detectSingleAttackChain(test.attackTrack, controlsLookup)
 			assert.NoError(t, err)
@@ -205,6 +224,11 @@ func TestDetectSingleAttackTrack(t *testing.T) {
 
 func TestDetectAllAttackChains(t *testing.T) {
 
+	Attributes := map[string]string{"cluster": "minikubesecurity1",
+		"kind":      "Pod",
+		"name":      "wowtest",
+		"namespace": "default"}
+
 	attackTrack1 := AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
 		Name: "A",
 		SubSteps: []v1alpha1.AttackTrackStep{
@@ -251,18 +275,18 @@ func TestDetectAllAttackChains(t *testing.T) {
 		attackTracks         []v1alpha1.IAttackTrack
 		FailedControls       []string
 		WarningControls      []string
-		Vul                  vulnerabilityTest
+		Vuls                 []*cscanlib.CommonContainerScanSummaryResult
 		ExpectedAttackTracks []v1alpha1.IAttackTrack
 	}{
 		{
 			name:            "Found one attack chain",
 			attackTracks:    []v1alpha1.IAttackTrack{attackTrack1},
 			FailedControls:  []string{"control1", "control2"},
 			WarningControls: []string{"control3", "control4"},
-			Vul: vulnerabilityTest{
-				HasRelevancyData: true,
-				RelevantLabel:    "yes",
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(true, containerscan.RelevantLabelYes, Attributes),
 			},
+
 			ExpectedAttackTracks: []v1alpha1.IAttackTrack{AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
 				Name: "A",
 				SubSteps: []v1alpha1.AttackTrackStep{
@@ -283,10 +307,10 @@ func TestDetectAllAttackChains(t *testing.T) {
 			attackTracks:    []v1alpha1.IAttackTrack{attackTrack1, attackTrack2},
 			FailedControls:  []string{"control1", "control2", "control6"},
 			WarningControls: []string{"control3", "control4"},
-			Vul: vulnerabilityTest{
-				HasRelevancyData: true,
-				RelevantLabel:    "yes",
+			Vuls: []*cscanlib.CommonContainerScanSummaryResult{
+				CommonContainerScanSummaryResultMock(true, containerscan.RelevantLabelYes, Attributes),
 			},
+
 			ExpectedAttackTracks: []v1alpha1.IAttackTrack{AttackTrackMock("attackchain1", v1alpha1.AttackTrackStep{
 				Name: "A",
 				SubSteps: []v1alpha1.AttackTrackStep{
@@ -335,9 +359,7 @@ func TestDetectAllAttackChains(t *testing.T) {
 
 			postureResourcesSummary := PostureResourcesSummaryMock(Attributes, test.FailedControls, test.WarningControls)
 
-			commonContainerScanSummaryResult := CommonContainerScanSummaryResultMock(test.Vul.HasRelevancyData, test.Vul.RelevantLabel, Attributes)
-
-			attackChains, err := attackChainHandler.DetectAllAttackChains(postureResourcesSummary, commonContainerScanSummaryResult)
+			attackChains, err := attackChainHandler.DetectAllAttackChains(postureResourcesSummary, test.Vuls)
 
 			assert.NoError(t, err)
 			assert.Equalf(t, len(test.ExpectedAttackTracks), len(attackChains), "Expected and actual attack chains are not equal")

diff --git a/attackchains/attackchainutils.go b/attackchains/attackchainutils.go
@@ -50,7 +50,7 @@ func convertVulToControl(vul *cscanlib.CommonContainerScanSummaryResult, tags []
 	}
 
 	return &reporthandling.Control{
-		ControlID: vul.ImageID,
+		ControlID: vul.ImageID + vul.ContainerName,
 		PortalBase: armotypes.PortalBase{
 			Attributes: map[string]interface{}{
 				"controlTypeTags":                    tags,
@@ -64,12 +64,12 @@ func convertVulToControl(vul *cscanlib.CommonContainerScanSummaryResult, tags []
 }
 
 // isVulnerableRelevantToAttackChain checks if the vulnerability is relevant to the attack chain
-func isVulnerableRelevantToAttackChain(vul *cscanlib.CommonContainerScanSummaryResult) (bool, error) {
+func isVulnerableRelevantToAttackChain(vul *cscanlib.CommonContainerScanSummaryResult) bool {
 	// validate relevancy
 	if !vul.HasRelevancyData || (vul.HasRelevancyData && vul.RelevantLabel == "yes") {
 		//validate severity
 		if vul.Severity == "Critical" {
-			return true, nil
+			return true
 		}
 
 		// TODO: figure out how to handle empty severity stats
@@ -79,11 +79,11 @@ func isVulnerableRelevantToAttackChain(vul *cscanlib.CommonContainerScanSummaryR
 
 		for _, stat := range vul.SeveritiesStats {
 			if stat.Severity == "Critical" && stat.TotalCount > 0 {
-				return true, nil
+				return true
 			}
 		}
 	}
-	return false, nil
+	return false
 }
 
 // validateWorkLoadMatch checks if the vulnerability and the posture resource summary are of the same workload
@@ -156,7 +156,10 @@ func ConvertAttackTrackStepToAttackChainNode(step v1alpha1.IAttackTrackStep) *ar
 				}
 			}
 
-			imageVulnerabilities = append(imageVulnerabilities, armotypes.Vulnerabilities{ImageScanID: containerScanID, Names: vulNames})
+			imageVulnerabilities = append(imageVulnerabilities, armotypes.Vulnerabilities{
+				ContainerName: vulControl.(*reporthandling.Control).Attributes[identifiers.AttributeContainerName].(string),
+				ImageScanID:   containerScanID,
+				Names:         vulNames})
 
 		}
 	} else {