From 4f71eff4194e52a23cab7af3caeb4182eb2d2992 Mon Sep 17 00:00:00 2001 From: kooomix Date: Wed, 12 Jul 2023 12:15:08 +0300 Subject: [PATCH] fixed following review --- .../attacktrack/v1alpha1/attacktrackmethods.go | 14 +++++++++----- .../v1alpha1/attacktrackmethods_test.go | 8 +++----- .../attacktrack/v1alpha1/attacktrackmocks.go | 4 ++++ .../attacktrack/v1alpha1/attacktrackstepmethods.go | 6 +++++- reporthandling/attacktrack/v1alpha1/interface.go | 1 + 5 files changed, 22 insertions(+), 11 deletions(-) diff --git a/reporthandling/attacktrack/v1alpha1/attacktrackmethods.go b/reporthandling/attacktrack/v1alpha1/attacktrackmethods.go index b9ae10b7..e70bace9 100644 --- a/reporthandling/attacktrack/v1alpha1/attacktrackmethods.go +++ b/reporthandling/attacktrack/v1alpha1/attacktrackmethods.go @@ -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) @@ -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 { @@ -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 } diff --git a/reporthandling/attacktrack/v1alpha1/attacktrackmethods_test.go b/reporthandling/attacktrack/v1alpha1/attacktrackmethods_test.go index 934769be..d1ab22f0 100644 --- a/reporthandling/attacktrack/v1alpha1/attacktrackmethods_test.go +++ b/reporthandling/attacktrack/v1alpha1/attacktrackmethods_test.go @@ -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) } } @@ -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) } }) diff --git a/reporthandling/attacktrack/v1alpha1/attacktrackmocks.go b/reporthandling/attacktrack/v1alpha1/attacktrackmocks.go index 80c75502..5821d68b 100644 --- a/reporthandling/attacktrack/v1alpha1/attacktrackmocks.go +++ b/reporthandling/attacktrack/v1alpha1/attacktrackmocks.go @@ -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 } diff --git a/reporthandling/attacktrack/v1alpha1/attacktrackstepmethods.go b/reporthandling/attacktrack/v1alpha1/attacktrackstepmethods.go index 294a0200..84e44fb6 100644 --- a/reporthandling/attacktrack/v1alpha1/attacktrackstepmethods.go +++ b/reporthandling/attacktrack/v1alpha1/attacktrackstepmethods.go @@ -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 { @@ -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 } } diff --git a/reporthandling/attacktrack/v1alpha1/interface.go b/reporthandling/attacktrack/v1alpha1/interface.go index b0f43486..3580997a 100644 --- a/reporthandling/attacktrack/v1alpha1/interface.go +++ b/reporthandling/attacktrack/v1alpha1/interface.go @@ -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