vulnerability.go

package table

import (
	"bytes"
	"fmt"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"github.com/samber/lo"
	"github.com/xlab/treeprint"
	"golang.org/x/exp/maps"
	"golang.org/x/exp/slices"

	"github.com/aquasecurity/table"
	"github.com/aquasecurity/tml"
	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
	ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
	"github.com/aquasecurity/trivy/pkg/log"
	"github.com/aquasecurity/trivy/pkg/types"
)

type vulnerabilityRenderer struct {
	w           *bytes.Buffer
	tableWriter *table.Table
	result      types.Result
	isTerminal  bool
	tree        bool
	severities  []dbTypes.Severity
	once        *sync.Once
}

func NewVulnerabilityRenderer(result types.Result, isTerminal, tree bool, severities []dbTypes.Severity) *vulnerabilityRenderer {
	buf := bytes.NewBuffer([]byte{})
	if !isTerminal {
		tml.DisableFormatting()
	}
	return &vulnerabilityRenderer{
		w:           buf,
		tableWriter: newTableWriter(buf, isTerminal),
		result:      result,
		isTerminal:  isTerminal,
		tree:        tree,
		severities:  severities,
		once:        new(sync.Once),
	}
}

func (r *vulnerabilityRenderer) Render() string {
	r.setHeaders()
	r.setVulnerabilityRows(r.result.Vulnerabilities)

	severityCount := r.countSeverities(r.result.Vulnerabilities)
	total, summaries := summarize(r.severities, severityCount)

	target := r.result.Target
	if r.result.Class == types.ClassLangPkg {
		target += fmt.Sprintf(" (%s)", r.result.Type)
	}
	RenderTarget(r.w, target, r.isTerminal)
	r.printf("Total: %d (%s)\n\n", total, strings.Join(summaries, ", "))

	r.tableWriter.Render()
	if r.tree {
		r.renderDependencyTree()
	}

	return r.w.String()
}

func (r *vulnerabilityRenderer) setHeaders() {
	if len(r.result.Vulnerabilities) == 0 {
		return
	}
	header := []string{
		"Library",
		"Vulnerability",
		"Severity",
		"Status",
		"Installed Version",
		"Fixed Version",
		"Title",
	}
	r.tableWriter.SetHeaders(header...)
}

func (r *vulnerabilityRenderer) setVulnerabilityRows(vulns []types.DetectedVulnerability) {
	for _, v := range vulns {
		lib := v.PkgName
		if v.PkgPath != "" {
			// get path to root jar
			// for other languages return unchanged path
			pkgPath := rootJarFromPath(v.PkgPath)
			fileName := filepath.Base(pkgPath)
			lib = fmt.Sprintf("%s (%s)", v.PkgName, fileName)
			r.once.Do(func() {
				log.Logger.Infof("Table result includes only package filenames. Use '--format json' option to get the full path to the package file.")
			})
		}

		title := v.Title
		if title == "" {
			title = v.Description
		}
		splitTitle := strings.Split(title, " ")
		if len(splitTitle) >= 12 {
			title = strings.Join(splitTitle[:12], " ") + "..."
		}

		if len(v.PrimaryURL) > 0 {
			if r.isTerminal {
				title = tml.Sprintf("%s\n<blue>%s</blue>", title, v.PrimaryURL)
			} else {
				title = fmt.Sprintf("%s\n%s", title, v.PrimaryURL)
			}
		}

		var row []string
		if r.isTerminal {
			row = []string{
				lib,
				v.VulnerabilityID,
				ColorizeSeverity(v.Severity, v.Severity),
				v.Status.String(),
				v.InstalledVersion,
				v.FixedVersion,
				strings.TrimSpace(title),
			}
		} else {
			row = []string{
				lib,
				v.VulnerabilityID,
				v.Severity,
				v.Status.String(),
				v.InstalledVersion,
				v.FixedVersion,
				strings.TrimSpace(title),
			}
		}

		r.tableWriter.AddRow(row...)
	}
}

func (r *vulnerabilityRenderer) countSeverities(vulns []types.DetectedVulnerability) map[string]int {
	severityCount := make(map[string]int)
	for _, v := range vulns {
		severityCount[v.Severity]++
	}
	return severityCount
}

func (r *vulnerabilityRenderer) renderDependencyTree() {
	// Get parents of each dependency
	parents := ftypes.Packages(r.result.Packages).ParentDeps()
	if len(parents) == 0 {
		return
	}
	ancestors := traverseAncestors(r.result.Packages, parents)

	root := treeprint.NewWithRoot(fmt.Sprintf(`
Dependency Origin Tree (Reversed)
=================================
%s`, r.result.Target))

	// This count is next to the package ID.
	// e.g. node-fetch@1.7.3 (MEDIUM: 2, HIGH: 1, CRITICAL: 3)
	pkgSeverityCount := make(map[string]map[string]int)
	for _, vuln := range r.result.Vulnerabilities {
		cnts, ok := pkgSeverityCount[vuln.PkgID]
		if !ok {
			cnts = make(map[string]int)
		}

		cnts[vuln.Severity]++
		pkgSeverityCount[vuln.PkgID] = cnts
	}

	// Extract vulnerable packages
	vulnPkgs := lo.Filter(r.result.Packages, func(pkg ftypes.Package, _ int) bool {
		return lo.ContainsBy(r.result.Vulnerabilities, func(vuln types.DetectedVulnerability) bool {
			return pkg.ID == vuln.PkgID
		})
	})

	// Render tree
	for _, vulnPkg := range vulnPkgs {
		_, summaries := summarize(r.severities, pkgSeverityCount[vulnPkg.ID])
		topLvlID := tml.Sprintf("<red>%s, (%s)</red>", vulnPkg.ID, strings.Join(summaries, ", "))

		branch := root.AddBranch(topLvlID)
		addParents(branch, vulnPkg, parents, ancestors, map[string]struct{}{vulnPkg.ID: {}}, 1)

	}
	r.printf(root.String())
}

func (r *vulnerabilityRenderer) printf(format string, args ...interface{}) {
	// nolint
	_ = tml.Fprintf(r.w, format, args...)
}

func addParents(topItem treeprint.Tree, pkg ftypes.Package, parentMap map[string]ftypes.Packages, ancestors map[string][]string,
	seen map[string]struct{}, depth int) {
	if !pkg.Indirect {
		return
	}

	roots := make(map[string]struct{})
	for _, parent := range parentMap[pkg.ID] {
		if _, ok := seen[parent.ID]; ok {
			continue
		}
		seen[parent.ID] = struct{}{} // to avoid infinite loops

		if depth == 1 && !parent.Indirect {
			topItem.AddBranch(parent.ID)
		} else {
			// We omit intermediate dependencies and show only direct dependencies
			// as this could make the dependency tree huge.
			for _, ancestor := range ancestors[parent.ID] {
				roots[ancestor] = struct{}{}
			}
		}
	}

	// Omitted
	rootIDs := lo.Filter(maps.Keys(roots), func(pkgID string, _ int) bool {
		_, ok := seen[pkgID]
		return !ok
	})
	sort.Strings(rootIDs)
	if len(rootIDs) > 0 {
		branch := topItem.AddBranch("...(omitted)...")
		for _, rootID := range rootIDs {
			branch.AddBranch(rootID)
		}
	}
}

func traverseAncestors(pkgs []ftypes.Package, parentMap map[string]ftypes.Packages) map[string][]string {
	ancestors := make(map[string][]string)
	for _, pkg := range pkgs {
		ancestors[pkg.ID] = findAncestor(pkg.ID, parentMap, make(map[string]struct{}))
	}
	return ancestors
}

func findAncestor(pkgID string, parentMap map[string]ftypes.Packages, seen map[string]struct{}) []string {
	ancestors := make(map[string]struct{})
	seen[pkgID] = struct{}{}
	for _, parent := range parentMap[pkgID] {
		if _, ok := seen[parent.ID]; ok {
			continue
		}
		switch {
		case !parent.Indirect:
			ancestors[parent.ID] = struct{}{}
		case len(parentMap[parent.ID]) == 0:
			// Direct dependencies cannot be identified in some package managers like "package-lock.json" v1,
			// then the "Indirect" field can be always true. We try to guess direct dependencies in this case.
			// A dependency with no parents must be a direct dependency.
			//
			// e.g.
			//   -> styled-components
			//     -> fbjs
			//       -> isomorphic-fetch
			//         -> node-fetch
			//
			// Even if `styled-components` is not marked as a direct dependency, it must be a direct dependency
			// as it has no parents. Note that it doesn't mean `fbjs` is an indirect dependency.
			ancestors[parent.ID] = struct{}{}
		default:
			for _, a := range findAncestor(parent.ID, parentMap, seen) {
				ancestors[a] = struct{}{}
			}
		}
	}
	return maps.Keys(ancestors)
}

var jarExtensions = []string{
	".jar",
	".war",
	".par",
	".ear",
}

func rootJarFromPath(path string) string {
	// File paths are always forward-slashed in Trivy
	paths := strings.Split(path, "/")
	for i, p := range paths {
		if slices.Contains(jarExtensions, filepath.Ext(p)) {
			return strings.Join(paths[:i+1], "/")
		}
	}
	return path
}