package vulnerability_test

import (
	"testing"

	"github.com/aquasecurity/trivy-db/pkg/db"
	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
	"github.com/aquasecurity/trivy-db/pkg/utils"
	"github.com/aquasecurity/trivy-db/pkg/vulnsrc/vulnerability"
	"github.com/aquasecurity/trivy/pkg/dbtest"
	"github.com/aquasecurity/trivy/pkg/types"
	vuln "github.com/aquasecurity/trivy/pkg/vulnerability"
	"github.com/stretchr/testify/assert"
)

func TestClient_FillInfo(t *testing.T) {
	tests := []struct {
		name                    string
		fixtures                []string
		vulns                   []types.DetectedVulnerability
		expectedVulnerabilities []types.DetectedVulnerability
	}{
		{
			name:     "happy path, with only OS vulnerability but no vendor severity, no NVD",
			fixtures: []string{"testdata/fixtures/vulnerability.yaml"},
			vulns: []types.DetectedVulnerability{
				{VulnerabilityID: "CVE-2019-0001"},
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					Vulnerability: dbTypes.Vulnerability{
						Title:            "dos",
						Description:      "dos vulnerability",
						Severity:         dbTypes.SeverityMedium.String(),
						References:       []string{"http://example.com"},
						LastModifiedDate: utils.MustTimeParse("2020-01-01T01:01:00Z"),
						PublishedDate:    utils.MustTimeParse("2001-01-01T01:01:00Z"),
					},
					PrimaryURL: "https://avd.aquasec.com/nvd/cve-2019-0001",
				},
			},
		},
		{
			name:     "happy path, with only OS vulnerability but no vendor severity, yes NVD",
			fixtures: []string{"testdata/fixtures/vulnerability.yaml"},
			vulns: []types.DetectedVulnerability{
				{
					VulnerabilityID: "RUSTSEC-2018-0017",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.OSV,
						Name: "RustSec Advisory Database",
						URL:  "https://github.com/RustSec/advisory-db",
					},
				},
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "RUSTSEC-2018-0017",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "dos",
						Description: "dos vulnerability",
						Severity:    dbTypes.SeverityLow.String(),
						VendorSeverity: dbTypes.VendorSeverity{
							vulnerability.NVD: dbTypes.SeverityLow,
						},
						LastModifiedDate: utils.MustTimeParse("2020-01-01T01:01:00Z"),
						PublishedDate:    utils.MustTimeParse("2001-01-01T01:01:00Z"),
					},
					SeveritySource: vulnerability.NVD,
					PrimaryURL:     "https://osv.dev/vulnerability/RUSTSEC-2018-0017",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.OSV,
						Name: "RustSec Advisory Database",
						URL:  "https://github.com/RustSec/advisory-db",
					},
				},
			},
		},
		{
			name:     "happy path, with only OS vulnerability but no severity, no vendor severity, no NVD",
			fixtures: []string{"testdata/fixtures/vulnerability.yaml"},
			vulns: []types.DetectedVulnerability{
				{
					VulnerabilityID: "GHSA-28fw-88hq-6jmm",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.GHSA,
						Name: "GitHub Security Advisory Pip",
						URL:  "https://github.com/advisories?query=type%3Areviewed+ecosystem%3Apip",
					},
				},
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "GHSA-28fw-88hq-6jmm",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "dos",
						Description: "dos vulnerability",
						Severity:    dbTypes.SeverityUnknown.String(),
						References:  []string{"http://example.com"},
					},
					PrimaryURL: "https://github.com/advisories/GHSA-28fw-88hq-6jmm",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.GHSA,
						Name: "GitHub Security Advisory Pip",
						URL:  "https://github.com/advisories?query=type%3Areviewed+ecosystem%3Apip",
					},
				},
			},
		},
		{
			name:     "happy path, with only OS vulnerability, yes vendor severity, with both NVD and CVSS info",
			fixtures: []string{"testdata/fixtures/vulnerability.yaml"},
			vulns: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0004",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.RedHat,
						Name: "Red Hat OVAL v2",
						URL:  "https://www.redhat.com/security/data/oval/v2/",
					},
				},
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0004",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "dos",
						Description: "dos vulnerability",
						Severity:    dbTypes.SeverityLow.String(),
						VendorSeverity: dbTypes.VendorSeverity{
							vulnerability.RedHat: dbTypes.SeverityLow,
						},
						CweIDs:     []string{"CWE-311"},
						References: []string{"http://example.com"},
						CVSS: map[dbTypes.SourceID]dbTypes.CVSS{
							vulnerability.NVD: {
								V2Vector: "AV:N/AC:L/Au:N/C:P/I:P/A:P",
								V2Score:  4.5,
								V3Vector: "CVSS:3.0/PR:N/UI:N/S:U/C:H/I:H/A:H",
								V3Score:  5.6,
							},
							vulnerability.RedHat: {
								V2Vector: "AV:N/AC:M/Au:N/C:N/I:P/A:N",
								V2Score:  7.8,
								V3Vector: "CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H",
								V3Score:  9.8,
							},
						},
					},
					SeveritySource: vulnerability.RedHat,
					PrimaryURL:     "https://avd.aquasec.com/nvd/cve-2019-0004",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.RedHat,
						Name: "Red Hat OVAL v2",
						URL:  "https://www.redhat.com/security/data/oval/v2/",
					},
				},
			},
		},
		{
			name:     "happy path, with only library vulnerability",
			fixtures: []string{"testdata/fixtures/vulnerability.yaml"},
			vulns: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0005",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.GHSA,
						Name: "GitHub Security Advisory Pip",
						URL:  "https://github.com/advisories?query=type%3Areviewed+ecosystem%3Apip",
					},
				},
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0005",
					Vulnerability: dbTypes.Vulnerability{
						Title:       "COVID-19",
						Description: "a nasty virus vulnerability for humans",
						Severity:    dbTypes.SeverityCritical.String(),
						VendorSeverity: dbTypes.VendorSeverity{
							vulnerability.GHSA: dbTypes.SeverityCritical,
						},
						References: []string{"https://www.who.int/emergencies/diseases/novel-coronavirus-2019"},
					},
					SeveritySource: vulnerability.GHSA,
					PrimaryURL:     "https://avd.aquasec.com/nvd/cve-2019-0005",
					DataSource: &dbTypes.DataSource{
						ID:   vulnerability.GHSA,
						Name: "GitHub Security Advisory Pip",
						URL:  "https://github.com/advisories?query=type%3Areviewed+ecosystem%3Apip",
					},
				},
			},
		},
		{
			name:     "happy path, with package-specific severity",
			fixtures: []string{"testdata/fixtures/vulnerability.yaml"},
			vulns: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					SeveritySource:  vulnerability.Debian,
					Vulnerability: dbTypes.Vulnerability{
						Severity: dbTypes.SeverityLow.String(),
					},
				},
			},
			expectedVulnerabilities: []types.DetectedVulnerability{
				{
					VulnerabilityID: "CVE-2019-0001",
					SeveritySource:  vulnerability.Debian,
					Vulnerability: dbTypes.Vulnerability{
						Title:            "dos",
						Description:      "dos vulnerability",
						Severity:         dbTypes.SeverityLow.String(),
						References:       []string{"http://example.com"},
						LastModifiedDate: utils.MustTimeParse("2020-01-01T01:01:00Z"),
						PublishedDate:    utils.MustTimeParse("2001-01-01T01:01:00Z"),
					},
					PrimaryURL: "https://avd.aquasec.com/nvd/cve-2019-0001",
				},
			},
		},
		{
			name:     "GetVulnerability returns an error",
			fixtures: []string{"testdata/fixtures/sad.yaml"},
			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) {
			dbtest.InitDB(t, tt.fixtures)
			defer db.Close()

			c := vuln.NewClient(db.Config{})
			c.FillInfo(tt.vulns)
			assert.Equal(t, tt.expectedVulnerabilities, tt.vulns, tt.name)
		})
	}
}