trivy/pkg/vulnerability/vulnerability_test.go

package vulnerability

import (
	"os"
	"testing"

	"golang.org/x/xerrors"

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

	"github.com/aquasecurity/trivy/pkg/log"
	"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 args struct {
		vulns      []types.DetectedVulnerability
		reportType string
	}
	tests := []struct {
		name                    string
		getSeverity             []db.GetSeverityExpectation
		getVulnerability        []db.GetVulnerabilityExpectation
		args                    args
		expectedVulnerabilities []types.DetectedVulnerability
	}{
		{
			name: "happy path, with only OS vulnerability but no vendor severity, no NVD",
			getVulnerability: []db.GetVulnerabilityExpectation{
				{
					Args: db.GetVulnerabilityArgs{
						VulnerabilityID: "CVE-2019-0001",
					},
					Returns: db.GetVulnerabilityReturns{
						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"},
				},
				reportType: vulnerability.Ubuntu,
			},
			expectedVulnerabilities: []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 only OS vulnerability but no vendor severity, yes NVD",
			getVulnerability: []db.GetVulnerabilityExpectation{
				{
					Args: db.GetVulnerabilityArgs{
						VulnerabilityID: "CVE-2019-0001",
					},
					Returns: db.GetVulnerabilityReturns{
						Vulnerability: dbTypes.Vulnerability{
							Title:       "dos",
							Description: "dos vulnerability",
							VendorSeverity: dbTypes.VendorSeverity{
								vulnerability.Nvd: dbTypes.SeverityLow,
							},
							References: []string{"http://example.com"},
						},
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0001"},
				},
				reportType: vulnerability.Ubuntu,
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "dos",
						Description: "dos vulnerability",
						Severity:    dbTypes.SeverityLow.String(),
						References:  []string{"http://example.com"},
					},
					SeveritySource: vulnerability.Nvd,
				},
			},
		},
		{
			name: "happy path, with only OS vulnerability but no severity, no vendor severity, no NVD",
			getVulnerability: []db.GetVulnerabilityExpectation{
				{
					Args: db.GetVulnerabilityArgs{
						VulnerabilityID: "CVE-2019-0001",
					},
					Returns: db.GetVulnerabilityReturns{
						Vulnerability: dbTypes.Vulnerability{
							Title:       "dos",
							Description: "dos vulnerability",
							References:  []string{"http://example.com"},
						},
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0001"},
				},
				reportType: vulnerability.Ubuntu,
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "dos",
						Description: "dos vulnerability",
						References:  []string{"http://example.com"},
					},
				},
			},
		},
		{
			name: "happy path, with only OS vulnerability, yes vendor severity",
			getVulnerability: []db.GetVulnerabilityExpectation{
				{
					Args: db.GetVulnerabilityArgs{
						VulnerabilityID: "CVE-2019-0001",
					},
					Returns: db.GetVulnerabilityReturns{
						Vulnerability: dbTypes.Vulnerability{
							Title:       "dos",
							Description: "dos vulnerability",
							Severity:    dbTypes.SeverityMedium.String(),
							VendorSeverity: dbTypes.VendorSeverity{
								vulnerability.RedHat: dbTypes.SeverityLow, // CentOS uses RedHat
							},
							References: []string{"http://example.com"},
						},
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0001"},
				},
				reportType: vulnerability.CentOS,
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "dos",
						Description: "dos vulnerability",
						Severity:    dbTypes.SeverityLow.String(),
						References:  []string{"http://example.com"},
					},
					SeveritySource: vulnerability.RedHat,
				},
			},
		},
		{
			name: "happy path light db, with only OS vulnerability, yes vendor severity",
			getVulnerability: []db.GetVulnerabilityExpectation{
				{
					Args: db.GetVulnerabilityArgs{
						VulnerabilityID: "CVE-2019-0001",
					},
					Returns: db.GetVulnerabilityReturns{
						Vulnerability: dbTypes.Vulnerability{
							Severity: dbTypes.SeverityMedium.String(),
							VendorSeverity: dbTypes.VendorSeverity{
								vulnerability.Ubuntu: dbTypes.SeverityLow,
							},
						},
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0001"},
				},
				reportType: vulnerability.Ubuntu,
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityLow.String(),
					},
					SeveritySource: vulnerability.Ubuntu,
				},
			},
		},
		{
			name: "happy path, with only library vulnerability",
			getVulnerability: []db.GetVulnerabilityExpectation{
				{
					Args: db.GetVulnerabilityArgs{
						VulnerabilityID: "CVE-2020-0001",
					},
					Returns: db.GetVulnerabilityReturns{
						Vulnerability: dbTypes.Vulnerability{
							Title:       "COVID-19",
							Description: "a nasty virus vulnerability for humans",
							Severity:    dbTypes.SeverityMedium.String(),
							VendorSeverity: dbTypes.VendorSeverity{
								vulnerability.PythonSafetyDB: dbTypes.SeverityCritical,
							},
							References: []string{"https://www.who.int/emergencies/diseases/novel-coronavirus-2019"},
						},
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2020-0001"},
				},
				reportType: "poetry",
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2020-0001",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "COVID-19",
						Description: "a nasty virus vulnerability for humans",
						Severity:    dbTypes.SeverityCritical.String(),
						References:  []string{"https://www.who.int/emergencies/diseases/novel-coronavirus-2019"},
					},
					SeveritySource: vulnerability.PythonSafetyDB,
				},
			},
		},
		{
			name: "GetVulnerability returns an error",
			getVulnerability: []db.GetVulnerabilityExpectation{
				{
					Args: db.GetVulnerabilityArgs{
						VulnerabilityID: "CVE-2019-0004",
					},
					Returns: db.GetVulnerabilityReturns{
						Err: xerrors.New("failed"),
					},
				},
			},
			args: args{
				vulns: []types.DetectedVulnerability{
					{VulnerabilityID: "CVE-2019-0004"},
				},
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{VulnerabilityID: "CVE-2019-0004"},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockDBConfig := new(db.MockOperation)
			mockDBConfig.ApplyGetSeverityExpectations(tt.getSeverity)
			mockDBConfig.ApplyGetVulnerabilityExpectations(tt.getVulnerability)

			c := Client{
				dbc: mockDBConfig,
			}

			c.FillInfo(tt.args.vulns, tt.args.reportType)
			assert.Equal(t, tt.expectedVulnerabilities, 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)
		})
	}
}