trivy/pkg/vulnerability/vulnerability_test.go

package vulnerability

import (
	"os"
	"testing"

	"github.com/aquasecurity/trivy/pkg/log"
	"golang.org/x/xerrors"

	"github.com/aquasecurity/trivy/pkg/types"

	"github.com/aquasecurity/trivy-db/pkg/db"

	"github.com/stretchr/testify/assert"

	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
)

func TestMain(m *testing.M) {
	if err := log.InitLogger(false, true); err != nil {
		log.Fatal(err)
	}

	code := m.Run()
	os.Exit(code)
}

func TestClient_FillInfo(t *testing.T) {
	type getSeverityOutput struct {
		severity dbTypes.Severity
		err      error
	}
	type getSeverity struct {
		input  string
		output getSeverityOutput
	}
	type getVulnerabilityOutput struct {
		vulnerability dbTypes.Vulnerability
		err           error
	}
	type getVulnerability struct {
		input  string
		output getVulnerabilityOutput
	}

	type args struct {
		vulns []types.DetectedVulnerability
		light bool
	}
	tests := []struct {
		name             string
		getSeverity      []getSeverity
		getVulnerability []getVulnerability
		args             args
		expected         []types.DetectedVulnerability
	}{
		{
			name: "happy path",
			getVulnerability: []getVulnerability{
				{
					input: "CVE-2019-0001",
					output: getVulnerabilityOutput{
						vulnerability: dbTypes.Vulnerability{
							Title:       "dos",
							Description: "dos vulnerability",
							Severity:    dbTypes.SeverityMedium.String(),
							References:  []string{"http://example.com"},
						},
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0001"},
				},
				light: false,
			},
			expected: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "dos",
						Description: "dos vulnerability",
						Severity:    dbTypes.SeverityMedium.String(),
						References:  []string{"http://example.com"},
					},
				},
			},
		},
		{
			name: "happy path with light option",
			getSeverity: []getSeverity{
				{
					input: "CVE-2019-0001",
					output: getSeverityOutput{
						severity: dbTypes.SeverityCritical,
					},
				},
				{
					input: "CVE-2019-0002",
					output: getSeverityOutput{
						severity: dbTypes.SeverityHigh,
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0001"},
					{VulnerabilityID: "CVE-2019-0002"},
				},
				light: true,
			},
			expected: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityCritical.String(),
					},
				},
				{
					VulnerabilityID: "CVE-2019-0002",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityHigh.String(),
					},
				},
			},
		},
		{
			name: "GetVulnerability returns an error",
			getVulnerability: []getVulnerability{
				{
					input: "CVE-2019-0004",
					output: getVulnerabilityOutput{
						err: xerrors.New("failed"),
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0004"},
				},
				light: false,
			},
			expected: []types.DetectedVulnerability{
				{VulnerabilityID: "CVE-2019-0004"},
			},
		},
		{
			name: "GetSeverity returns an error",
			getSeverity: []getSeverity{
				{
					input: "CVE-2019-0003",
					output: getSeverityOutput{
						err: xerrors.New("failed"),
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0003"},
				},
				light: true,
			},
			expected: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0003",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityUnknown.String(),
					},
				},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockDBConfig := new(db.MockDBConfig)
			c := Client{
				dbc: mockDBConfig,
			}
			for _, gs := range tt.getSeverity {
				mockDBConfig.On("GetSeverity", gs.input).Return(
					gs.output.severity, gs.output.err)
			}
			for _, gv := range tt.getVulnerability {
				mockDBConfig.On("GetVulnerability", gv.input).Return(
					gv.output.vulnerability, gv.output.err)
			}

			c.FillInfo(tt.args.vulns, tt.args.light)
			assert.Equal(t, tt.expected, tt.args.vulns, tt.name)
			mockDBConfig.AssertExpectations(t)
		})
	}
}

func TestClient_Filter(t *testing.T) {
	type args struct {
		vulns         []types.DetectedVulnerability
		severities    []dbTypes.Severity
		ignoreUnfixed bool
		ignoreFile    string
	}
	tests := []struct {
		name string
		args args
		want []types.DetectedVulnerability
	}{
		{
			name: "happy path",
			args: args{
				vulns: []types.DetectedVulnerability{
					{
						VulnerabilityID:  "CVE-2019-0001",
						PkgName:          "foo",
						InstalledVersion: "1.2.3",
						FixedVersion:     "1.2.4",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityLow.String(),
						},
					},
					{
						VulnerabilityID:  "CVE-2019-0002",
						PkgName:          "bar",
						InstalledVersion: "1.2.3",
						FixedVersion:     "1.2.4",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityCritical.String(),
						},
					},
					{
						VulnerabilityID:  "CVE-2018-0001",
						PkgName:          "baz",
						InstalledVersion: "1.2.3",
						FixedVersion:     "",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityHigh.String(),
						},
					},
					{
						VulnerabilityID:  "CVE-2018-0001",
						PkgName:          "bar",
						InstalledVersion: "1.2.3",
						FixedVersion:     "",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityCritical.String(),
						},
					},
					{
						VulnerabilityID:  "CVE-2018-0002",
						PkgName:          "bar",
						InstalledVersion: "1.2.3",
						FixedVersion:     "",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityHigh.String(),
						},
					},
				},
				severities:    []dbTypes.Severity{dbTypes.SeverityCritical, dbTypes.SeverityHigh},
				ignoreUnfixed: false,
			},
			want: []types.DetectedVulnerability{
				{
					VulnerabilityID:  "CVE-2018-0001",
					PkgName:          "bar",
					InstalledVersion: "1.2.3",
					FixedVersion:     "",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityCritical.String(),
					},
				},
				{
					VulnerabilityID:  "CVE-2019-0002",
					PkgName:          "bar",
					InstalledVersion: "1.2.3",
					FixedVersion:     "1.2.4",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityCritical.String(),
					},
				},
				{
					VulnerabilityID:  "CVE-2018-0002",
					PkgName:          "bar",
					InstalledVersion: "1.2.3",
					FixedVersion:     "",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityHigh.String(),
					},
				},
				{
					VulnerabilityID:  "CVE-2018-0001",
					PkgName:          "baz",
					InstalledVersion: "1.2.3",
					FixedVersion:     "",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityHigh.String(),
					},
				},
			},
		},
		{
			name: "happy path with ignore-unfixed",
			args: args{
				vulns: []types.DetectedVulnerability{
					{
						VulnerabilityID:  "CVE-2019-0001",
						PkgName:          "foo",
						InstalledVersion: "1.2.3",
						FixedVersion:     "1.2.4",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityLow.String(),
						},
					},
					{
						VulnerabilityID:  "CVE-2018-0002",
						PkgName:          "bar",
						InstalledVersion: "1.2.3",
						FixedVersion:     "",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityHigh.String(),
						},
					},
				},
				severities:    []dbTypes.Severity{dbTypes.SeverityHigh},
				ignoreUnfixed: true,
			},
		},
		{
			name: "happy path with ignore-file",
			args: args{
				vulns: []types.DetectedVulnerability{
					{
						// this vulnerability is ignored
						VulnerabilityID:  "CVE-2019-0001",
						PkgName:          "foo",
						InstalledVersion: "1.2.3",
						FixedVersion:     "1.2.4",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityLow.String(),
						},
					},
					{
						// this vulnerability is ignored
						VulnerabilityID:  "CVE-2019-0002",
						PkgName:          "foo",
						InstalledVersion: "1.2.3",
						FixedVersion:     "1.2.4",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityLow.String(),
						},
					},
					{
						VulnerabilityID:  "CVE-2019-0003",
						PkgName:          "foo",
						InstalledVersion: "1.2.3",
						FixedVersion:     "1.2.4",
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityLow.String(),
						},
					},
				},
				severities:    []dbTypes.Severity{dbTypes.SeverityLow},
				ignoreUnfixed: false,
				ignoreFile:    "testdata/.trivyignore",
			},
			want: []types.DetectedVulnerability{
				{
					VulnerabilityID:  "CVE-2019-0003",
					PkgName:          "foo",
					InstalledVersion: "1.2.3",
					FixedVersion:     "1.2.4",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityLow.String(),
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := Client{}
			got := c.Filter(tt.args.vulns, tt.args.severities, tt.args.ignoreUnfixed, tt.args.ignoreFile)
			assert.Equal(t, tt.want, got, tt.name)
		})
	}
}