Skip to content

Commit

Permalink
Support checks vuls (#119)
Browse files Browse the repository at this point in the history
* support GetSubstepsWithVulnerabilities

* minor

* minor fix
  • Loading branch information
kooomix authored Jul 17, 2023
1 parent c09ae47 commit 8334eca
Show file tree
Hide file tree
Showing 4 changed files with 87 additions and 0 deletions.
19 changes: 19 additions & 0 deletions reporthandling/attacktrack/v1alpha1/attacktrackmethods.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,25 @@ func (at *AttackTrack) Iterator() IAttackTrackIterator {
}
}

// GetSubstepsWithVulnerabilities returns a list of substeps names that check for vulnerabilities
func (at *AttackTrack) GetSubstepsWithVulnerabilities() []string {
var substepNames []string

var traverse func(step AttackTrackStep)
traverse = func(step AttackTrackStep) {
if step.DoesCheckVulnerabilities() {
substepNames = append(substepNames, step.Name)
}
for _, substep := range step.SubSteps {
traverse(substep)
}
}

traverse(at.Spec.Data)

return substepNames
}

func (iter *AttackTrackIterator) HasNext() bool {
return !iter.stack.IsEmpty()
}
Expand Down
47 changes: 47 additions & 0 deletions reporthandling/attacktrack/v1alpha1/attacktrackmethods_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -685,3 +685,50 @@ func TestFilterNodesWithControls(t *testing.T) {
})
}
}

func TestGetSubstepsWithVulnerabilities(t *testing.T) {
// Create an AttackTrack object with substeps having different values for ChecksVulnerabilities
attackTrack := AttackTrack{
ApiVersion: "v1",
Kind: "AttackTrack",
Metadata: map[string]interface{}{},
Spec: AttackTrackSpecification{
Version: "1.0",
Description: "Example attack track",
Data: AttackTrackStep{
Name: "Step 1",
Description: "First step",
ChecksVulnerabilities: true,
SubSteps: []AttackTrackStep{
{
Name: "Substep 1.1",
Description: "Substep 1.1 description",
ChecksVulnerabilities: true,
},
{
Name: "Substep 1.2",
Description: "Substep 1.2 description",
ChecksVulnerabilities: false,
},
},
},
},
}

// Call the method being tested
substepNames := attackTrack.GetSubstepsWithVulnerabilities()

// Define the expected substep names with ChecksVulnerabilities set to true
expectedSubstepNames := []string{"Step 1", "Substep 1.1"}

// Check if the returned substep names match the expected substep names
if len(substepNames) != len(expectedSubstepNames) {
t.Errorf("Unexpected number of substep names. Expected: %d, Got: %d", len(expectedSubstepNames), len(substepNames))
}

for i, name := range substepNames {
if name != expectedSubstepNames[i] {
t.Errorf("Mismatched substep name. Expected: %s, Got: %s", expectedSubstepNames[i], name)
}
}
}
20 changes: 20 additions & 0 deletions reporthandling/attacktrack/v1alpha1/attacktrackmocks.go
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,26 @@ func (at AttackTrackMock) Iterator() IAttackTrackIterator {
}
}

// GetSubstepsWithVulnerabilities returns a list of substeps names that check for vulnerabilities
func (at AttackTrackMock) GetSubstepsWithVulnerabilities() []string {
var substepNames []string

var traverse func(step AttackTrackStep)
traverse = func(step AttackTrackStep) {
if step.DoesCheckVulnerabilities() {
substepNames = append(substepNames, step.Name)
}
for _, substep := range step.SubSteps {
traverse(substep)
}
}

t := at.Spec.Data.(*AttackTrackStep)
traverse(*t)

return substepNames
}

type MockAttackTrackSpecification struct {
Version string `json:"version,omitempty"`
Description string `json:"description,omitempty"`
Expand Down
1 change: 1 addition & 0 deletions reporthandling/attacktrack/v1alpha1/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ type IAttackTrack interface {
GetData() IAttackTrackStep
Iterator() IAttackTrackIterator
IsValid() bool
GetSubstepsWithVulnerabilities() []string
}

// A step in an attack track
Expand Down

0 comments on commit 8334eca

Please sign in to comment.