trivy/pkg/vulnerability/vulnerability.go

package vulnerability

import (
	"bufio"
	"context"
	"io/ioutil"
	"os"
	"sort"
	"strings"

	"github.com/google/wire"
	"github.com/open-policy-agent/opa/rego"
	"golang.org/x/xerrors"

	"github.com/aquasecurity/trivy-db/pkg/db"
	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
	"github.com/aquasecurity/trivy-db/pkg/vulnsrc/vulnerability"
	"github.com/aquasecurity/trivy/pkg/log"
	"github.com/aquasecurity/trivy/pkg/types"
	"github.com/aquasecurity/trivy/pkg/utils"
)

const (
	// DefaultIgnoreFile is the file name to be ignored
	DefaultIgnoreFile = ".trivyignore"
)

// SuperSet binds the dependencies
var SuperSet = wire.NewSet(
	wire.Struct(new(db.Config)),
	NewClient,
	wire.Bind(new(Operation), new(Client)),
)

// Operation defines the vulnerability operations
type Operation interface {
	FillInfo(vulns []types.DetectedVulnerability, reportType string)
	Filter(ctx context.Context, vulns []types.DetectedVulnerability, severities []dbTypes.Severity,
		ignoreUnfixed bool, ignoreFile string, policy string) ([]types.DetectedVulnerability, error)
}

// Client implements db operations
type Client struct {
	dbc db.Operation
}

// NewClient is the factory method for Client
func NewClient(dbc db.Config) Client {
	return Client{dbc: dbc}
}

// FillInfo fills extra info rin vulnerability objects
func (c Client) FillInfo(vulns []types.DetectedVulnerability, reportType string) {
	var err error

	for i := range vulns {
		vulns[i].Vulnerability, err = c.dbc.GetVulnerability(vulns[i].VulnerabilityID)
		if err != nil {
			log.Logger.Warnf("Error while getting vulnerability details: %s\n", err)
			continue
		}

		var source string
		switch reportType {
		case vulnerability.Ubuntu, vulnerability.Alpine, vulnerability.RedHat, vulnerability.RedHatOVAL, vulnerability.Debian, vulnerability.DebianOVAL, vulnerability.Fedora, vulnerability.Amazon, vulnerability.OracleOVAL, vulnerability.SuseCVRF, vulnerability.OpenSuseCVRF, vulnerability.Photon:
			source = reportType
		case vulnerability.CentOS: // CentOS doesn't have its own so we use RedHat
			source = vulnerability.RedHat
		case "npm", "yarn":
			source = vulnerability.NodejsSecurityWg
		case "pipenv", "poetry":
			source = vulnerability.PythonSafetyDB
		case "bundler":
			source = vulnerability.RubySec
		case "cargo":
			source = vulnerability.RustSec
		case "composer":
			source = vulnerability.PhpSecurityAdvisories
		}
		c.getVendorSeverity(&vulns[i], source)
		vulns[i].Vulnerability.VendorSeverity = nil // Remove VendorSeverity from Results
	}
}

func (c Client) getVendorSeverity(vuln *types.DetectedVulnerability, reportType string) {
	if vs, ok := vuln.Vulnerability.VendorSeverity[reportType]; ok {
		vuln.Vulnerability.Severity = vs.String()
		vuln.SeveritySource = reportType
		return
	}

	// Try NVD as a fallback if it exists
	if vs, ok := vuln.Vulnerability.VendorSeverity[vulnerability.Nvd]; ok {
		vuln.Vulnerability.Severity = vs.String()
		vuln.SeveritySource = vulnerability.Nvd
		return
	}
}

// Filter filter out the vulnerabilities
func (c Client) Filter(ctx context.Context, vulns []types.DetectedVulnerability, severities []dbTypes.Severity,
	ignoreUnfixed bool, ignoreFile string, policyFile string) ([]types.DetectedVulnerability, error) {
	ignoredIDs := getIgnoredIDs(ignoreFile)
	var vulnerabilities []types.DetectedVulnerability
	for _, vuln := range vulns {
		// Filter vulnerabilities by severity
		for _, s := range severities {
			if s.String() == vuln.Severity {
				// Ignore unfixed vulnerabilities
				if ignoreUnfixed && vuln.FixedVersion == "" {
					continue
				} else if utils.StringInSlice(vuln.VulnerabilityID, ignoredIDs) {
					continue
				}
				vulnerabilities = append(vulnerabilities, vuln)
				break
			}
		}
	}

	if policyFile != "" {
		var err error
		vulnerabilities, err = applyPolicy(ctx, vulnerabilities, policyFile)
		if err != nil {
			return nil, xerrors.Errorf("failed to apply the policy: %w", err)
		}
	}
	sort.Sort(types.BySeverity(vulnerabilities))
	return vulnerabilities, nil
}

func applyPolicy(ctx context.Context, vulns []types.DetectedVulnerability, policyFile string) ([]types.DetectedVulnerability, error) {
	policy, err := ioutil.ReadFile(policyFile)
	if err != nil {
		return nil, xerrors.Errorf("unable to read the policy file: %w", err)
	}

	query, err := rego.New(
		rego.Query("data.trivy.ignore"),
		rego.Module("lib.rego", module),
		rego.Module("trivy.rego", string(policy)),
	).PrepareForEval(ctx)
	if err != nil {
		return nil, xerrors.Errorf("unable to prepare for eval: %w", err)
	}

	var filtered []types.DetectedVulnerability
	for _, vuln := range vulns {
		results, err := query.Eval(ctx, rego.EvalInput(vuln))
		if err != nil {
			return nil, xerrors.Errorf("unable to evaluate the policy: %w", err)
		} else if len(results) == 0 {
			// Handle undefined result.
			filtered = append(filtered, vuln)
			continue
		}
		ignore, ok := results[0].Expressions[0].Value.(bool)
		if !ok {
			// Handle unexpected result type.
			return nil, xerrors.New("the policy must return boolean")
		}
		if ignore {
			continue
		}
		filtered = append(filtered, vuln)
	}
	return filtered, nil
}

func getIgnoredIDs(ignoreFile string) []string {
	f, err := os.Open(ignoreFile)
	if err != nil {
		// trivy must work even if no .trivyignore exist
		return nil
	}

	var ignoredIDs []string
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "#") || line == "" {
			continue
		}
		ignoredIDs = append(ignoredIDs, line)
	}
	return ignoredIDs
}