Skip to content

Commit

Permalink
fixed following review
Browse files Browse the repository at this point in the history
  • Loading branch information
kooomix committed Jul 12, 2023
1 parent 0b103fd commit 4f71eff
Show file tree
Hide file tree
Showing 5 changed files with 22 additions and 11 deletions.
14 changes: 9 additions & 5 deletions reporthandling/attacktrack/v1alpha1/attacktrackmethods.go
Original file line number Diff line number Diff line change
Expand Up @@ -192,9 +192,9 @@ func (handler *AttackTrackAllPathsHandler) CalculatePathsRootToLeaf() [][]IAttac
currentPath = append(currentPath, step)
}

if step.Length() == 0 {
if step.IsLeaf() {
// Reached a leaf node
if len(step.GetControls()) > 0 {
if step.IsPartOfAttackTrackPath() {
// Add current path to paths only if controls are not empty
path := make([]IAttackTrackStep, len(currentPath))
copy(path, currentPath)
Expand All @@ -213,7 +213,7 @@ func (handler *AttackTrackAllPathsHandler) CalculatePathsRootToLeaf() [][]IAttac
handler.visited[subStep.GetName()] = true

// Only include nodes with controls in the path
if len(subStep.GetControls()) > 0 {
if step.IsPartOfAttackTrackPath() {
traverse(subStep)
traversedSubstepWithControls = true
} else {
Expand All @@ -225,20 +225,24 @@ func (handler *AttackTrackAllPathsHandler) CalculatePathsRootToLeaf() [][]IAttac
}

// Add the current path to paths only if it ends with a leaf node with controls
if len(step.GetControls()) > 0 && !traversedSubstepWithControls && step.Length() == 0 {
if step.IsPartOfAttackTrackPath() && !traversedSubstepWithControls && step.IsLeaf() {
path := make([]IAttackTrackStep, len(currentPath))
copy(path, currentPath)
paths = append(paths, path)
}
}

if len(step.GetControls()) > 0 {
if step.IsPartOfAttackTrackPath() {
currentPath = currentPath[:len(currentPath)-1] // Remove last step from current path in order to explore other paths
}
}

traverse(handler.attackTrack.GetData())

if len(paths) == 0 {
return nil
}

return paths
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -406,13 +406,12 @@ func TestCalculatePathsRootToLeaf(t *testing.T) {
t.Run(tt.name, func(t *testing.T) {
handler := NewAttackTrackAllPathsHandler(tt.attackTrack, &tt.controlsMap)
paths := handler.CalculatePathsRootToLeaf()
// newAttackTrack := handler.GenerateAttackTrackFromPaths(paths)

if !(tt.want == nil && paths == nil) {
assert.Equalf(t, len(tt.want), len(paths), "CalculatePathsToLeaf should return the correct number of paths. expected: %v, actual: %v", len(tt.want), len(paths))
if tt.want != nil || paths != nil {
assert.Equalf(t, len(tt.want), len(paths), "CalculatePathsRootToLeaf should return the correct number of paths. expected: %v, actual: %v", len(tt.want), len(paths))
for i, path := range paths {
for j, step := range path {
assert.Equalf(t, tt.want[i][j], step.GetName(), "CalculatePathsToLeaf should return the correct paths. expected: %v, actual: %v", tt.want, paths)
assert.Equalf(t, tt.want[i][j], step.GetName(), "CalculatePathsRootToLeaf should return the correct paths. expected: %v, actual: %v", tt.want, paths)
}
}

Expand Down Expand Up @@ -681,7 +680,6 @@ func TestFilterNodesWithControls(t *testing.T) {
result := handler.filterNodesWithControls(handler.attackTrack.GetData(), pathsCopy)

if !(result == nil && tc.expectedResult == nil) && result.Equal(tc.expectedResult, true) == false {
// if !compareAttackTrackStep(result, tc.expectedResult) {
t.Errorf("Unexpected result.\nExpected: %+v\nGot: %+v", tc.expectedResult, result)
}
})
Expand Down
4 changes: 4 additions & 0 deletions reporthandling/attacktrack/v1alpha1/attacktrackmocks.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,10 @@ func (s AttackTrackStepMock) SubStepAt(index int) IAttackTrackStep {
return s.SubSteps[index]
}

func (s AttackTrackStepMock) IsLeaf() bool {
return len(s.SubSteps) == 0
}

func (a AttackTrackStepMock) IsPartOfAttackTrackPath() bool {
return len(a.Controls) > 0
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@ func (step *AttackTrackStep) SubStepAt(index int) IAttackTrackStep {
return &step.SubSteps[index]
}

func (step *AttackTrackStep) IsLeaf() bool {
return step.Length() == 0
}

// Equal checks if the given attack track step is equal to the current one
// If compareControls is true, the controls are also compared
func (s *AttackTrackStep) Equal(other *AttackTrackStep, compareControls bool) bool {
Expand All @@ -50,7 +54,7 @@ func (s *AttackTrackStep) Equal(other *AttackTrackStep, compareControls bool) bo

for i := range s.Controls {

if !(s.Controls[i] == other.Controls[i]) {
if s.Controls[i] != other.Controls[i] {
return false
}
}
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 @@ -20,6 +20,7 @@ type IAttackTrackStep interface {
Length() int // returns the number of sub steps
SubStepAt(index int) IAttackTrackStep // returns a sub step at the given index
IsPartOfAttackTrackPath() bool // checks if the step can be a part of an attack track path
IsLeaf() bool // checks if the step is a leaf node
}

// A control related to an attack track step
Expand Down

0 comments on commit 4f71eff

Please sign in to comment.