// Copyright 2017 clair 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 vulnerability

import (
	"database/sql"
	"math/rand"
	"strconv"
	"testing"

	"github.com/pborman/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/coreos/clair/database"
	"github.com/coreos/clair/database/pgsql/feature"
	"github.com/coreos/clair/database/pgsql/namespace"
	"github.com/coreos/clair/database/pgsql/testutil"
	"github.com/coreos/clair/ext/versionfmt"
	"github.com/coreos/clair/ext/versionfmt/dpkg"
	"github.com/coreos/clair/pkg/strutil"
)

func TestInsertVulnerabilities(t *testing.T) {
	store, cleanup := testutil.CreateTestDBWithFixture(t, "InsertVulnerabilities")
	defer cleanup()

	ns1 := database.Namespace{
		Name:          "name",
		VersionFormat: "random stuff",
	}

	ns2 := database.Namespace{
		Name:          "debian:7",
		VersionFormat: "dpkg",
	}

	// invalid vulnerability
	v1 := database.Vulnerability{
		Name:      "invalid",
		Namespace: ns1,
	}

	vwa1 := database.VulnerabilityWithAffected{
		Vulnerability: v1,
	}
	// valid vulnerability
	v2 := database.Vulnerability{
		Name:      "valid",
		Namespace: ns2,
		Severity:  database.UnknownSeverity,
	}

	vwa2 := database.VulnerabilityWithAffected{
		Vulnerability: v2,
	}

	tx, err := store.Begin()
	require.Nil(t, err)

	// empty
	err = InsertVulnerabilities(tx, []database.VulnerabilityWithAffected{})
	assert.Nil(t, err)

	// invalid content: vwa1 is invalid
	err = InsertVulnerabilities(tx, []database.VulnerabilityWithAffected{vwa1, vwa2})
	assert.NotNil(t, err)

	tx = testutil.RestartTransaction(store, tx, false)
	// invalid content: duplicated input
	err = InsertVulnerabilities(tx, []database.VulnerabilityWithAffected{vwa2, vwa2})
	assert.NotNil(t, err)

	tx = testutil.RestartTransaction(store, tx, false)
	// valid content
	err = InsertVulnerabilities(tx, []database.VulnerabilityWithAffected{vwa2})
	assert.Nil(t, err)

	tx = testutil.RestartTransaction(store, tx, true)
	// ensure the content is in database
	vulns, err := FindVulnerabilities(tx, []database.VulnerabilityID{{Name: "valid", Namespace: "debian:7"}})
	if assert.Nil(t, err) && assert.Len(t, vulns, 1) {
		assert.True(t, vulns[0].Valid)
	}

	tx = testutil.RestartTransaction(store, tx, false)
	// valid content: vwa2 removed and inserted
	err = DeleteVulnerabilities(tx, []database.VulnerabilityID{{Name: vwa2.Name, Namespace: vwa2.Namespace.Name}})
	assert.Nil(t, err)

	err = InsertVulnerabilities(tx, []database.VulnerabilityWithAffected{vwa2})
	assert.Nil(t, err)

	require.Nil(t, tx.Rollback())
}

func TestCachingVulnerable(t *testing.T) {
	tx, cleanup := testutil.CreateTestTxWithFixtures(t, "CachingVulnerable")
	defer cleanup()

	ns := database.Namespace{
		Name:          "debian:8",
		VersionFormat: dpkg.ParserName,
	}

	f := database.NamespacedFeature{
		Feature: database.Feature{
			Name:          "openssl",
			Version:       "1.0",
			VersionFormat: dpkg.ParserName,
			Type:          database.SourcePackage,
		},
		Namespace: ns,
	}

	vuln := database.VulnerabilityWithAffected{
		Vulnerability: database.Vulnerability{
			Name:      "CVE-YAY",
			Namespace: ns,
			Severity:  database.HighSeverity,
		},
		Affected: []database.AffectedFeature{
			{
				Namespace:       ns,
				FeatureName:     "openssl",
				FeatureType:     database.SourcePackage,
				AffectedVersion: "2.0",
				FixedInVersion:  "2.1",
			},
		},
	}

	vuln2 := database.VulnerabilityWithAffected{
		Vulnerability: database.Vulnerability{
			Name:      "CVE-YAY2",
			Namespace: ns,
			Severity:  database.HighSeverity,
		},
		Affected: []database.AffectedFeature{
			{
				Namespace:       ns,
				FeatureName:     "openssl",
				FeatureType:     database.SourcePackage,
				AffectedVersion: "2.1",
				FixedInVersion:  "2.2",
			},
		},
	}

	vulnFixed1 := database.VulnerabilityWithFixedIn{
		Vulnerability: database.Vulnerability{
			Name:      "CVE-YAY",
			Namespace: ns,
			Severity:  database.HighSeverity,
		},
		FixedInVersion: "2.1",
	}

	vulnFixed2 := database.VulnerabilityWithFixedIn{
		Vulnerability: database.Vulnerability{
			Name:      "CVE-YAY2",
			Namespace: ns,
			Severity:  database.HighSeverity,
		},
		FixedInVersion: "2.2",
	}

	require.Nil(t, InsertVulnerabilities(tx, []database.VulnerabilityWithAffected{vuln, vuln2}))
	r, err := FindAffectedNamespacedFeatures(tx, []database.NamespacedFeature{f})
	assert.Nil(t, err)
	assert.Len(t, r, 1)
	for _, anf := range r {
		if assert.True(t, anf.Valid) && assert.Len(t, anf.AffectedBy, 2) {
			for _, a := range anf.AffectedBy {
				if a.Name == "CVE-YAY" {
					assert.Equal(t, vulnFixed1, a)
				} else if a.Name == "CVE-YAY2" {
					assert.Equal(t, vulnFixed2, a)
				} else {
					t.FailNow()
				}
			}
		}
	}
}

func TestFindVulnerabilities(t *testing.T) {
	tx, cleanup := testutil.CreateTestTxWithFixtures(t, "FindVulnerabilities")
	defer cleanup()

	vuln, err := FindVulnerabilities(tx, []database.VulnerabilityID{
		{Name: "CVE-OPENSSL-1-DEB7", Namespace: "debian:7"},
		{Name: "CVE-NOPE", Namespace: "debian:7"},
		{Name: "CVE-NOT HERE"},
	})

	ns := database.Namespace{
		Name:          "debian:7",
		VersionFormat: "dpkg",
	}

	expectedExisting := []database.VulnerabilityWithAffected{
		{
			Vulnerability: database.Vulnerability{
				Namespace:   ns,
				Name:        "CVE-OPENSSL-1-DEB7",
				Description: "A vulnerability affecting OpenSSL < 2.0 on Debian 7.0",
				Link:        "http://google.com/#q=CVE-OPENSSL-1-DEB7",
				Severity:    database.HighSeverity,
			},
			Affected: []database.AffectedFeature{
				{
					FeatureName:     "openssl",
					FeatureType:     database.SourcePackage,
					AffectedVersion: "2.0",
					FixedInVersion:  "2.0",
					Namespace:       ns,
				},
				{
					FeatureName:     "libssl",
					FeatureType:     database.SourcePackage,
					AffectedVersion: "1.9-abc",
					FixedInVersion:  "1.9-abc",
					Namespace:       ns,
				},
			},
		},
		{
			Vulnerability: database.Vulnerability{
				Namespace:   ns,
				Name:        "CVE-NOPE",
				Description: "A vulnerability affecting nothing",
				Severity:    database.UnknownSeverity,
			},
		},
	}

	expectedExistingMap := map[database.VulnerabilityID]database.VulnerabilityWithAffected{}
	for _, v := range expectedExisting {
		expectedExistingMap[database.VulnerabilityID{Name: v.Name, Namespace: v.Namespace.Name}] = v
	}

	nonexisting := database.VulnerabilityWithAffected{
		Vulnerability: database.Vulnerability{Name: "CVE-NOT HERE"},
	}

	if assert.Nil(t, err) {
		for _, v := range vuln {
			if v.Valid {
				key := database.VulnerabilityID{
					Name:      v.Name,
					Namespace: v.Namespace.Name,
				}

				expected, ok := expectedExistingMap[key]
				if assert.True(t, ok, "vulnerability not found: "+key.Name+":"+key.Namespace) {
					testutil.AssertVulnerabilityWithAffectedEqual(t, expected, v.VulnerabilityWithAffected)
				}
			} else if !assert.Equal(t, nonexisting, v.VulnerabilityWithAffected) {
				t.FailNow()
			}
		}
	}

	// same vulnerability
	r, err := FindVulnerabilities(tx, []database.VulnerabilityID{
		{Name: "CVE-OPENSSL-1-DEB7", Namespace: "debian:7"},
		{Name: "CVE-OPENSSL-1-DEB7", Namespace: "debian:7"},
	})

	if assert.Nil(t, err) {
		for _, vuln := range r {
			if assert.True(t, vuln.Valid) {
				expected, _ := expectedExistingMap[database.VulnerabilityID{Name: "CVE-OPENSSL-1-DEB7", Namespace: "debian:7"}]
				testutil.AssertVulnerabilityWithAffectedEqual(t, expected, vuln.VulnerabilityWithAffected)
			}
		}
	}
}

func TestDeleteVulnerabilities(t *testing.T) {
	tx, cleanup := testutil.CreateTestTxWithFixtures(t, "DeleteVulnerabilities")
	defer cleanup()

	remove := []database.VulnerabilityID{}
	// empty case
	assert.Nil(t, DeleteVulnerabilities(tx, remove))
	// invalid case
	remove = append(remove, database.VulnerabilityID{})
	assert.NotNil(t, DeleteVulnerabilities(tx, remove))

	// valid case
	validRemove := []database.VulnerabilityID{
		{Name: "CVE-OPENSSL-1-DEB7", Namespace: "debian:7"},
		{Name: "CVE-NOPE", Namespace: "debian:7"},
	}

	assert.Nil(t, DeleteVulnerabilities(tx, validRemove))
	vuln, err := FindVulnerabilities(tx, validRemove)
	if assert.Nil(t, err) {
		for _, v := range vuln {
			assert.False(t, v.Valid)
		}
	}
}

func TestFindVulnerabilityIDs(t *testing.T) {
	tx, cleanup := testutil.CreateTestTxWithFixtures(t, "FindVulnerabilityIDs")
	defer cleanup()

	ids, err := FindLatestDeletedVulnerabilityIDs(tx, []database.VulnerabilityID{{Name: "CVE-DELETED", Namespace: "debian:7"}})
	if assert.Nil(t, err) {
		if !(assert.Len(t, ids, 1) && assert.True(t, ids[0].Valid) && assert.Equal(t, 3, int(ids[0].Int64))) {
			assert.Fail(t, "")
		}
	}

	ids, err = FindNotDeletedVulnerabilityIDs(tx, []database.VulnerabilityID{{Name: "CVE-NOPE", Namespace: "debian:7"}})
	if assert.Nil(t, err) {
		if !(assert.Len(t, ids, 1) && assert.True(t, ids[0].Valid) && assert.Equal(t, 2, int(ids[0].Int64))) {
			assert.Fail(t, "")
		}
	}
}

func TestFindAffectedNamespacedFeatures(t *testing.T) {
	tx, cleanup := testutil.CreateTestTxWithFixtures(t, "FindAffectedNamespacedFeatures")
	defer cleanup()

	ns := database.NamespacedFeature{
		Feature: database.Feature{
			Name:          "openssl",
			Version:       "1.0",
			VersionFormat: "dpkg",
			Type:          database.SourcePackage,
		},
		Namespace: database.Namespace{
			Name:          "debian:7",
			VersionFormat: "dpkg",
		},
	}

	ans, err := FindAffectedNamespacedFeatures(tx, []database.NamespacedFeature{ns})
	if assert.Nil(t, err) &&
		assert.Len(t, ans, 1) &&
		assert.True(t, ans[0].Valid) &&
		assert.Len(t, ans[0].AffectedBy, 1) {
		assert.Equal(t, "CVE-OPENSSL-1-DEB7", ans[0].AffectedBy[0].Name)
	}
}

func genRandomVulnerabilityAndNamespacedFeature(t *testing.T, store *sql.DB) ([]database.NamespacedFeature, []database.VulnerabilityWithAffected) {
	tx, err := store.Begin()
	if err != nil {
		panic(err)
	}

	numFeatures := 100
	numVulnerabilities := 100

	featureName := "TestFeature"
	featureVersionFormat := dpkg.ParserName
	// Insert the namespace on which we'll work.
	ns := database.Namespace{
		Name:          "TestRaceAffectsFeatureNamespace1",
		VersionFormat: dpkg.ParserName,
	}

	if !assert.Nil(t, namespace.PersistNamespaces(tx, []database.Namespace{ns})) {
		t.FailNow()
	}

	// Generate Distinct random features
	features := make([]database.Feature, numFeatures)
	nsFeatures := make([]database.NamespacedFeature, numFeatures)
	for i := 0; i < numFeatures; i++ {
		version := rand.Intn(numFeatures)

		features[i] = *database.NewSourcePackage(featureName, strconv.Itoa(version), featureVersionFormat)
		nsFeatures[i] = database.NamespacedFeature{
			Namespace: ns,
			Feature:   features[i],
		}
	}

	if !assert.Nil(t, feature.PersistFeatures(tx, features)) {
		t.FailNow()
	}

	// Generate vulnerabilities.
	vulnerabilities := []database.VulnerabilityWithAffected{}
	for i := 0; i < numVulnerabilities; i++ {
		// any version less than this is vulnerable
		version := rand.Intn(numFeatures) + 1

		vulnerability := database.VulnerabilityWithAffected{
			Vulnerability: database.Vulnerability{
				Name:      uuid.New(),
				Namespace: ns,
				Severity:  database.UnknownSeverity,
			},
			Affected: []database.AffectedFeature{
				{
					Namespace:       ns,
					FeatureName:     featureName,
					FeatureType:     database.SourcePackage,
					AffectedVersion: strconv.Itoa(version),
					FixedInVersion:  strconv.Itoa(version),
				},
			},
		}

		vulnerabilities = append(vulnerabilities, vulnerability)
	}
	tx.Commit()

	return nsFeatures, vulnerabilities
}

func TestVulnChangeAffectsVulnerableFeatures(t *testing.T) {
	db, cleanup := testutil.CreateTestDB(t, "caching")
	defer cleanup()

	nsFeatures, vulnerabilities := genRandomVulnerabilityAndNamespacedFeature(t, db)
	tx, err := db.Begin()
	require.Nil(t, err)

	require.Nil(t, feature.PersistNamespacedFeatures(tx, nsFeatures))
	require.Nil(t, tx.Commit())

	tx, err = db.Begin()
	require.Nil(t, InsertVulnerabilities(tx, vulnerabilities))
	require.Nil(t, tx.Commit())

	tx, err = db.Begin()
	require.Nil(t, err)
	defer tx.Rollback()

	affected, err := FindAffectedNamespacedFeatures(tx, nsFeatures)
	require.Nil(t, err)

	for _, ansf := range affected {
		require.True(t, ansf.Valid)

		expectedAffectedNames := []string{}
		for _, vuln := range vulnerabilities {
			if ok, err := versionfmt.InRange(dpkg.ParserName, ansf.Version, vuln.Affected[0].AffectedVersion); err == nil {
				if ok {
					expectedAffectedNames = append(expectedAffectedNames, vuln.Name)
				}
			}
		}

		actualAffectedNames := []string{}
		for _, s := range ansf.AffectedBy {
			actualAffectedNames = append(actualAffectedNames, s.Name)
		}

		require.Len(t, strutil.Difference(expectedAffectedNames, actualAffectedNames), 0, "\nvulns: %#v\nfeature:%#v\nexpected:%#v\nactual:%#v", vulnerabilities, ansf.NamespacedFeature, expectedAffectedNames, actualAffectedNames)
		require.Len(t, strutil.Difference(actualAffectedNames, expectedAffectedNames), 0)
	}
}