Print attack tree (optional, with argument) (#997)

* Print attack tree with argument

* fix

cmd/scan/scan.go

@@ -91,12 +91,14 @@ func GetScanCommand(ks meta.IKubescape) *cobra.Command {
 	scanCmd.PersistentFlags().StringVar(&scanInfo.CustomClusterName, "cluster-name", "", "Set the custom name of the cluster. Not same as the kube-context flag")
 	scanCmd.PersistentFlags().BoolVarP(&scanInfo.Submit, "submit", "", false, "Submit the scan results to Kubescape SaaS where you can see the results in a user-friendly UI, choose your preferred compliance framework, check risk results history and trends, manage exceptions, get remediation recommendations and much more. By default the results are not submitted")
 	scanCmd.PersistentFlags().BoolVarP(&scanInfo.OmitRawResources, "omit-raw-resources", "", false, "Omit raw resources from the output. By default the raw resources are included in the output")
+	scanCmd.PersistentFlags().BoolVarP(&scanInfo.PrintAttackTree, "print-attack-tree", "", false, "Print attack tree")
 
 	scanCmd.PersistentFlags().MarkDeprecated("silent", "use '--logger' flag instead. Flag will be removed at 1.May.2022")
 
 	// hidden flags
 	scanCmd.PersistentFlags().MarkHidden("host-scan-yaml") // this flag should be used very cautiously. We prefer users will not use it at all unless the DaemonSet can not run pods on the nodes
 	scanCmd.PersistentFlags().MarkHidden("omit-raw-resources")
+	scanCmd.PersistentFlags().MarkHidden("print-attack-tree")
 
 	// Retrieve --kubeconfig flag from https://github.com/kubernetes/kubectl/blob/master/pkg/cmd/cmd.go
 	scanCmd.PersistentFlags().AddGoFlag(flag.Lookup("kubeconfig"))

core/cautils/datastructures.go

@@ -5,6 +5,7 @@ import (
 	"github.com/kubescape/k8s-interface/workloadinterface"
 	"github.com/kubescape/opa-utils/reporthandling"
 	apis "github.com/kubescape/opa-utils/reporthandling/apis"
+	"github.com/kubescape/opa-utils/reporthandling/attacktrack/v1alpha1"
 	"github.com/kubescape/opa-utils/reporthandling/results/v1/prioritization"
 	"github.com/kubescape/opa-utils/reporthandling/results/v1/resourcesresults"
 	reporthandlingv2 "github.com/kubescape/opa-utils/reporthandling/v2"
@@ -22,8 +23,10 @@ type OPASessionObj struct {
 	ResourcesResult       map[string]resourcesresults.Result            // resources scan results, map[<resource ID>]<resource result>
 	ResourceSource        map[string]reporthandling.Source              // resources sources, map[<resource ID>]<resource result>
 	ResourcesPrioritized  map[string]prioritization.PrioritizedResource // resources prioritization information, map[<resource ID>]<prioritized resource>
-	Report                *reporthandlingv2.PostureReport               // scan results v2 - Remove
-	RegoInputData         RegoInputData                                 // input passed to rego for scanning. map[<control name>][<input arguments>]
+	ResourceAttackTracks  map[string]v1alpha1.IAttackTrack              // resources attack tracks, map[<resource ID>]<attack track>
+	AttackTracks          map[string]v1alpha1.IAttackTrack
+	Report                *reporthandlingv2.PostureReport // scan results v2 - Remove
+	RegoInputData         RegoInputData                   // input passed to rego for scanning. map[<control name>][<input arguments>]
 	Metadata              *reporthandlingv2.Metadata
 	InfoMap               map[string]apis.StatusInfo         // Map errors of resources to StatusInfo
 	ResourceToControlsMap map[string][]string                // map[<apigroup/apiversion/resource>] = [<control_IDs>]

core/cautils/scaninfo.go

@@ -130,6 +130,7 @@ type ScanInfo struct {
 	FrameworkScan         bool               // false if scanning control
 	ScanAll               bool               // true if scan all frameworks
 	OmitRawResources      bool               // true if omit raw resources from the output
+	PrintAttackTree       bool               // true if print attack tree
 }
 
 type Getters struct {

core/core/initutils.go

@@ -263,8 +263,8 @@ func getAttackTracksGetter(accountID string, downloadReleasedPolicy *getter.Down
 }
 
 // getUIPrinter returns a printer that will be used to print to the program’s UI (terminal)
-func getUIPrinter(verboseMode bool, formatVersion string, viewType cautils.ViewTypes) printer.IPrinter {
-	p := printerv2.NewPrettyPrinter(verboseMode, formatVersion, viewType)
+func getUIPrinter(verboseMode bool, formatVersion string, attackTree bool, viewType cautils.ViewTypes) printer.IPrinter {
+	p := printerv2.NewPrettyPrinter(verboseMode, formatVersion, attackTree, viewType)
 
 	// Since the UI of the program is a CLI (Stdout), it means that it should always print to Stdout
 	p.SetWriter(os.Stdout.Name())

core/core/initutils_test.go

@@ -17,7 +17,7 @@ func Test_getUIPrinter(t *testing.T) {
 	wantVerboseMode := scanInfo.VerboseMode
 	wantViewType := cautils.ViewTypes(scanInfo.View)
 
-	got := getUIPrinter(scanInfo.VerboseMode, scanInfo.FormatVersion, cautils.ViewTypes(scanInfo.View))
+	got := getUIPrinter(scanInfo.VerboseMode, scanInfo.FormatVersion, scanInfo.PrintAttackTree, cautils.ViewTypes(scanInfo.View))
 
 	gotValue := reflect.ValueOf(got).Elem()
 	gotFormatVersion := gotValue.FieldByName("formatVersion").String()

core/core/scan.go

@@ -99,12 +99,12 @@ func getInterfaces(scanInfo *cautils.ScanInfo) componentInterfaces {
 
 	outputPrinters := make([]printer.IPrinter, 0)
 	for _, format := range formats {
-		printerHandler := resultshandling.NewPrinter(format, scanInfo.FormatVersion, scanInfo.VerboseMode, cautils.ViewTypes(scanInfo.View))
+		printerHandler := resultshandling.NewPrinter(format, scanInfo.FormatVersion, scanInfo.PrintAttackTree, scanInfo.VerboseMode, cautils.ViewTypes(scanInfo.View))
 		printerHandler.SetWriter(scanInfo.Output)
 		outputPrinters = append(outputPrinters, printerHandler)
 	}
 
-	uiPrinter := getUIPrinter(scanInfo.VerboseMode, scanInfo.FormatVersion, cautils.ViewTypes(scanInfo.View))
+	uiPrinter := getUIPrinter(scanInfo.VerboseMode, scanInfo.FormatVersion, scanInfo.PrintAttackTree, cautils.ViewTypes(scanInfo.View))
 
 	// ================== return interface ======================================
 
@@ -170,7 +170,7 @@ func (ks *Kubescape) Scan(scanInfo *cautils.ScanInfo) (*resultshandling.ResultsH
 
 	// ======================== prioritization ===================
 
-	if priotizationHandler, err := resourcesprioritization.NewResourcesPrioritizationHandler(scanInfo.Getters.AttackTracksGetter); err != nil {
+	if priotizationHandler, err := resourcesprioritization.NewResourcesPrioritizationHandler(scanInfo.Getters.AttackTracksGetter, scanInfo.PrintAttackTree); err != nil {
 		logger.L().Warning("failed to get attack tracks, this may affect the scanning results", helpers.Error(err))
 	} else if err := priotizationHandler.PrioritizeResources(scanData); err != nil {
 		return resultsHandling, fmt.Errorf("%w", err)

core/pkg/resourcesprioritization/prioritizationhandler.go

@@ -1,6 +1,7 @@
 package resourcesprioritization
 
 import (
+	"encoding/json"
 	"fmt"
 
 	logger "github.com/kubescape/go-logger"
@@ -13,12 +14,16 @@ import (
 )
 
 type ResourcesPrioritizationHandler struct {
-	attackTracks []v1alpha1.IAttackTrack
+	resourceToAttackTracks map[string]v1alpha1.IAttackTrack
+	attackTracks           []v1alpha1.IAttackTrack
+	buildResourcesMap      bool
 }
 
-func NewResourcesPrioritizationHandler(attackTracksGetter getter.IAttackTracksGetter) (*ResourcesPrioritizationHandler, error) {
+func NewResourcesPrioritizationHandler(attackTracksGetter getter.IAttackTracksGetter, buildResourcesMap bool) (*ResourcesPrioritizationHandler, error) {
 	handler := &ResourcesPrioritizationHandler{
-		attackTracks: make([]v1alpha1.IAttackTrack, 0),
+		attackTracks:           make([]v1alpha1.IAttackTrack, 0),
+		resourceToAttackTracks: make(map[string]v1alpha1.IAttackTrack),
+		buildResourcesMap:      buildResourcesMap,
 	}
 
 	tracks, err := attackTracksGetter.GetAttackTracks()
@@ -64,7 +69,6 @@ func (handler *ResourcesPrioritizationHandler) PrioritizeResources(sessionObj *c
 		resourcePriorityVector := []prioritization.ControlsVector{}
 		resource, exist := sessionObj.AllResources[resourceId]
 		if !exist {
-			logger.L().Error("resource not found in resources map", helpers.String("resource ID", resourceId))
 			continue
 		}
 
@@ -86,6 +90,12 @@ func (handler *ResourcesPrioritizationHandler) PrioritizeResources(sessionObj *c
 					// Load the failed controls into the attack track
 					allPathsHandler := v1alpha1.NewAttackTrackAllPathsHandler(attackTrack, &controlsLookup)
 
+					// only build the map if the user requested it
+					if handler.buildResourcesMap {
+						// Store the attack track for returning to the caller
+						handler.resourceToAttackTracks[resourceId] = handler.copyAttackTrack(attackTrack, &controlsLookup)
+					}
+
 					// Calculate all the paths for the attack track
 					allAttackPaths := allPathsHandler.CalculateAllPaths()
 
@@ -128,6 +138,8 @@ func (handler *ResourcesPrioritizationHandler) PrioritizeResources(sessionObj *c
 		sessionObj.ResourcesPrioritized[resourceId] = prioritizedResource
 	}
 
+	sessionObj.ResourceAttackTracks = handler.resourceToAttackTracks
+
 	return nil
 }
 
@@ -147,3 +159,18 @@ func (handler *ResourcesPrioritizationHandler) isSupportedKind(obj workloadinter
 	}
 	return false
 }
+
+func (handler *ResourcesPrioritizationHandler) copyAttackTrack(attackTrack v1alpha1.IAttackTrack, lookup v1alpha1.IAttackTrackControlsLookup) v1alpha1.IAttackTrack {
+	copyBytes, _ := json.Marshal(attackTrack)
+	var copyObj v1alpha1.AttackTrack
+	json.Unmarshal(copyBytes, &copyObj)
+
+	iter := copyObj.Iterator()
+	for iter.HasNext() {
+		step := iter.Next()
+		failedControls := lookup.GetAssociatedControls(copyObj.GetName(), step.GetName())
+		step.SetControls(failedControls)
+	}
+
+	return &copyObj
+}

core/pkg/resourcesprioritization/prioritizationhandler_test.go

@@ -97,7 +97,7 @@ func ResourceAssociatedControlMock(controlID string, status apis.ScanningStatus)
 }
 
 func TestNewResourcesPrioritizationHandler(t *testing.T) {
-	handler, err := NewResourcesPrioritizationHandler(&AttackTracksGetterMock{})
+	handler, err := NewResourcesPrioritizationHandler(&AttackTracksGetterMock{}, false)
 	assert.NoError(t, err)
 	assert.Len(t, handler.attackTracks, 2)
 	assert.Equal(t, handler.attackTracks[0].GetName(), "TestAttackTrack")
@@ -182,7 +182,7 @@ func TestResourcesPrioritizationHandler_PrioritizeResources(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			handler, _ := NewResourcesPrioritizationHandler(&AttackTracksGetterMock{})
+			handler, _ := NewResourcesPrioritizationHandler(&AttackTracksGetterMock{}, false)
 			sessionObj := OPASessionObjMock(tt.allPoliciesControls, tt.results, tt.controls, tt.resources)
 			err := handler.PrioritizeResources(sessionObj)
 			assert.NoError(t, err, "expected to have no errors in PrioritizeResources()")

core/pkg/resultshandling/gotree/gotree.go

@@ -0,0 +1,128 @@
+package gotree
+
+import (
+	"strings"
+)
+
+const (
+	newLine      = "\n"
+	emptySpace   = "    "
+	middleItem   = "├── "
+	continueItem = "│   "
+	lastItem     = "└── "
+)
+
+type (
+	tree struct {
+		text  string
+		items []Tree
+	}
+
+	// Tree is tree interface
+	Tree interface {
+		Add(text string) Tree
+		AddTree(tree Tree)
+		Items() []Tree
+		Text() string
+		Print() string
+	}
+
+	printer struct {
+	}
+
+	// Printer is printer interface
+	Printer interface {
+		Print(Tree) string
+	}
+)
+
+// New returns a new GoTree.Tree
+func New(text string) Tree {
+	return &tree{
+		text:  text,
+		items: []Tree{},
+	}
+}
+
+// Add adds a node to the tree
+func (t *tree) Add(text string) Tree {
+	n := New(text)
+	t.items = append(t.items, n)
+	return n
+}
+
+// AddTree adds a tree as an item
+func (t *tree) AddTree(tree Tree) {
+	t.items = append(t.items, tree)
+}
+
+// Text returns the node's value
+func (t *tree) Text() string {
+	return t.text
+}
+
+// Items returns all items in the tree
+func (t *tree) Items() []Tree {
+	return t.items
+}
+
+// Print returns an visual representation of the tree
+func (t *tree) Print() string {
+	return newPrinter().Print(t)
+}
+
+func newPrinter() Printer {
+	return &printer{}
+}
+
+// Print prints a tree to a string
+func (p *printer) Print(t Tree) string {
+	return t.Text() + newLine + p.printItems(t.Items(), []bool{})
+}
+
+func (p *printer) printText(text string, spaces []bool, last bool) string {
+	var result string
+	for _, space := range spaces {
+		if space {
+			result += emptySpace
+		} else {
+			result += continueItem
+		}
+	}
+
+	indicator := middleItem
+	if last {
+		indicator = lastItem
+	}
+
+	var out string
+	lines := strings.Split(text, "\n")
+	for i := range lines {
+		text := lines[i]
+		if i == 0 {
+			out += result + indicator + text + newLine
+			continue
+		}
+		if last {
+			indicator = emptySpace
+		} else {
+			indicator = continueItem
+		}
+		out += result + indicator + text + newLine
+	}
+
+	return out
+}
+
+func (p *printer) printItems(t []Tree, spaces []bool) string {
+	var result string
+	for i, f := range t {
+		last := i == len(t)-1
+		result += p.printText(f.Text(), spaces, last)
+		if len(f.Items()) > 0 {
+			spacesChild := append(spaces, last)
+			result += p.printItems(f.Items(), spacesChild)
+		}
+	}
+	return result
+}

core/pkg/resultshandling/printer/v2/attacktracks.go

@@ -0,0 +1,138 @@
+package printer
+
+import (
+	"fmt"
+	"os"
+	"sort"
+	"strconv"
+	"strings"
+
+	"github.com/fatih/color"
+	"github.com/kubescape/kubescape/v2/core/cautils"
+	"github.com/kubescape/kubescape/v2/core/pkg/resultshandling/gotree"
+	"github.com/kubescape/opa-utils/reporthandling/apis"
+	"github.com/kubescape/opa-utils/reporthandling/attacktrack/v1alpha1"
+	"github.com/kubescape/opa-utils/reporthandling/results/v1/prioritization"
+)
+
+const TOP_RESOURCE_COUNT = 15
+const TOP_VECTOR_COUNT = 10
+
+func (prettyPrinter *PrettyPrinter) printAttackTreeNode(node v1alpha1.IAttackTrackStep, depth int) {
+	prefix := strings.Repeat("\t", depth)
+	text := prefix + node.GetName() + "\n"
+	if len(node.GetControls()) > 0 {
+		color.Red(text)
+	} else {
+		color.Green(text)
+	}
+
+	for i := 0; i < node.Length(); i++ {
+		prettyPrinter.printAttackTreeNode(node.SubStepAt(i), depth+1)
+	}
+}
+
+func (prettyPrinter *PrettyPrinter) createFailedControlList(node v1alpha1.IAttackTrackStep) string {
+	var r string
+	for i, control := range node.GetControls() {
+		if i == 0 {
+			r = control.GetControlId()
+		} else {
+			r = fmt.Sprintf("%s, %s", r, control.GetControlId())
+		}
+	}
+	return r
+}
+
+func (prettyPrinter *PrettyPrinter) buildTreeFromAttackTrackStep(tree gotree.Tree, node v1alpha1.IAttackTrackStep) gotree.Tree {
+	nodeName := node.GetName()
+	if len(node.GetControls()) > 0 {
+		red := color.New(color.Bold, color.FgRed).SprintFunc()
+		nodeName = red(nodeName)
+	}
+
+	controlText := prettyPrinter.createFailedControlList(node)
+	if len(controlText) > 0 {
+		controlStyle := color.New(color.FgWhite, color.Faint).SprintFunc()
+		controlText = controlStyle(fmt.Sprintf(" (%s)", controlText))
+	}
+
+	subTree := gotree.New(nodeName + controlText)
+	for i := 0; i < node.Length(); i++ {
+		subTree.AddTree(prettyPrinter.buildTreeFromAttackTrackStep(tree, node.SubStepAt(i)))
+	}
+
+	if tree == nil {
+		return subTree
+	}
+
+	tree.AddTree(subTree)
+	return tree
+}
+
+func (prettyPrinter *PrettyPrinter) printResourceAttackGraph(attackTrack v1alpha1.IAttackTrack) {
+	tree := prettyPrinter.buildTreeFromAttackTrackStep(nil, attackTrack.GetData())
+	fmt.Fprintln(prettyPrinter.writer, tree.Print())
+}
+
+func getNumericValueFromEnvVar(envVar string, defaultValue int) int {
+	value := os.Getenv(envVar)
+	if value != "" {
+		if value, err := strconv.Atoi(value); err == nil {
+			return value
+		}
+	}
+	return defaultValue
+}
+func (prettyPrinter *PrettyPrinter) printAttackTracks(opaSessionObj *cautils.OPASessionObj) {
+	if prettyPrinter.printAttackTree == false || opaSessionObj.ResourceAttackTracks == nil {
+		return
+	}
+
+	// check if counters are set in env vars and use them, otherwise use default values
+	topResourceCount := getNumericValueFromEnvVar("ATTACK_TREE_TOP_RESOURCES", TOP_RESOURCE_COUNT)
+	topVectorCount := getNumericValueFromEnvVar("ATTACK_TREE_TOP_VECTORS", TOP_VECTOR_COUNT)
+
+	prioritizedResources := opaSessionObj.ResourcesPrioritized
+	resourceToAttackTrack := opaSessionObj.ResourceAttackTracks
+
+	resources := make([]prioritization.PrioritizedResource, 0, len(prioritizedResources))
+	for _, value := range prioritizedResources {
+		resources = append(resources, value)
+	}
+
+	sort.Slice(resources, func(i, j int) bool {
+		return resources[i].Score > resources[j].Score
+	})
+
+	for i := 0; i < topResourceCount && i < len(resources); i++ {
+		fmt.Fprintf(prettyPrinter.writer, "\n"+getSeparator("^")+"\n")
+		resource := resources[i]
+		resourceObj := opaSessionObj.AllResources[resource.ResourceID]
+
+		fmt.Fprintf(prettyPrinter.writer, "Name: %s\n", resourceObj.GetName())
+		fmt.Fprintf(prettyPrinter.writer, "Kind: %s\n", resourceObj.GetKind())
+		fmt.Fprintf(prettyPrinter.writer, "Namespace: %s\n\n", resourceObj.GetNamespace())
+
+		fmt.Fprintf(prettyPrinter.writer, "Score: %.2f\n", resource.Score)
+		fmt.Fprintf(prettyPrinter.writer, "Severity: %s\n", apis.SeverityNumberToString(resource.Severity))
+		fmt.Fprintf(prettyPrinter.writer, "Total vectors: %v\n\n", len(resources[i].PriorityVector))
+
+		prettyPrinter.printResourceAttackGraph(resourceToAttackTrack[resource.ResourceID])
+
+		sort.Slice(resource.PriorityVector, func(x, y int) bool {
+			return resource.PriorityVector[x].Score > resource.PriorityVector[y].Score
+		})
+
+		for j := 0; j < topVectorCount && j < len(resources[i].PriorityVector); j++ {
+			priorityVector := resource.PriorityVector[j]
+
+			vectorStrings := []string{}
+			for _, controlId := range priorityVector.ListControls() {
+				vectorStrings = append(vectorStrings, fmt.Sprintf("%s (%s)", controlId.Category, controlId.ControlID))
+			}
+
+			fmt.Fprintf(prettyPrinter.writer, "%v) [%.2f] [Severity: %v] [Attack Track: %v]: %v \n", j+1, priorityVector.Score, apis.SeverityNumberToString(priorityVector.Severity), priorityVector.AttackTrackName, strings.Join(vectorStrings, " -> "))
+		}
+	}
+}

core/pkg/resultshandling/printer/v2/prettyprinter.go

@@ -24,17 +24,19 @@ const (
 )
 
 type PrettyPrinter struct {
-	formatVersion string
-	viewType      cautils.ViewTypes
-	writer        *os.File
-	verboseMode   bool
+	writer          *os.File
+	formatVersion   string
+	viewType        cautils.ViewTypes
+	verboseMode     bool
+	printAttackTree bool
 }
 
-func NewPrettyPrinter(verboseMode bool, formatVersion string, viewType cautils.ViewTypes) *PrettyPrinter {
+func NewPrettyPrinter(verboseMode bool, formatVersion string, attackTree bool, viewType cautils.ViewTypes) *PrettyPrinter {
 	return &PrettyPrinter{
-		verboseMode:   verboseMode,
-		formatVersion: formatVersion,
-		viewType:      viewType,
+		verboseMode:     verboseMode,
+		formatVersion:   formatVersion,
+		viewType:        viewType,
+		printAttackTree: attackTree,
 	}
 }
 
@@ -60,6 +62,7 @@ func (pp *PrettyPrinter) ActionPrint(opaSessionObj *cautils.OPASessionObj) {
 		printer.LogOutputFile(pp.writer.Name())
 	}
 
+	pp.printAttackTracks(opaSessionObj)
 }
 
 func (pp *PrettyPrinter) SetWriter(outputFile string) {

core/pkg/resultshandling/results.go

@@ -92,7 +92,7 @@ func (rh *ResultsHandler) HandleResults() error {
 }
 
 // NewPrinter returns a new printer for a given format and configuration options
-func NewPrinter(printFormat, formatVersion string, verboseMode bool, viewType cautils.ViewTypes) printer.IPrinter {
+func NewPrinter(printFormat, formatVersion string, verboseMode bool, attackTree bool, viewType cautils.ViewTypes) printer.IPrinter {
 
 	switch printFormat {
 	case printer.JsonFormat:
@@ -117,6 +117,6 @@ func NewPrinter(printFormat, formatVersion string, verboseMode bool, viewType ca
 		if printFormat != printer.PrettyFormat {
 			logger.L().Error(fmt.Sprintf("Invalid format \"%s\", default format \"pretty-printer\" is applied", printFormat))
 		}
-		return printerv2.NewPrettyPrinter(verboseMode, formatVersion, viewType)
+		return printerv2.NewPrettyPrinter(verboseMode, formatVersion, attackTree, viewType)
 	}
 }