// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package scan

import (
	"fmt"
	"testing"

	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/lib/q"
	"github.com/goharbor/harbor/src/pkg/scan/dao/scanner"
	v1 "github.com/goharbor/harbor/src/pkg/scan/rest/v1"
	htesting "github.com/goharbor/harbor/src/testing"
	"github.com/stretchr/testify/suite"
)

const sampleReportWithCompleteVulnData = `{
	"generated_at": "2020-08-01T18:28:49.072885592Z",
	"artifact": {
	  "repository": "library/ubuntu",
	  "digest": "sha256:d5b40885539615b9aeb7119516427959a158386af13e00d79a7da43ad1b3fb87",
	  "mime_type": "application/vnd.docker.distribution.manifest.v2+json"
	},
	"scanner": {
	  "name": "Trivy",
	  "vendor": "Aqua Security",
	  "version": "v0.9.1"
	},
	"severity": "Medium",
	"vulnerabilities": [
	  {
		"id": "CVE-2019-18276",
		"package": "bash",
		"version": "5.0-6ubuntu1.1",
		"severity": "Low",
		"description": "An issue was discovered in disable_priv_mode in shell.c in GNU Bash through 5.0 patch 11. By default, if Bash is run with its effective UID not equal to its real UID, it will drop privileges by setting its effective UID to its real UID. However, it does so incorrectly. On Linux and other systems that support \"saved UID\" functionality, the saved UID is not dropped. An attacker with command execution in the shell can use \"enable -f\" for runtime loading of a new builtin, which can be a shared object that calls setuid() and therefore regains privileges. However, binaries running with an effective UID of 0 are unaffected.",
		"links": [
		  "http://packetstormsecurity.com/files/155498/Bash-5.0-Patch-11-Privilege-Escalation.html",
		  "https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2019-18276",
		  "https://github.com/bminor/bash/commit/951bdaad7a18cc0dc1036bba86b18b90874d39ff",
		  "https://security.netapp.com/advisory/ntap-20200430-0003/",
		  "https://www.youtube.com/watch?v=-wGtxJ8opa8"
		],
		"layer": {
		  "digest": "sha256:4739cd2f4f486596c583c79f6033f1a9dee019389d512603609494678c8ccd53",
		  "diff_id": "sha256:f66829086c450acd5f67d0529a58f7120926c890f04e17aa7f0e9365da86480a"
		},
		"cwe_ids": ["CWE-476", "CWE-345"],
		"cvss":{
			"score_v3": 3.2,
			"score_v2": 2.3,
			"vector_v3": "CVSS:3.0/AV:L/AC:L/PR:L/UI:N/S:U/C:H/I:N/A:N",
			"vector_v2": "AV:L/AC:M/Au:N/C:P/I:N/A:N"
		},
		"vendor_attributes":[ {
			"key": "foo",
			"value": "bar"
		},
		{
			"key": "foo1",
			"value": "bar1"
		}
		]
	  }
	]
}`

// VulnerabilityTestSuite is test suite of testing vulnerability DAO.
type VulnerabilityTestSuite struct {
	htesting.Suite
	rpUUID                 string
	vulnerabilityRecordDao VulnerabilityRecordDao
	dao                    DAO
	registrationID         string
}

// TestVulnerabilityItem is the entry of ReportTestSuite.
func TestVulnerabilityItem(t *testing.T) {
	suite.Run(t, &VulnerabilityTestSuite{})
}

// SetupSuite prepares env for test suite.
func (suite *VulnerabilityTestSuite) SetupSuite() {
	suite.Suite.SetupSuite()
	suite.rpUUID = "uuid"
	suite.vulnerabilityRecordDao = NewVulnerabilityRecordDao()
	suite.dao = New()

}

// SetupTest prepares env for test case.
func (suite *VulnerabilityTestSuite) SetupTest() {
	r := &Report{
		UUID:             "uuid",
		Digest:           "digest1001",
		RegistrationUUID: "scannerId1",
		MimeType:         v1.MimeTypeNativeReport,
		Status:           job.PendingStatus.String(),
		Report:           sampleReportWithCompleteVulnData,
	}
	suite.registerScanner(r.RegistrationUUID)
	suite.createReport(r)
	vulns := generateVulnerabilityRecordsForReport("uuid", "scannerId1", 10)
	for _, v := range vulns {
		suite.insertVulnRecordForReport("uuid", v)
	}

}

// TearDownTest clears enf for test case.
func (suite *VulnerabilityTestSuite) TearDownTest() {
	registrations, err := scanner.ListRegistrations(suite.Context(), &q.Query{})
	suite.NoError(err, "Failed to cleanup scanner registrations")
	for _, registration := range registrations {
		err = scanner.DeleteRegistration(suite.Context(), registration.UUID)
		suite.NoError(err, "Error when cleaning up scanner registrations")
	}
	reports, err := suite.dao.List(suite.Context(), &q.Query{})
	suite.NoError(err)
	for _, report := range reports {
		suite.cleanUpAdditionalData(report.UUID, report.RegistrationUUID)
	}
}

// TestVulnerabilityRecordsListForReport tests listing of vulnerability record for reports
func (suite *VulnerabilityTestSuite) TestVulnerabilityRecordsListForReport() {
	// create a second report and associate the same  vulnerability record set to the report
	r := &Report{
		UUID:             "uuid1",
		Digest:           "digest1002",
		RegistrationUUID: "scannerId2",
		MimeType:         v1.MimeTypeNativeReport,
		Status:           job.PendingStatus.String(),
		Report:           sampleReportWithCompleteVulnData,
	}
	suite.registerScanner(r.RegistrationUUID)
	suite.createReport(r)
	// insert a set of vulnerability records for this report. the vulnerability records
	// belong to the same scanner
	vulns := generateVulnerabilityRecordsForReport("uuid1", "scannerId2", 10)
	for _, v := range vulns {
		suite.insertVulnRecordForReport("uuid1", v)
	}

	// fetch the records for the first report. Additionally assert that these records
	// indeed belong to the same report being fetched and not to another report
	{
		vulns, err := suite.vulnerabilityRecordDao.GetForReport(suite.Context(), "uuid")
		suite.NoError(err, "Error when fetching vulnerability records for report")
		suite.True(len(vulns) > 0)
	}
	{
		vulns, err := suite.vulnerabilityRecordDao.GetForReport(suite.Context(), "uuid1")
		suite.NoError(err, "Error when fetching vulnerability records for report")
		suite.True(len(vulns) > 0)
	}

}

// TestGetVulnerabilityRecordsForScanner gets vulnerability records for scanner
func (suite *VulnerabilityTestSuite) TestGetVulnerabilityRecordsForScanner() {

	vulns, err := suite.vulnerabilityRecordDao.GetForScanner(suite.Context(), "scannerId1")
	suite.NoError(err, "Error when fetching vulnerability records for report")
	suite.True(len(vulns) > 0)
}

// TestGetVulnerabilityRecordIdsForScanner gets vulnerability records for scanner
func (suite *VulnerabilityTestSuite) TestGetVulnerabilityRecordIdsForScanner() {
	vulns, err := suite.vulnerabilityRecordDao.GetRecordIdsForScanner(suite.Context(), "scannerId1")
	suite.NoError(err, "Error when fetching vulnerability records for report")
	suite.True(len(vulns) > 0)
}

// TestDeleteForDigest tests deleting vulnerability report for a specific digest
func (suite *VulnerabilityTestSuite) TestDeleteForDigest() {
	// create a second report and associate the same  vulnerability record set to the report
	r := &Report{
		UUID:             "uuid1",
		Digest:           "digest1",
		RegistrationUUID: "scannerId2",
		MimeType:         v1.MimeTypeNativeReport,
		Status:           job.PendingStatus.String(),
		Report:           sampleReportWithCompleteVulnData,
	}
	suite.registerScanner(r.RegistrationUUID)
	suite.createReport(r)
	// insert a set of vulnerability records for this report. the vulnerability records
	// belong to the same scanner
	vulns := generateVulnerabilityRecordsForReport("uuid1", "scannerId2", 10)
	for _, v := range vulns {
		suite.insertVulnRecordForReport("uuid1", v)
	}
	delCount, err := suite.vulnerabilityRecordDao.DeleteForDigests(suite.Context(), "digest1")
	suite.NoError(err)
	suite.Equal(int64(10), delCount)
}

func (suite *VulnerabilityTestSuite) TestDuplicateRecords() {
	r := &Report{
		UUID:             "uuid2",
		Digest:           "digest1002",
		RegistrationUUID: "scannerId1",
		MimeType:         v1.MimeTypeNativeReport,
		Status:           job.PendingStatus.String(),
		Report:           sampleReportWithCompleteVulnData,
	}
	suite.createReport(r)
	vulns := generateVulnerabilityRecordsForReport("uuid2", "scannerId1", 10)
	for _, v := range vulns {
		suite.insertVulnRecordForReport("uuid2", v)
	}
}

// TestDeleteVulnerabilityRecord gets vulnerability records for scanner
func (suite *VulnerabilityTestSuite) TestDeleteVulnerabilityRecord() {

	vulns, err := suite.vulnerabilityRecordDao.GetForScanner(suite.Context(), "scannerId1")
	suite.NoError(err, "Error when fetching vulnerability records for report")
	suite.True(len(vulns) > 0)
	for _, vuln := range vulns {
		err = suite.vulnerabilityRecordDao.Delete(suite.Context(), vuln)
		suite.NoError(err)
	}
}

// TestListVulnerabilityRecord gets vulnerability records for scanner
func (suite *VulnerabilityTestSuite) TestListVulnerabilityRecord() {

	vulns, err := suite.vulnerabilityRecordDao.List(suite.Context(), &q.Query{Keywords: map[string]interface{}{"CVEID": "CVE-ID1"}})
	suite.NoError(err, "Error when fetching vulnerability records for report")
	suite.True(len(vulns) > 0)
}

func (suite *VulnerabilityTestSuite) createReport(r *Report) {
	id, err := suite.dao.Create(suite.Context(), r)
	suite.NoError(err)
	suite.Condition(func() (success bool) {
		success = id > 0
		return
	})
}

func (suite *VulnerabilityTestSuite) insertVulnRecordForReport(reportUUID string, vr *VulnerabilityRecord) {
	id, err := suite.vulnerabilityRecordDao.Create(suite.Context(), vr)
	suite.NoError(err, "Failed to create vulnerability record")

	err = suite.vulnerabilityRecordDao.InsertForReport(suite.Context(), reportUUID, id)
	suite.NoError(err, "Failed to insert vulnerability record row for report %s", reportUUID)
}

func (suite *VulnerabilityTestSuite) cleanUpAdditionalData(reportID string, scannerID string) {
	_, err := suite.dao.DeleteMany(suite.Context(), q.Query{Keywords: q.KeyWords{"uuid": reportID}})

	suite.NoError(err)
	_, err = suite.vulnerabilityRecordDao.DeleteForReport(suite.Context(), reportID)
	suite.NoError(err, "Failed to cleanup records")
	_, err = suite.vulnerabilityRecordDao.DeleteForScanner(suite.Context(), scannerID)
	suite.NoError(err, "Failed to delete vulnerability records")
}

func (suite *VulnerabilityTestSuite) registerScanner(registrationUUID string) {
	r := &scanner.Registration{
		UUID:        registrationUUID,
		Name:        registrationUUID,
		Description: "sample registration",
		URL:         fmt.Sprintf("https://sample.scanner.com/%s", registrationUUID),
	}

	_, err := scanner.AddRegistration(suite.Context(), r)
	suite.NoError(err, "add new registration")
}

func generateVulnerabilityRecordsForReport(reportUUID string, registrationUUID string, numRecords int) []*VulnerabilityRecord {
	vulns := make([]*VulnerabilityRecord, 0)
	for i := 1; i <= numRecords; i++ {
		vulnV2 := new(VulnerabilityRecord)
		vulnV2.CVEID = fmt.Sprintf("CVE-ID%d", i)
		vulnV2.Package = fmt.Sprintf("Package%d", i)
		vulnV2.PackageVersion = "NotAvailable"
		vulnV2.PackageType = "Unknown"
		vulnV2.Fix = "1.0.0"
		vulnV2.URLs = "url1"
		vulnV2.RegistrationUUID = registrationUUID
		if i%2 == 0 {
			vulnV2.Severity = "High"
		} else if i%3 == 0 {
			vulnV2.Severity = "Medium"
		} else if i%4 == 0 {
			vulnV2.Severity = "Critical"
		} else {
			vulnV2.Severity = "Low"
		}
		vulns = append(vulns, vulnV2)
	}

	return vulns
}
